From 0d42916d6774ff175538c0be65d9fd3b04b2d0ac Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Mon, 28 Sep 2026 23:32:21 +0100 Subject: [PATCH] Hand the MCP SDK an authorization result with the issuer mcp 2.x reads code, state and the RFC 9207 iss off the callback handler's result; the handler still returned a (code, state) tuple, so every MCP sign-in failed after the user approved it. The callback route now keeps the issuer too, since servers that advertise it (Linear) refuse a sign-in without it. --- docsgpt/agents/tools/mcp_tool.py | 33 +++++++++++++++++++++++----- docsgpt/api/user/tools/mcp.py | 2 +- tests/agents/tools/test_mcp_tool.py | 24 +++++++++++++++----- tests/api/user/test_tools_mcp_pg.py | 20 +++++++++++++++++ tests/connectors/test_mcp_refresh.py | 30 +++++++++++++++++++++++++ 5 files changed, 97 insertions(+), 12 deletions(-) diff --git a/docsgpt/agents/tools/mcp_tool.py b/docsgpt/agents/tools/mcp_tool.py index c1d8a888..23d77d72 100644 --- a/docsgpt/agents/tools/mcp_tool.py +++ b/docsgpt/agents/tools/mcp_tool.py @@ -15,7 +15,12 @@ from fastmcp.client.transports import ( StreamableHttpTransport, ) from mcp.client.auth import OAuthClientProvider, TokenStorage -from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken +from mcp.shared.auth import ( + AuthorizationCodeResult, + OAuthClientInformationFull, + OAuthClientMetadata, + OAuthToken, +) from pydantic import AnyHttpUrl, ValidationError from redis import Redis @@ -811,8 +816,13 @@ class DocsGPTOAuth(OAuthClientProvider): exc_info=True, ) - async def callback_handler(self) -> tuple[str, str | None]: - """Wait for auth code from Redis using the state value.""" + async def callback_handler(self) -> AuthorizationCodeResult: + """Wait for auth code from Redis using the state value. + + Returns: + The code, the state it came back with, and the RFC 9207 issuer + when the authorization server sent one. + """ if not self.redis_client or not self.extracted_state: raise Exception("Redis client or state not configured for OAuth") poll_interval = 1 @@ -825,15 +835,22 @@ class DocsGPTOAuth(OAuthClientProvider): if code_data: code = code_data.decode() returned_state = self.extracted_state + iss_key = f"{self.redis_prefix}iss:{self.extracted_state}" + iss_data = self.redis_client.get(iss_key) self.redis_client.delete(code_key) + self.redis_client.delete(iss_key) self.redis_client.delete( f"{self.redis_prefix}auth_url:{self.extracted_state}" ) self.redis_client.delete( f"{self.redis_prefix}state:{self.extracted_state}" ) - return code, returned_state + return AuthorizationCodeResult( + code=code, + state=returned_state, + iss=iss_data.decode() if iss_data else None, + ) error_key = f"{self.redis_prefix}error:{self.extracted_state}" error_data = self.redis_client.get(error_key) if error_data: @@ -869,7 +886,7 @@ class NonInteractiveOAuth(DocsGPTOAuth): "OAuth session expired — please re-authorize this MCP server in tool settings." ) - async def callback_handler(self) -> tuple[str, str | None]: + async def callback_handler(self) -> AuthorizationCodeResult: raise Exception( "OAuth session expired — please re-authorize this MCP server in tool settings." ) @@ -1037,7 +1054,7 @@ class MCPOAuthManager: self.redis_prefix = redis_prefix def handle_oauth_callback( - self, state: str, code: str, error: Optional[str] = None + self, state: str, code: str, error: Optional[str] = None, iss: Optional[str] = None ) -> bool: """ Handle OAuth callback from provider. @@ -1046,6 +1063,7 @@ class MCPOAuthManager: state: The state parameter from OAuth callback code: The authorization code from OAuth callback error: Error message if OAuth failed + iss: The RFC 9207 issuer from the callback, if the server sent one Returns: True if successful, False otherwise @@ -1057,6 +1075,9 @@ class MCPOAuthManager: error_key = f"{self.redis_prefix}error:{state}" self.redis_client.setex(error_key, 300, error) raise Exception(f"OAuth error received: {error}") + # The issuer goes first: the waiting sign-in reads it once the code lands. + if iss: + self.redis_client.setex(f"{self.redis_prefix}iss:{state}", 300, iss) code_key = f"{self.redis_prefix}code:{state}" self.redis_client.setex(code_key, 300, code) diff --git a/docsgpt/api/user/tools/mcp.py b/docsgpt/api/user/tools/mcp.py index d443a800..8738d579 100644 --- a/docsgpt/api/user/tools/mcp.py +++ b/docsgpt/api/user/tools/mcp.py @@ -539,7 +539,7 @@ class MCPOAuthCallback(Resource): "/api/connectors/callback-status?status=error&message=Internal+server+error:+Redis+not+available.&provider=mcp_tool" ) manager = MCPOAuthManager(redis_client) - success = manager.handle_oauth_callback(state, code, error) + success = manager.handle_oauth_callback(state, code, error, iss=request.args.get("iss")) if success: return redirect( "/api/connectors/callback-status?status=success&message=Authorization+code+received+successfully.+You+can+close+this+window.&provider=mcp_tool" diff --git a/tests/agents/tools/test_mcp_tool.py b/tests/agents/tools/test_mcp_tool.py index 5eb30465..6104e28f 100644 --- a/tests/agents/tools/test_mcp_tool.py +++ b/tests/agents/tools/test_mcp_tool.py @@ -840,6 +840,15 @@ class TestMCPOAuthManager: assert result is True mock_redis.setex.assert_called() + def test_handle_callback_keeps_the_issuer(self): + from docsgpt.agents.tools.mcp_tool import MCPOAuthManager + + mock_redis = MagicMock() + manager = MCPOAuthManager(mock_redis) + + assert manager.handle_oauth_callback(state="s", code="c", iss="https://issuer.example") is True + mock_redis.setex.assert_any_call("mcp_oauth:iss:s", 300, "https://issuer.example") + def test_handle_callback_no_redis(self): from docsgpt.agents.tools.mcp_tool import MCPOAuthManager @@ -2026,8 +2035,11 @@ class TestDocsGPTOAuthExtended: mock_db.__getitem__ = MagicMock(return_value=mock_collection) mock_redis = MagicMock() - # First get returns the code - mock_redis.get.return_value = b"auth_code_123" + stored = { + "mcp_oauth:code:mystate": b"auth_code_123", + "mcp_oauth:iss:mystate": b"https://auth.example.com", + } + mock_redis.get.side_effect = stored.get oauth = DocsGPTOAuth( mcp_url="https://mcp.example.com/api", @@ -2043,9 +2055,11 @@ class TestDocsGPTOAuthExtended: loop = asyncio.new_event_loop() try: - code, state = loop.run_until_complete(oauth.callback_handler()) - assert code == "auth_code_123" - assert state == "mystate" + # The MCP SDK reads ``.code``, ``.state`` and the RFC 9207 ``.iss``. + result = loop.run_until_complete(oauth.callback_handler()) + assert result.code == "auth_code_123" + assert result.state == "mystate" + assert result.iss == "https://auth.example.com" finally: loop.close() diff --git a/tests/api/user/test_tools_mcp_pg.py b/tests/api/user/test_tools_mcp_pg.py index 3718763e..0fc602fb 100644 --- a/tests/api/user/test_tools_mcp_pg.py +++ b/tests/api/user/test_tools_mcp_pg.py @@ -502,3 +502,23 @@ class TestMCPAuthStatus: request.decoded_token = None response = MCPAuthStatus().get() assert response.status_code == 401 + + +class TestMCPOAuthCallbackIssuer: + def test_passes_the_issuer_to_the_waiting_sign_in(self, app): + from docsgpt.api.user.tools.mcp import MCPOAuthCallback + + manager = MagicMock() + manager.handle_oauth_callback.return_value = True + with patch( + "docsgpt.api.user.tools.mcp.get_redis_instance", return_value=MagicMock(), + ), patch( + "docsgpt.api.user.tools.mcp.MCPOAuthManager", return_value=manager, + ), app.test_request_context( + "/api/mcp_server/callback?code=c&state=s&iss=https%3A%2F%2Fmcp.linear.app", + ): + response = MCPOAuthCallback().get() + assert response.status_code == 302 + manager.handle_oauth_callback.assert_called_once_with( + "s", "c", None, iss="https://mcp.linear.app", + ) diff --git a/tests/connectors/test_mcp_refresh.py b/tests/connectors/test_mcp_refresh.py index c5492498..50f21edc 100644 --- a/tests/connectors/test_mcp_refresh.py +++ b/tests/connectors/test_mcp_refresh.py @@ -142,3 +142,33 @@ def test_oauth_resource_is_the_full_mcp_endpoint(): assert check_resource_allowed(requested_resource=requested, configured_resource="https://mcp.linear.app") # Stored tokens stay keyed by the server's origin. assert oauth.context.storage.server_url == "https://mcp.linear.app" + + +def test_callback_hands_the_sdk_an_authorization_result(): + """The MCP SDK reads ``.code``, ``.state`` and the RFC 9207 ``.iss`` off the + callback's result; Linear advertises ``iss`` and refuses a sign-in without it.""" + import asyncio + + from mcp.shared.auth import AuthorizationCodeResult + + from docsgpt.agents.tools.mcp_tool import DocsGPTOAuth + + stored = { + "mcp_oauth:code:st": b"the-code", + "mcp_oauth:iss:st": b"https://mcp.linear.app", + } + redis_client = MagicMock() + redis_client.get.side_effect = stored.get + oauth = DocsGPTOAuth( + mcp_url="https://mcp.linear.app/mcp", + redis_client=redis_client, + redirect_uri="https://example.com/callback", + user_id="user1", + ) + oauth.extracted_state = "st" + + result = asyncio.run(oauth.callback_handler()) + + assert isinstance(result, AuthorizationCodeResult) + assert (result.code, result.state, result.iss) == ("the-code", "st", "https://mcp.linear.app") + redis_client.delete.assert_any_call("mcp_oauth:iss:st")