diff --git a/app/Http/Middleware/TokenAuth.php b/app/Http/Middleware/TokenAuth.php
index e921459ff2..b68524a156 100644
--- a/app/Http/Middleware/TokenAuth.php
+++ b/app/Http/Middleware/TokenAuth.php
@@ -33,13 +33,15 @@ class TokenAuth
     public function handle($request, Closure $next)
     {
 
-         if ($request->header('X-API-TOKEN') && ($company_token = CompanyToken::with([
+        if (config('ninja.db.multi_db_enabled') &&
+            $request->header('X-API-TOKEN') &&
+             ($company_token = MultiDB::getCompanyToken($request->header('X-API-TOKEN')))) {
+        } elseif ($request->header('X-API-TOKEN') && ($company_token = CompanyToken::with([
             'user.account',
             'company',
             'account',
-            'cu',
+            'cu'
             ])->where('token', $request->header('X-API-TOKEN'))->first())) {
-
         } else {
             return response()->json(['message' => 'Invalid token'], 403);
         }
diff --git a/app/Libraries/MultiDB.php b/app/Libraries/MultiDB.php
index e879b9120f..4d0e21e01d 100644
--- a/app/Libraries/MultiDB.php
+++ b/app/Libraries/MultiDB.php
@@ -336,6 +336,7 @@ class MultiDB
                 'user.account',
                 'company',
                 'account',
+                'cu'
             ])->where('token', $token)->first()) {
 
                 self::setDB($db);
diff --git a/app/Models/CompanyGateway.php b/app/Models/CompanyGateway.php
index b29e99a1be..4bacfc382b 100644
--- a/app/Models/CompanyGateway.php
+++ b/app/Models/CompanyGateway.php
@@ -203,7 +203,7 @@ class CompanyGateway extends BaseModel
     }
 
     /* This is the public entry point into the payment superclass */
-    public function driver(Client $client = null)
+    public function driver(?Client $client = null)
     {
         // $class = static::driver_class();
         $class = self::driver_class();
diff --git a/app/Models/CompanyToken.php b/app/Models/CompanyToken.php
index 19672fd4ba..27f108a3ad 100644
--- a/app/Models/CompanyToken.php
+++ b/app/Models/CompanyToken.php
@@ -95,7 +95,6 @@ class CompanyToken extends BaseModel
         return $this->hasOne(CompanyUser::class, ['user_id', 'company_id'], ['user_id', 'company_id']);
     }
 
-
     // public function company_user(): \Illuminate\Database\Eloquent\Relations\HasOne
     // {
     //     return $this->hasOne(CompanyUser::class, 'user_id', 'user_id')
diff --git a/app/Models/User.php b/app/Models/User.php
index 33eddc3737..c9b3eace7b 100644
--- a/app/Models/User.php
+++ b/app/Models/User.php
@@ -234,12 +234,16 @@ class User extends Authenticatable implements MustVerifyEmail
             return $truth->getCompanyToken();
         }
 
-        // if (request()->header('X-API-TOKEN')) {
         if (request()->header('X-API-TOKEN')) {
-            return CompanyToken::with(['cu'])->where('token', request()->header('X-API-TOKEN'))->first();
+            
+            $token = CompanyToken::with(['cu'])->where('token', request()->header('X-API-TOKEN'))->first();
+            if ($token) {
+                return $token;
+            }
+
         }
 
-        return $this->tokens()->first();
+        return $this->tokens()->with(['cu'])->first();
     }
 
     /**