Gateway Fixes

Author: Asaif Ali

agent_service/app/infrastructure/agents_backend/model_provider.py

@@ -20,19 +20,15 @@
 # Environment Variables
 # --------------------------------------------------
 MODEL_TYPE = os.getenv("MODEL_TYPE", "OpenAI")
-OPENAI_MODEL_ID = os.getenv("OPENAI_MODEL_ID", "gemini-2.5-flash")
-OPENAI_BASE_URL = os.getenv("LLM_BASE_URL") or os.getenv("OPENAI_BASE_URL")
+OPENAI_MODEL_ID = os.getenv("LLM_MODEL") or os.getenv("OPENAI_MODEL_ID", "gateway-managed")
 LLM_GATEWAY_URL = os.getenv("LLM_GATEWAY_URL", "https://portfolio-llm-gateway.onrender.com/v1").strip()
 LLM_GATEWAY_TIMEOUT = float(os.getenv("LLM_GATEWAY_TIMEOUT", "180"))
 # Demo/portfolio gateways commonly rate-limit bursts. Avoid immediate SDK retries
 # that amplify a 429; make the value configurable for production.
-LLM_GATEWAY_MAX_RETRIES = max(0, int(os.getenv("LLM_GATEWAY_MAX_RETRIES", "2")))
+LLM_GATEWAY_MAX_RETRIES = max(0, int(os.getenv("LLM_GATEWAY_MAX_RETRIES", "0")))
 VLLM_BASE_URL = os.getenv("VLLM_BASE_URL")
 VLLM_CHAT_MODEL_ID = os.getenv("VLLM_CHAT_MODEL_ID")
 VLLM_API_KEY = os.getenv("VLLM_API_KEY") or "local"
-DIRECT_PROVIDER = os.getenv("LLM_DIRECT_PROVIDER", "google").strip().lower()
-GOOGLE_OPENAI_BASE_URL = os.getenv("GOOGLE_OPENAI_BASE_URL", "https://generativelanguage.googleapis.com/v1beta/openai/").strip()
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "").strip()
 # --------------------------------------------------
 # Message Normalization (vLLM compatibility)
 # --------------------------------------------------
@@ -92,80 +88,47 @@ def _gateway_config(self):
 
     def get_client(self):
         token, gateway_url = self._gateway_config()
-        if token:
-            if not gateway_url:
-                raise RuntimeError("LLM gateway token present but LLM_GATEWAY_URL is not configured")
-            logger.info(
-                "Using request-scoped Portfolio LLM Gateway for model=%s base_url=%s",
-                self.id,
-                gateway_url,
+        if not token:
+            raise RuntimeError(
+                "LLM gateway session token is required. Pass X-LLM-Gateway-Token from the portfolio session."
             )
-            return OpenAI(
-                api_key=token,
-                base_url=gateway_url,
-                timeout=max(float(self.timeout or 0), LLM_GATEWAY_TIMEOUT),
-                max_retries=LLM_GATEWAY_MAX_RETRIES if self.max_retries in (None, 0) else int(self.max_retries),
-                default_headers=self.default_headers,
-                default_query=self.default_query,
-            )
-        # Standalone LegacyLens deployments do not have a portfolio session token.
-        # Use a remote OpenAI-compatible free-tier provider instead of silently
-        # falling through to an unauthenticated OpenRouter endpoint.
-        if DIRECT_PROVIDER == "google":
-            if not GOOGLE_API_KEY:
-                raise RuntimeError(
-                    "No LLM gateway session token and GOOGLE_API_KEY is not configured for direct Gemini fallback"
-                )
-            logger.info(
-                "No gateway token; using direct Google Gemini OpenAI-compatible endpoint for model=%s",
-                self.id,
-            )
-            return OpenAI(
-                api_key=GOOGLE_API_KEY,
-                base_url=GOOGLE_OPENAI_BASE_URL,
-                timeout=max(float(self.timeout or 0), LLM_GATEWAY_TIMEOUT),
-                max_retries=LLM_GATEWAY_MAX_RETRIES if self.max_retries in (None, 0) else int(self.max_retries),
-                default_headers=self.default_headers,
-                default_query=self.default_query,
-            )
-        return super().get_client()
+        if not gateway_url:
+            raise RuntimeError("LLM gateway token present but LLM_GATEWAY_URL is not configured")
+        logger.info(
+            "Using request-scoped Portfolio LLM Gateway for model=%s base_url=%s",
+            self.id,
+            gateway_url,
+        )
+        return OpenAI(
+            api_key=token,
+            base_url=gateway_url,
+            timeout=max(float(self.timeout or 0), LLM_GATEWAY_TIMEOUT),
+            max_retries=LLM_GATEWAY_MAX_RETRIES if self.max_retries in (None, 0) else int(self.max_retries),
+            default_headers=self.default_headers,
+            default_query=self.default_query,
+        )
 
     def get_async_client(self):
         token, gateway_url = self._gateway_config()
-        if token:
-            if not gateway_url:
-                raise RuntimeError("LLM gateway token present but LLM_GATEWAY_URL is not configured")
-            logger.info(
-                "Using request-scoped Portfolio LLM Gateway (async) for model=%s base_url=%s",
-                self.id,
-                gateway_url,
+        if not token:
+            raise RuntimeError(
+                "LLM gateway session token is required. Pass X-LLM-Gateway-Token from the portfolio session."
             )
-            return AsyncOpenAI(
-                api_key=token,
-                base_url=gateway_url,
-                timeout=max(float(self.timeout or 0), LLM_GATEWAY_TIMEOUT),
-                max_retries=LLM_GATEWAY_MAX_RETRIES if self.max_retries in (None, 0) else int(self.max_retries),
-                default_headers=self.default_headers,
-                default_query=self.default_query,
-            )
-        if DIRECT_PROVIDER == "google":
-            if not GOOGLE_API_KEY:
-                raise RuntimeError(
-                    "No LLM gateway session token and GOOGLE_API_KEY is not configured for direct Gemini fallback"
-                )
-            logger.info(
-                "No gateway token; using direct Google Gemini OpenAI-compatible endpoint (async) for model=%s",
-                self.id,
-            )
-            return AsyncOpenAI(
-                api_key=GOOGLE_API_KEY,
-                base_url=GOOGLE_OPENAI_BASE_URL,
-                timeout=max(float(self.timeout or 0), LLM_GATEWAY_TIMEOUT),
-                max_retries=LLM_GATEWAY_MAX_RETRIES if self.max_retries in (None, 0) else int(self.max_retries),
-                default_headers=self.default_headers,
-                default_query=self.default_query,
-            )
-        return super().get_async_client()
+        if not gateway_url:
+            raise RuntimeError("LLM gateway token present but LLM_GATEWAY_URL is not configured")
+        logger.info(
+            "Using request-scoped Portfolio LLM Gateway (async) for model=%s base_url=%s",
+            self.id,
+            gateway_url,
+        )
+        return AsyncOpenAI(
+            api_key=token,
+            base_url=gateway_url,
+            timeout=max(float(self.timeout or 0), LLM_GATEWAY_TIMEOUT),
+            max_retries=LLM_GATEWAY_MAX_RETRIES if self.max_retries in (None, 0) else int(self.max_retries),
+            default_headers=self.default_headers,
+            default_query=self.default_query,
+        )
 
 
 
@@ -177,8 +140,8 @@ def create_model():
     if MODEL_TYPE == "OpenAI":
         model = GatewayAwareOpenAIChat(
             id=OPENAI_MODEL_ID,
-            api_key=os.getenv("OPENAI_API_KEY"),
-            base_url=OPENAI_BASE_URL,
+            api_key="gateway-session",
+            base_url=LLM_GATEWAY_URL,
             temperature=0.1,
         )
         logger.info("Using OpenAI-compatible model with request-scoped gateway support")

agent_service/app/presentation/routes/agent_router.py

@@ -179,7 +179,7 @@ async def upload_team_files(
 
 async def execute_agent_team(task_id: str, request: RunTeamRequest, owner_user_id: str, llm_gateway_token: str = ""):
     logger.info("[TASK %s] STARTED | source_path=%r owner=%r", task_id, request.source_path, owner_user_id)
-    update_task(task_id, status=AgentConstants.TASK_STATUS_RUNNING)
+    create_task(task_id, user_id=owner_user_id, status=AgentConstants.TASK_STATUS_RUNNING)
     # Connect workflow progress events to this persisted task so the UI
     # receives live updates through GET /v1/tasks/{task_id}.
     bind_task(task_id)
@@ -337,6 +337,21 @@ async def download_migration_v1(migration_name: str, user=Depends(get_current_us
     )
 
 
+def _require_llm_gateway_token(token: str | None) -> str:
+    """Require the portfolio-issued short-lived LLM gateway JWT."""
+    value = (token or "").strip()
+    if not value:
+        raise HTTPException(
+            status_code=401,
+            detail={
+                "code": "LLM_GATEWAY_SESSION_REQUIRED",
+                "message": "A Portfolio LLM Gateway session token is required for AI operations.",
+                "header": "X-LLM-Gateway-Token",
+            },
+        )
+    return value
+
+
 @router.post("/v1/teams/run", response_model=TaskAcceptedResponse)
 async def run_agent_team(
     request: RunTeamRequest,
@@ -345,16 +360,14 @@ async def run_agent_team(
     user=Depends(get_current_user),
 ):
     current_user.set(user)
+    gateway_token = _require_llm_gateway_token(x_llm_gateway_token)
     task_id = str(uuid.uuid4())
     logger.info("POST /v1/teams/run | task_id=%s", task_id)
-    # Persist the task before returning 202 so an immediate UI poll cannot
-    # race the BackgroundTasks startup and receive a false 404.
-    create_task(task_id, user_id=str(user.id), status=AgentConstants.TASK_STATUS_ACCEPTED)
 
     # Stack detection uses an LLM and can take longer than the UI's request
     # timeout.  Do it inside execute_agent_team, where it belongs, after this
     # endpoint has returned a task ID.
-    background_tasks.add_task(execute_agent_team, task_id, request, str(user.id), (x_llm_gateway_token or "").strip())
+    background_tasks.add_task(execute_agent_team, task_id, request, str(user.id), gateway_token)
     return TaskAcceptedResponse(
         task_id=task_id,
         message=AgentConstants.AGENT_TEAM_EXECUTION_QUEUED,
@@ -403,6 +416,7 @@ async def health_check():
 @router.post("/v1/chat/ask", response_model=ChatAskResponse)
 async def chat_ask(
     request: ChatAskRequest,
+    x_llm_gateway_token: str | None = Header(default=None, alias="X-LLM-Gateway-Token"),
     user=Depends(get_current_user),
 ) -> ChatAskResponse:
     try:
@@ -411,6 +425,7 @@ async def chat_ask(
         except Exception:
             pass
         _set_migration_context(request.migration_name, user)
+        set_llm_gateway_token(_require_llm_gateway_token(x_llm_gateway_token))
         from app.application.agents.chat.chat_tools import ask_kb_impl
 
         ask_kb_callable = _resolve_tool_callable(ask_kb_impl)
@@ -428,6 +443,7 @@ async def chat_ask(
 @router.post("/v1/report/migration")
 async def migration_report(
     request: MigrationReportRequest,
+    x_llm_gateway_token: str | None = Header(default=None, alias="X-LLM-Gateway-Token"),
     user=Depends(get_current_user),
 ):
     try:
@@ -436,6 +452,7 @@ async def migration_report(
         except Exception:
             pass
         _set_migration_context(request.migration_name, user)
+        set_llm_gateway_token(_require_llm_gateway_token(x_llm_gateway_token))
         from app.infrastructure.utils.reporting_manager import generate_migration_comparison_report
 
         report_callable = _resolve_tool_callable(generate_migration_comparison_report)
@@ -454,6 +471,7 @@ async def migration_report(
 @router.post("/v1/showcase/migration")
 async def migration_showcase(
     request: MigrationShowcaseRequest,
+    x_llm_gateway_token: str | None = Header(default=None, alias="X-LLM-Gateway-Token"),
     user=Depends(get_current_user),
 ):
     try:
@@ -462,6 +480,7 @@ async def migration_showcase(
         except Exception:
             pass
         _set_migration_context(request.migration_name, user)
+        set_llm_gateway_token(_require_llm_gateway_token(x_llm_gateway_token))
         from app.infrastructure.utils.showcase_manager import generate_showcase_bundle
 
         showcase_callable = _resolve_tool_callable(generate_showcase_bundle)
@@ -477,6 +496,7 @@ def _run_post_migration_background(
     migrated_code_path,
     persist: bool,
     user,
+    llm_gateway_token: str = "",
 ) -> None:
     """Background task: sets context vars and runs the full post-migration workflow."""
     try:
@@ -485,6 +505,7 @@ def _run_post_migration_background(
         pass
     try:
         _set_migration_context(migration_name, user)
+        set_llm_gateway_token((llm_gateway_token or "").strip())
     except Exception:
         pass
     try:
@@ -505,18 +526,21 @@ def _run_post_migration_background(
 async def run_post_migration(
     request: PostMigrationRunRequest,
     background_tasks: BackgroundTasks,
+    x_llm_gateway_token: str | None = Header(default=None, alias="X-LLM-Gateway-Token"),
     user=Depends(get_current_user),
 ):
     try:
         current_user.set(user)
     except Exception:
         pass
+    gateway_token = _require_llm_gateway_token(x_llm_gateway_token)
     background_tasks.add_task(
         _run_post_migration_background,
         request.migration_name,
         request.migrated_code_path,
         bool(request.persist),
         user,
+        gateway_token,
     )
     return {
         "status": "accepted",

render.yaml

@@ -15,16 +15,14 @@ services:
         value: "10000"
       - key: MODEL_TYPE
         value: "OpenAI"
-      - key: OPENAI_BASE_URL
-        value: "https://openrouter.ai/api/v1"
       - key: OPENAI_MODEL_ID
-        value: "openai/gpt-4o-mini"
+        value: "gateway-managed"
       - key: LLM_GATEWAY_URL
         value: "https://portfolio-llm-gateway.onrender.com/v1"
       - key: LLM_GATEWAY_TIMEOUT
         value: "180"
-      - key: LLM_MODEL
-        value: "gemini-3.5-flash-lite"
+      - key: LLM_GATEWAY_REQUIRED
+        value: "true"
       - key: QDRANT_URL
         sync: false
       - key: QDRANT_API_KEY
@@ -39,8 +37,6 @@ services:
         value: "qdrant/bm25"
       - key: QDRANT_DENSE_DIMENSIONS
         value: "384"
-      - key: OPENAI_API_KEY
-        sync: false
       - key: DATABASE_URL
         sync: false
       - key: ALLOWED_ORIGINS

tests/test_gateway_jwt_enforcement.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+MODEL_PROVIDER = ROOT / "agent_service/app/infrastructure/agents_backend/model_provider.py"
+ROUTER = ROOT / "agent_service/app/presentation/routes/agent_router.py"
+RENDER = ROOT / "render.yaml"
+
+
+def test_render_is_gateway_only():
+    text = RENDER.read_text(encoding="utf-8")
+    assert 'LLM_GATEWAY_URL' in text
+    assert 'LLM_GATEWAY_REQUIRED' in text
+    assert 'OPENAI_BASE_URL' not in text
+    assert 'OPENAI_API_KEY' not in text
+    assert 'LLM_MODEL' not in text
+
+
+def test_model_provider_has_no_direct_openai_fallback():
+    text = MODEL_PROVIDER.read_text(encoding="utf-8")
+    assert 'return super().get_client()' not in text
+    assert 'return super().get_async_client()' not in text
+    assert 'LLM gateway session token is required' in text
+    assert 'api_key="gateway-session"' in text
+
+
+def test_ai_routes_require_and_propagate_gateway_token():
+    text = ROUTER.read_text(encoding="utf-8")
+    assert text.count('alias="X-LLM-Gateway-Token"') >= 5
+    assert '_require_llm_gateway_token' in text
+    assert 'background_tasks.add_task(execute_agent_team, task_id, request, str(user.id), gateway_token)' in text
+    assert 'gateway_token,' in text