refactor(mcp): unify OAuth flow resolution and distinguish missing params from expired state (#5598)

routes/mcp/mcp_routes.py

@@ -114,6 +114,18 @@ def _mcp_oauth_redirect_uri() -> str:
     return REDIRECT_URI
 
 
+def _is_legacy_google_flow(identifier: str) -> bool:
+    """Determine whether an identifier is the server_id of a configured Google OAuth MCP server."""
+    if not identifier:
+        return False
+    db = SessionLocal()
+    try:
+        srv = db.query(McpServer).filter(McpServer.id == identifier).first()
+        return bool(srv and srv.oauth_config)
+    finally:
+        db.close()
+
+
 def setup_mcp_routes(mcp_manager: McpManager):
     """Setup MCP routes with the provided manager."""
 
@@ -490,24 +502,23 @@ def setup_mcp_routes(mcp_manager: McpManager):
     async def oauth_callback(code: str = "", state: str = "", request: Request = None):
         """Handle OAuth callback. Generic MCP OAuth flows resolve via the
         pending-state registry; Google flows fall through to the legacy path."""
+        if not state:
+            return HTMLResponse(_oauth_result_page("Error", "No OAuth state parameter found in the callback URL."), status_code=400)
+        if not code:
+            return HTMLResponse(_oauth_result_page("Error", "No authorization code found in the URL. Make sure you copied the full URL from your browser."), status_code=400)
+
         from src.mcp_oauth import resolve_pending
-        if state and code and resolve_pending(state, code):
+        if resolve_pending(state, code):
             return HTMLResponse(_oauth_result_page(
                 "Authorization Successful",
                 "The MCP server is connecting. You can close this window and return to Odysseus.",
                 success=True,
             ))
 
-        # Legacy Google path: state is the server_id
+        # Check if this callback belongs to a legacy Google flow vs expired generic flow
-        if state:
+        if _is_legacy_google_flow(state):
-            db = SessionLocal()
+            require_admin(request)
-            try:
+            return await _exchange_and_connect(state, code, request)
-                srv = db.query(McpServer).filter(McpServer.id == state).first()
-                if srv and srv.oauth_config:
-                    require_admin(request)
-                    return await _exchange_and_connect(state, code, request)
-            finally:
-                db.close()
 
         return HTMLResponse(_oauth_result_page(
             "Authorization Expired",
@@ -527,31 +538,28 @@ def setup_mcp_routes(mcp_manager: McpManager):
         except Exception:
             return HTMLResponse(_oauth_result_page("Error", "Invalid URL format."), status_code=400)
 
+        state = params.get("state", [None])[0]
+        if not state:
+            return HTMLResponse(_oauth_result_page("Error", "No OAuth state parameter found in the callback URL."), status_code=400)
+
         # Generic MCP OAuth: if the pasted URL carries a state we are waiting on,
         # resolve it directly (the background connect finishes the handshake).
-        state = params.get("state", [None])[0]
         from src.mcp_oauth import resolve_pending
-        if state and resolve_pending(state, code):
+        if resolve_pending(state, code):
             return HTMLResponse(_oauth_result_page(
                 "Authorization Successful",
                 "The MCP server is connecting. You can close this window and return to Odysseus.",
                 success=True,
             ))
 
-        db = SessionLocal()
+        # Check if this exchange belongs to a legacy Google flow vs expired generic flow
-        try:
+        if _is_legacy_google_flow(server_id):
-            srv = db.query(McpServer).filter(McpServer.id == server_id).first()
+            return await _exchange_and_connect(server_id, code, request)
-            if not srv:
-                return HTMLResponse(_oauth_result_page("Error", "Server not found."), status_code=404)
-            if not srv.oauth_config:
-                return HTMLResponse(_oauth_result_page(
-                    "Authorization Expired",
-                    "This authorization round has expired. Press Reconnect in Odysseus and try again.",
-                ), status_code=400)
-        finally:
-            db.close()
 
-        return await _exchange_and_connect(server_id, code, request)
+        return HTMLResponse(_oauth_result_page(
+            "Authorization Expired",
+            "This authorization round has expired. Press Reconnect in Odysseus and try again.",
+        ), status_code=400)
 
     async def _exchange_and_connect(server_id: str, code: str, request: Request):
         """Exchange auth code for tokens and connect the MCP server."""

tests/test_mcp_oauth.py

@@ -450,3 +450,45 @@ def test_oauth_exchange_expired_generic_state_returns_expired_page(monkeypatch):
         in resp.text
     )
     assert "No OAuth config." not in resp.text
+
+
+def test_oauth_callback_missing_state_returns_malformed_error_not_expired():
+    """Missing state parameter in callback must return a bad-request error and must
+    not be treated as an expired authorization round."""
+    from unittest.mock import MagicMock
+    from fastapi import FastAPI
+    from fastapi.testclient import TestClient
+    from routes.mcp import mcp_routes
+
+    app = FastAPI()
+    manager = MagicMock()
+    router = mcp_routes.setup_mcp_routes(manager)
+    app.include_router(router)
+    client = TestClient(app)
+
+    resp = client.get("/api/mcp/oauth/callback?code=some-code")
+    assert resp.status_code == 400
+    assert "No OAuth state parameter found in the callback URL." in resp.text
+    assert "Authorization Expired" not in resp.text
+    assert "This authorization round has expired" not in resp.text
+
+
+def test_oauth_callback_missing_code_returns_malformed_error_not_expired():
+    """Missing code parameter in callback must return a bad-request error and must
+    not be treated as an expired authorization round."""
+    from unittest.mock import MagicMock
+    from fastapi import FastAPI
+    from fastapi.testclient import TestClient
+    from routes.mcp import mcp_routes
+
+    app = FastAPI()
+    manager = MagicMock()
+    router = mcp_routes.setup_mcp_routes(manager)
+    app.include_router(router)
+    client = TestClient(app)
+
+    resp = client.get("/api/mcp/oauth/callback?state=some-state")
+    assert resp.status_code == 400
+    assert "No authorization code found in the URL." in resp.text
+    assert "Authorization Expired" not in resp.text
+    assert "This authorization round has expired" not in resp.text