llm_be/chat_backend/oauth.py

@@ -338,6 +338,11 @@ def upsert_drive_connection(
     *, user: CustomUser, kind: str, profile: ProviderProfile
 ) -> DriveConnection:
     """Create/refresh a DriveConnection from a link_drive/link_company_drive callback (#47)."""
+    if not user.company_id:
+        raise OAuthError(
+            "no_company",
+            "A company is required before connecting a Drive account.",
+        )
     expires_at = _token_expiry(profile.expires_in)
     lookup_user = user if kind == DriveConnection.Kind.PERSONAL else None
     connection = DriveConnection.objects.filter(

llm_be/chat_backend/tests/test_oauth.py

@@ -431,6 +431,25 @@ class OAuthCallbackDriveLinkTestCase(APITestCase):
         self.assertFalse(DriveConnection.objects.exists())
         mock_exchange.assert_not_called()
 
+    @patch("chat_backend.views_oauth.exchange_code_for_profile")
+    def test_callback_rejects_user_without_company(self, mock_exchange):
+        """#55 — users with no company must not hit DriveConnection.company NOT NULL."""
+        self.user.company = None
+        self.user.save(update_fields=["company"])
+        mock_exchange.return_value = _google_profile()
+        state = self._state(intent="link_drive")
+
+        response = self.client.get(
+            reverse("oauth_callback", kwargs={"provider": "google"}),
+            {"code": "auth-code", "state": state},
+        )
+
+        self.assertEqual(response.status_code, status.HTTP_302_FOUND)
+        params = parse_qs(urlparse(response["Location"]).query)
+        self.assertEqual(params["error"], ["no_company"])
+        self.assertFalse(DriveConnection.objects.exists())
+        mock_exchange.assert_not_called()
+
     def test_callback_missing_user_id_in_state_is_rejected(self):
         # Simulates a forged/legacy state payload without the linking user.
         from django.core import signing

llm_be/chat_backend/views_oauth.py

@@ -229,6 +229,11 @@ class OAuthCallbackView(APIView):
                     "forbidden",
                     "Only a company manager can connect a company Drive.",
                 )
+            if not user.company_id:
+                raise OAuthError(
+                    "no_company",
+                    "A company is required before connecting a Drive account.",
+                )
             assert_feature_allowed(user, "rag")
 
             redirect_uri = _callback_redirect_uri(request, provider)