diff --git a/routes/mcp/mcp_routes.py b/routes/mcp/mcp_routes.py index 385cf09e3..5d14ba8c3 100644 --- a/routes/mcp/mcp_routes.py +++ b/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 - if state: - db = SessionLocal() - try: - 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() + # Check if this callback belongs to a legacy Google flow vs expired generic flow + if _is_legacy_google_flow(state): + require_admin(request) + return await _exchange_and_connect(state, code, request) 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() - try: - srv = db.query(McpServer).filter(McpServer.id == server_id).first() - 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() + # Check if this exchange belongs to a legacy Google flow vs expired generic flow + if _is_legacy_google_flow(server_id): + return await _exchange_and_connect(server_id, code, request) - 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.""" diff --git a/tests/test_mcp_oauth.py b/tests/test_mcp_oauth.py index 24f52820e..a0ba8a6a2 100644 --- a/tests/test_mcp_oauth.py +++ b/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