From b8fb1feb1e93c09955ffd82a5fe209b7f62aadcd Mon Sep 17 00:00:00 2001 From: GitHub Action Date: Mon, 20 Jul 2026 16:55:28 -0700 Subject: [PATCH 1/4] fix(keycardai-mcp): run OAuth completion cleanup synchronously before returning Completion cleanup (PKCE state delete, pending-auth clear, and the completion-route delete in CompletionRouter) previously ran via bare asyncio.create_task on the background path. The event loop keeps only a weak reference to such tasks, so they could be garbage-collected before running (or cancelled at shutdown with the CancelledError swallowed), leaving a stale pending-auth record that connected sessions kept advertising as an auth challenge. Both cleanup steps are fast storage calls, so they now always run synchronously before the completion result is returned, for every coordinator: - oauth_completion_handler awaits its cleanup unconditionally; the run_cleanup_in_background parameter is deprecated and ignored (a DeprecationWarning fires when it is passed explicitly) - CompletionRouter.route_completion awaits the route-metadata delete instead of scheduling a fire-and-forget task - a failure deleting the PKCE state no longer skips clearing the pending-auth record; each cleanup step is attempted independently - CancelledError is no longer swallowed by cleanup helpers - AuthCoordinator.requires_synchronous_cleanup is deprecated (kept for backward compatibility, returns True); OAuthStrategy no longer consults it and LocalAuthCoordinator's override is removed Co-Authored-By: Claude Fable 5 --- .../src/keycardai/mcp/client/CONTRIBUTORS.md | 2 +- .../mcp/client/auth/coordinators/base.py | 11 +- .../mcp/client/auth/coordinators/local.py | 11 -- .../src/keycardai/mcp/client/auth/handlers.py | 80 ++++---- .../mcp/client/auth/strategies/oauth.py | 12 +- .../mcp/client/auth/strategies/test_oauth.py | 23 +-- .../mcp/client/auth/test_handlers.py | 173 ++++++++++++++++++ .../client/auth/test_stateless_completions.py | 6 - 8 files changed, 236 insertions(+), 82 deletions(-) diff --git a/packages/mcp/src/keycardai/mcp/client/CONTRIBUTORS.md b/packages/mcp/src/keycardai/mcp/client/CONTRIBUTORS.md index c133843a..056a2ad6 100644 --- a/packages/mcp/src/keycardai/mcp/client/CONTRIBUTORS.md +++ b/packages/mcp/src/keycardai/mcp/client/CONTRIBUTORS.md @@ -1408,7 +1408,7 @@ Clear boundaries help maintain clean architecture. - Discovery and dynamic client registration - PKCE flow with state validation - Stores PKCE params in storage (works across processes) -- Adapts cleanup behavior based on `coordinator.requires_synchronous_cleanup` +- Cleans up flow state synchronously before the completion result is returned - Works with both LocalAuthCoordinator and StarletteAuthCoordinator **ApiKeyStrategy** (subclass): diff --git a/packages/mcp/src/keycardai/mcp/client/auth/coordinators/base.py b/packages/mcp/src/keycardai/mcp/client/auth/coordinators/base.py index 797e0d81..bdf26c69 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/coordinators/base.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/coordinators/base.py @@ -92,14 +92,15 @@ def requires_synchronous_cleanup(self) -> bool: """ Whether this coordinator requires synchronous callback cleanup. - Override in subclasses if they need cleanup to complete before - callback response (e.g., LocalAuthCoordinator needs this to avoid - race conditions with blocking wait patterns). + Deprecated: completion cleanup always runs synchronously before the + completion result is returned, for every coordinator. The SDK no + longer consults this property; it is kept for backward compatibility + with code that reads or overrides it. Returns: - False by default (asynchronous cleanup is fine) + True (cleanup is always synchronous) """ - return False + return True def create_context(self, context_id: str, metadata: dict[str, Any] | None = None) -> Context: """ diff --git a/packages/mcp/src/keycardai/mcp/client/auth/coordinators/local.py b/packages/mcp/src/keycardai/mcp/client/auth/coordinators/local.py index 3d44d36d..f180802a 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/coordinators/local.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/coordinators/local.py @@ -21,7 +21,6 @@ class LocalAuthCoordinator(AuthCoordinator): Key behaviors: - Opens browser to authorization URL (configurable) - Blocks in handle_redirect() until completion arrives (configurable) - - Requires synchronous cleanup to avoid race conditions (when blocking) """ def __init__( @@ -65,16 +64,6 @@ def endpoint_type(self) -> str: """Type of endpoint: local HTTP server.""" return "local" - @property - def requires_synchronous_cleanup(self) -> bool: - """ - LocalAuthCoordinator requires synchronous cleanup. - - This prevents race conditions with the blocking wait pattern - in handle_redirect() which waits for completion to arrive. - """ - return True - async def shutdown(self) -> None: """ Stop local completion server. diff --git a/packages/mcp/src/keycardai/mcp/client/auth/handlers.py b/packages/mcp/src/keycardai/mcp/client/auth/handlers.py index 1444bd3f..dfe3832d 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/handlers.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/handlers.py @@ -1,6 +1,6 @@ """Auth completion routing and handlers.""" -import asyncio +import warnings from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any @@ -55,8 +55,8 @@ async def route_completion( Route completion to appropriate handler. Retrieves routing metadata from storage, invokes the registered handler, - and schedules cleanup. The coordinator is passed to handlers for context - creation and other coordinator operations. + and cleans up the routing metadata before returning. The coordinator is + passed to handlers for context creation and other coordinator operations. Args: coordinator: AuthCoordinator instance @@ -83,9 +83,10 @@ async def route_completion( completion_metadata ) - # Schedule cleanup as background task (don't block HTTP response) - # Note: Cleanup happens asynchronously after response is sent - asyncio.create_task(self._cleanup_completion_route(state)) + # Clean up routing metadata before returning so no stale route survives + # the completion response. The delete is a fast storage call; a bare + # background task here can be garbage-collected before it ever runs. + await self._cleanup_completion_route(state) logger.info(f"Completion handled successfully for state: {state[:8]}...") @@ -171,10 +172,11 @@ def _navigate_to_namespace( async def _cleanup_completion_route(self, state: str) -> None: """ - Clean up completion routing metadata as background task. + Clean up completion routing metadata. - This runs asynchronously after the HTTP response is sent to avoid - blocking the completion response. Handles cancellation gracefully. + Runs before the completion result is returned. Storage errors are + logged and not propagated because the completion itself already + succeeded; cancellation propagates normally. Args: state: OAuth state parameter to clean up @@ -182,11 +184,8 @@ async def _cleanup_completion_route(self, state: str) -> None: try: await self.state_storage.delete_completion_route(state) logger.debug(f"Cleaned up completion route for state: {state[:8]}...") - except asyncio.CancelledError: - # This is expected if the task is cancelled during shutdown - logger.debug(f"Completion cleanup cancelled for state: {state[:8]}...") except Exception as e: - # Log but don't propagate errors from background cleanup + # Log but don't propagate cleanup errors - the completion succeeded logger.warning(f"Failed to cleanup completion route for state {state[:8]}...: {e}") @@ -449,7 +448,7 @@ async def oauth_completion_handler( storage: "NamespacedStorage", params: dict[str, str], client_factory: Callable[[], "AsyncClient"] | None = None, - run_cleanup_in_background: bool = True + run_cleanup_in_background: bool | None = None ) -> dict[str, Any]: """ Handle OAuth authorization code completion. @@ -462,7 +461,7 @@ async def oauth_completion_handler( 2. Load PKCE state from storage 3. Exchange authorization code for tokens 4. Store tokens in strategy storage - 5. Clear pending auth state + 5. Clear PKCE and pending auth state (synchronously, before returning) 6. Return success result Args: @@ -471,9 +470,10 @@ async def oauth_completion_handler( params: Completion parameters (e.g., {"code": "...", "state": "..."}) client_factory: Optional factory function that returns an AsyncClient instance. If None, uses default factory that creates AsyncClient() - run_cleanup_in_background: If True (default), cleanup tasks run asynchronously - after response. If False, cleanup runs synchronously. - Use False in unit tests to ensure cleanup completes. + run_cleanup_in_background: Deprecated and ignored. Cleanup always runs + synchronously before the result is returned, so + a completed authorization never leaves a stale + pending-auth record behind. Returns: Result dict with success status and metadata @@ -498,6 +498,15 @@ def my_factory(): coordinator, storage, params, client_factory=my_factory ) """ + if run_cleanup_in_background is not None: + warnings.warn( + "run_cleanup_in_background is deprecated and has no effect: " + "OAuth completion cleanup always runs synchronously before the " + "result is returned.", + DeprecationWarning, + stacklevel=2, + ) + code = params.get("code") state = params.get("state") @@ -555,16 +564,14 @@ def my_factory(): namespace_parts = storage._namespace.split(":") context_id = namespace_parts[1] if len(namespace_parts) > 1 else "unknown" - # Cleanup: Run as background task (async) or synchronously based on parameter + # Cleanup runs synchronously before returning: both deletes are fast storage + # calls, and a fire-and-forget task can be garbage-collected (or cancelled at + # shutdown) before it runs, leaving a stale pending-auth record that sessions + # would keep advertising as an auth challenge. # Note: completion metadata cleanup is handled by coordinator - if run_cleanup_in_background: - asyncio.create_task(_cleanup_oauth_completion_state( - storage, coordinator, state, context_id, server_name - )) - else: - await _cleanup_oauth_completion_state( - storage, coordinator, state, context_id, server_name - ) + await _cleanup_oauth_completion_state( + storage, coordinator, state, context_id, server_name + ) logger.info(f"OAuth completion completed for {server_name}") @@ -579,10 +586,14 @@ async def _cleanup_oauth_completion_state( server_name: str ) -> None: """ - Clean up OAuth completion state as a background task. + Clean up transient OAuth flow state after a successful completion. - This runs asynchronously after the HTTP response is sent to avoid - blocking the completion response. + Runs before the completion result is returned so a completed authorization + never leaves a stale pending-auth record behind. Each step is attempted + independently: a failure deleting the PKCE state must not skip clearing + the pending-auth record. Storage errors are logged and not propagated + because the token exchange already succeeded; cancellation propagates + normally. Args: storage: Strategy storage @@ -593,11 +604,12 @@ async def _cleanup_oauth_completion_state( """ try: await storage.delete(f"_pkce_state:{state}") + except Exception as e: + logger.warning(f"Failed to delete PKCE state for {server_name}: {e}") + + try: await coordinator.clear_auth_pending(context_id=context_id, server_name=server_name) logger.debug(f"Cleaned up OAuth completion state for {server_name}") - except asyncio.CancelledError: - logger.debug(f"OAuth completion cleanup cancelled for {server_name}") except Exception as e: - # Log but don't propagate errors from background cleanup - logger.warning(f"Failed to cleanup OAuth completion state for {server_name}: {e}") + logger.warning(f"Failed to clear pending auth for {server_name}: {e}") diff --git a/packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py b/packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py index ac4628f6..c12e8204 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py @@ -272,15 +272,8 @@ async def handle_challenge( ) # Step 4: Register completion route for token exchange - # Coordinators that require synchronous cleanup (e.g., LocalAuthCoordinator - # with blocking wait patterns) will have it, others use background cleanup - handler_kwargs = {} - if self.coordinator.requires_synchronous_cleanup: - handler_kwargs["run_cleanup_in_background"] = False - logger.debug(f"[{self.server_name}] Using synchronous cleanup for {type(self.coordinator).__name__}") - else: - logger.debug(f"[{self.server_name}] Using background cleanup for {type(self.coordinator).__name__}") - + # Completion cleanup always runs synchronously in the handler, so no + # cleanup-mode kwargs are needed regardless of coordinator type await self.coordinator.register_completion_route( routing_key=flow_metadata.state, handler_name="oauth_completion", @@ -288,7 +281,6 @@ async def handle_challenge( context_id=self.context.id, server_name=self.server_name, routing_param="state", - handler_kwargs=handler_kwargs, metadata=self.context.metadata, ttl=timedelta(minutes=10) ) diff --git a/packages/mcp/tests/keycardai/mcp/client/auth/strategies/test_oauth.py b/packages/mcp/tests/keycardai/mcp/client/auth/strategies/test_oauth.py index de78c96e..16972d73 100644 --- a/packages/mcp/tests/keycardai/mcp/client/auth/strategies/test_oauth.py +++ b/packages/mcp/tests/keycardai/mcp/client/auth/strategies/test_oauth.py @@ -62,12 +62,6 @@ def __init__(self, redirect_uris: list[str] | None | object = None): # Mock storage for callback cleanup backend = InMemoryBackend() self.storage = NamespacedStorage(backend, "auth_coordinator") - self._requires_synchronous_cleanup = False - - @property - def requires_synchronous_cleanup(self) -> bool: - """Mock property for cleanup behavior.""" - return self._requires_synchronous_cleanup async def get_redirect_uris(self) -> list[str] | None: """Return mock redirect URIs.""" @@ -508,8 +502,12 @@ def registration_response(): assert "state" in challenge @pytest.mark.asyncio - async def test_handle_challenge_coordinator_cleanup_behavior(self): - """Test that handle_challenge respects coordinator cleanup requirements.""" + async def test_handle_challenge_registers_route_without_cleanup_kwargs(self): + """Test that handle_challenge registers no cleanup-mode handler kwargs. + + Completion cleanup always runs synchronously in the handler, so the + registered route must not carry a run_cleanup_in_background flag. + """ # Mock responses def resource_response(): @@ -551,9 +549,7 @@ def registration_response(): post_responses=[registration_response], ) - # Test with coordinator requiring synchronous cleanup coordinator = MockCoordinator() - coordinator._requires_synchronous_cleanup = True _, connection_storage, strategy_storage = create_test_storage() context = create_test_context(coordinator) @@ -576,11 +572,10 @@ def registration_response(): ) assert result is True - # Verify handler_kwargs includes synchronous cleanup flag + # The route must not carry any cleanup-mode flag state = list(coordinator.registered_routes.keys())[0] handler_kwargs = coordinator.registered_routes[state]["handler_kwargs"] - assert "run_cleanup_in_background" in handler_kwargs - assert handler_kwargs["run_cleanup_in_background"] is False + assert "run_cleanup_in_background" not in handler_kwargs class TestOAuthStrategyCallbackCompletion: @@ -737,7 +732,6 @@ def mock_token_response(): storage=strategy_storage, params={"code": "auth_code_123", "state": state}, client_factory=client_factory, - run_cleanup_in_background=False ) assert result["success"] is True @@ -840,7 +834,6 @@ def token_response(): storage=strategy_storage, params={"code": "auth_code_123", "state": state}, client_factory=client_factory, - run_cleanup_in_background=False ) # Verify final state diff --git a/packages/mcp/tests/keycardai/mcp/client/auth/test_handlers.py b/packages/mcp/tests/keycardai/mcp/client/auth/test_handlers.py index a74f45e2..27aa717d 100644 --- a/packages/mcp/tests/keycardai/mcp/client/auth/test_handlers.py +++ b/packages/mcp/tests/keycardai/mcp/client/auth/test_handlers.py @@ -169,6 +169,179 @@ async def test_oauth_completion_handler_missing_pkce_state(self): ) +def _token_client_factory(): + """Client factory returning a mock client whose POST yields a token response.""" + response = MagicMock() + response.status_code = 200 + response.json = MagicMock( + return_value={"access_token": "access_xyz", "expires_in": 3600} + ) + + client = MagicMock() + client.__aenter__ = AsyncMock(return_value=client) + client.__aexit__ = AsyncMock(return_value=None) + client.post = AsyncMock(return_value=response) + return client + + +async def _seed_pkce_state(storage: NamespacedStorage, state: str) -> None: + """Store the PKCE state the completion handler expects.""" + await storage.set( + f"_pkce_state:{state}", + { + "code_verifier": "verifier", + "redirect_uri": "http://localhost:8080/callback", + "client_id": "client_123", + "token_endpoint": "https://auth.example.com/token", + "server_name": "test_server", + }, + ) + + +class TestOAuthCompletionCleanup: + """Test that OAuth completion cleanup is synchronous and failure-tolerant.""" + + @pytest.mark.asyncio + async def test_cleanup_runs_before_handler_returns(self): + """PKCE state and pending auth are cleared before the handler returns. + + The assertions on the coordinator mock run synchronously after the + await, without yielding to the event loop, so they fail if cleanup + were deferred to a fire-and-forget task. + """ + coordinator = AsyncMock() + backend = InMemoryBackend() + storage = NamespacedStorage( + backend, "client:test_user:server:test_server:connection:oauth" + ) + state = "state_sync_cleanup" + await _seed_pkce_state(storage, state) + + result = await oauth_completion_handler( + coordinator, + storage, + {"code": "auth_code", "state": state}, + client_factory=_token_client_factory, + ) + + assert result["success"] is True + coordinator.clear_auth_pending.assert_awaited_once_with( + context_id="test_user", server_name="test_server" + ) + + assert await storage.get(f"_pkce_state:{state}") is None + tokens = await storage.get("tokens") + assert tokens is not None + assert tokens["access_token"] == "access_xyz" + + @pytest.mark.asyncio + async def test_pkce_delete_failure_still_clears_auth_pending(self): + """A failure deleting PKCE state must not skip clearing pending auth.""" + coordinator = AsyncMock() + backend = InMemoryBackend() + storage = NamespacedStorage( + backend, "client:test_user:server:test_server:connection:oauth" + ) + state = "state_pkce_failure" + await _seed_pkce_state(storage, state) + + storage.delete = AsyncMock(side_effect=RuntimeError("storage down")) + + result = await oauth_completion_handler( + coordinator, + storage, + {"code": "auth_code", "state": state}, + client_factory=_token_client_factory, + ) + + assert result["success"] is True + storage.delete.assert_awaited_once_with(f"_pkce_state:{state}") + coordinator.clear_auth_pending.assert_awaited_once_with( + context_id="test_user", server_name="test_server" + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("legacy_value", [True, False]) + async def test_run_cleanup_in_background_is_deprecated_and_ignored( + self, legacy_value + ): + """Passing run_cleanup_in_background warns and cleanup stays synchronous.""" + coordinator = AsyncMock() + backend = InMemoryBackend() + storage = NamespacedStorage( + backend, "client:test_user:server:test_server:connection:oauth" + ) + state = f"state_deprecated_{legacy_value}" + await _seed_pkce_state(storage, state) + + with pytest.warns(DeprecationWarning, match="run_cleanup_in_background"): + result = await oauth_completion_handler( + coordinator, + storage, + {"code": "auth_code", "state": state}, + client_factory=_token_client_factory, + run_cleanup_in_background=legacy_value, + ) + + assert result["success"] is True + coordinator.clear_auth_pending.assert_awaited_once_with( + context_id="test_user", server_name="test_server" + ) + assert await storage.get(f"_pkce_state:{state}") is None + + +class TestCompletionRouterCleanup: + """Test that CompletionRouter cleans up routing metadata synchronously.""" + + @pytest.mark.asyncio + async def test_route_metadata_deleted_before_completion_returns(self): + """The completion route is deleted before handle_completion returns.""" + from keycardai.mcp.client.auth.coordinators.base import AuthCoordinator + + class TestCoordinator(AuthCoordinator): + @property + def endpoint_type(self) -> str: + return "test" + + registry = CompletionHandlerRegistry() + + @registry.register("stub_completion") + async def stub_completion(coordinator, storage, params, **kwargs): + return {"success": True} + + coordinator = TestCoordinator( + backend=InMemoryBackend(), handler_registry=registry + ) + + state = "state_router_cleanup" + await coordinator.register_completion_route( + routing_key=state, + handler_name="stub_completion", + storage_namespace="client:test_user:server:test_server:connection:oauth", + context_id="test_user", + server_name="test_server", + ) + + state_storage = coordinator.completion_router.state_storage + deleted_routes: list[str] = [] + original_delete = state_storage.delete_completion_route + + async def recording_delete(routing_key: str) -> None: + deleted_routes.append(routing_key) + await original_delete(routing_key) + + state_storage.delete_completion_route = recording_delete + + result = await coordinator.handle_completion( + {"code": "auth_code", "state": state} + ) + + assert result["success"] is True + # Synchronous assertion, no event loop yield: the delete already ran + assert deleted_routes == [state] + assert await state_storage.get_completion_route(state) is None + + class TestClientFactoryManagement: """Test client factory management in the registry.""" diff --git a/packages/mcp/tests/keycardai/mcp/client/auth/test_stateless_completions.py b/packages/mcp/tests/keycardai/mcp/client/auth/test_stateless_completions.py index 4e375c52..98723780 100644 --- a/packages/mcp/tests/keycardai/mcp/client/auth/test_stateless_completions.py +++ b/packages/mcp/tests/keycardai/mcp/client/auth/test_stateless_completions.py @@ -22,12 +22,6 @@ def __init__(self): self.handle_redirect_calls = [] self.auth_pending = {} self.storage = NamespacedStorage(InMemoryBackend(), "auth_coordinator") - self._requires_synchronous_cleanup = False - - @property - def requires_synchronous_cleanup(self) -> bool: - """Mock property for cleanup behavior.""" - return self._requires_synchronous_cleanup async def get_redirect_uris(self) -> list[str]: """Return mock callback URIs.""" From fd6e5e3bc1d3090e397e5e009d938d828d4852a1 Mon Sep 17 00:00:00 2001 From: GitHub Action Date: Mon, 20 Jul 2026 16:59:19 -0700 Subject: [PATCH 2/4] fix(keycardai-mcp): never surface an auth challenge for a connected session Session.get_auth_challenge() was a pure storage read, so a pending-auth record that outlived a completed authorization (for example when cleanup did not run) kept being advertised as an auth challenge after the session auto-reconnected. Consumers of get_auth_challenges() had to filter it out themselves until the record's TTL expired, and custom storage backends that ignore ttl surfaced it indefinitely. A CONNECTED session has no pending challenge by definition: get_auth_challenge() now returns None in that state and lazily clears any stale stored record. All other statuses keep the storage read, so a session in AUTH_PENDING still surfaces its challenge. This is defense in depth on top of the synchronous completion cleanup; the lazy clear only helps in-process, which is fine for a secondary guard. Co-Authored-By: Claude Fable 5 --- .../mcp/src/keycardai/mcp/client/session.py | 21 +++++ .../tests/keycardai/mcp/client/test_client.py | 39 +++++++++ .../keycardai/mcp/client/test_session.py | 84 +++++++++++++++++++ 3 files changed, 144 insertions(+) diff --git a/packages/mcp/src/keycardai/mcp/client/session.py b/packages/mcp/src/keycardai/mcp/client/session.py index 90a4b078..47c12909 100644 --- a/packages/mcp/src/keycardai/mcp/client/session.py +++ b/packages/mcp/src/keycardai/mcp/client/session.py @@ -489,11 +489,32 @@ async def get_auth_challenge(self) -> dict[str, str] | None: An auth challenge is created by the auth strategy when authentication is required but not yet complete (e.g., waiting for OAuth callback). + A connected session has no pending challenge by definition, so this + returns None when the session status is CONNECTED. Any record still + stored for this session at that point is stale (e.g. completion + cleanup did not run) and is cleared lazily. + Returns: Dict with challenge details (strategy-specific) or None if no pending challenge. For OAuth: {'authorization_url': str, 'state': str} For other strategies: may contain different fields """ + if self.status == SessionStatus.CONNECTED: + stale_challenge = await self.coordinator.get_auth_pending( + context_id=self.context.id, + server_name=self.server_name + ) + if stale_challenge: + logger.debug( + f"Session {self.server_name}: Clearing stale auth challenge " + f"for connected session" + ) + await self.coordinator.clear_auth_pending( + context_id=self.context.id, + server_name=self.server_name + ) + return None + # Auth challenge is stored in coordinator return await self.coordinator.get_auth_pending( context_id=self.context.id, diff --git a/packages/mcp/tests/keycardai/mcp/client/test_client.py b/packages/mcp/tests/keycardai/mcp/client/test_client.py index 6982707a..844bbc27 100644 --- a/packages/mcp/tests/keycardai/mcp/client/test_client.py +++ b/packages/mcp/tests/keycardai/mcp/client/test_client.py @@ -1016,6 +1016,45 @@ async def test_get_auth_challenges_includes_server_name(self): assert challenges[0]["state"] == "abc123" assert challenges[0]["other_field"] == "value" + @pytest.mark.asyncio + async def test_get_auth_challenges_empty_for_connected_session_with_stale_record(self): + """Test that a connected session's stale pending-auth record is not surfaced. + + Simulates a completed authorization and auto-reconnect where the + completion cleanup task was dropped before running: the pending-auth + record is still in coordinator storage, but the session is CONNECTED. + get_auth_challenges() must return an empty list without any + consumer-side filtering. + """ + servers = {"test_server": {"url": "http://localhost:3000"}} + mock_coordinator = MockAuthCoordinator() + + client = Client(servers=servers, auth_coordinator=mock_coordinator) + + # Real Session (not a mock) so the CONNECTED short-circuit is exercised + session = client.sessions["test_server"] + + # Stale record left behind by a dropped cleanup task + await mock_coordinator.set_auth_pending( + context_id=client.context.id, + server_name="test_server", + auth_metadata={ + "authorization_url": "http://auth.example.com", + "state": "state123" + } + ) + session.status = SessionStatus.CONNECTED + + challenges = await client.get_auth_challenges() + + assert challenges == [] + # The stale record was cleared lazily + stored = await mock_coordinator.get_auth_pending( + context_id=client.context.id, + server_name="test_server" + ) + assert stored is None + class TestClientIdGeneration: """Test Client ID generation.""" diff --git a/packages/mcp/tests/keycardai/mcp/client/test_session.py b/packages/mcp/tests/keycardai/mcp/client/test_session.py index 71e36dcb..b2a5d36c 100644 --- a/packages/mcp/tests/keycardai/mcp/client/test_session.py +++ b/packages/mcp/tests/keycardai/mcp/client/test_session.py @@ -867,6 +867,90 @@ async def test_get_auth_challenge_with_different_strategies(self): result = await session.get_auth_challenge() assert result == custom_challenge + @pytest.mark.asyncio + async def test_get_auth_challenge_returns_challenge_when_auth_pending_status(self): + """Test that a session in AUTH_PENDING status surfaces the stored challenge.""" + storage = InMemoryBackend() + coordinator = MockAuthCoordinator(storage) + context = coordinator.create_context("user:alice") + + server_config = {"url": "http://localhost:3000"} + session = Session("test_server", server_config, context, coordinator) + session.status = SessionStatus.AUTH_PENDING + + challenge = { + "authorization_url": "http://auth.example.com", + "state": "state123" + } + await coordinator.set_auth_pending( + context_id=context.id, + server_name="test_server", + auth_metadata=challenge + ) + + result = await session.get_auth_challenge() + + assert result == challenge + + @pytest.mark.asyncio + async def test_get_auth_challenge_returns_none_when_connected(self): + """Test that a connected session never advertises an auth challenge. + + Simulates a completed authorization whose cleanup task was dropped or + cancelled: the pending-auth record is still in storage, but the session + is CONNECTED, so no challenge is returned and the stale record is + cleared lazily. + """ + storage = InMemoryBackend() + coordinator = MockAuthCoordinator(storage) + context = coordinator.create_context("user:alice") + + server_config = {"url": "http://localhost:3000"} + session = Session("test_server", server_config, context, coordinator) + + # Stale record left behind by a dropped cleanup task + await coordinator.set_auth_pending( + context_id=context.id, + server_name="test_server", + auth_metadata={ + "authorization_url": "http://auth.example.com", + "state": "state123" + } + ) + session.status = SessionStatus.CONNECTED + + result = await session.get_auth_challenge() + + assert result is None + # The stale record was cleared from storage + stored = await coordinator.get_auth_pending( + context_id=context.id, + server_name="test_server" + ) + assert stored is None + + @pytest.mark.asyncio + async def test_requires_auth_returns_false_when_connected_with_stale_record(self): + """Test that requires_auth is False for a connected session with a stale record.""" + storage = InMemoryBackend() + coordinator = MockAuthCoordinator(storage) + context = coordinator.create_context("user:alice") + + server_config = {"url": "http://localhost:3000"} + session = Session("test_server", server_config, context, coordinator) + + await coordinator.set_auth_pending( + context_id=context.id, + server_name="test_server", + auth_metadata={ + "authorization_url": "http://auth.example.com", + "state": "state123" + } + ) + session.status = SessionStatus.CONNECTED + + assert await session.requires_auth() is False + class TestSessionStorageIsolation: """Test Session storage isolation between different sessions.""" From d737b81337d9ab1abf1de1fe41c48cff8c938ec6 Mon Sep 17 00:00:00 2001 From: GitHub Action Date: Tue, 21 Jul 2026 14:14:16 -0700 Subject: [PATCH 3/4] fix(keycardai-mcp): clear stale auth-pending records on connect, not in the getter The CONNECTED short-circuit in Session.get_auth_challenge() clobbered fresh re-auth challenges. A mid-session 401 (token expiry on a live tool call) is handled by the httpx transport, which writes a new pending-auth record without ever changing the session status, so a long-lived CONNECTED session with an expired token had its fresh challenge deleted by the getter and get_auth_challenges() returned [], hiding the re-auth URL from the user. A destructive write inside a getter that polling APIs hit repeatedly was also the wrong shape. get_auth_challenge() is a plain storage read again. The stale-record clear moves to the transition into CONNECTED in _initialize_session: reaching CONNECTED means auth works, so any record stored at that point is stale by definition. The clear runs once at the transition, is best-effort (storage errors never fail the connection), and cannot clobber a mid-session re-auth challenge because that is written while the session is already CONNECTED. The auto-reconnect path in on_completion_handled flows through connect() into _initialize_session, so one hook covers both paths. Co-Authored-By: Claude Fable 5 --- .../mcp/src/keycardai/mcp/client/session.py | 54 ++++++-- .../tests/keycardai/mcp/client/test_client.py | 84 ++++++++++--- .../keycardai/mcp/client/test_session.py | 119 ++++++++++++++---- 3 files changed, 208 insertions(+), 49 deletions(-) diff --git a/packages/mcp/src/keycardai/mcp/client/session.py b/packages/mcp/src/keycardai/mcp/client/session.py index 47c12909..a1aae5d4 100644 --- a/packages/mcp/src/keycardai/mcp/client/session.py +++ b/packages/mcp/src/keycardai/mcp/client/session.py @@ -415,6 +415,16 @@ async def _initialize_session(self, retry_after_auth: bool) -> None: try: await self._session.initialize() self._set_status(SessionStatus.CONNECTED, "session initialized") + # Reaching CONNECTED means auth works, so any pending-auth record + # still stored for this session is stale by definition (e.g. a + # completion cleanup that never ran). Clearing at the transition + # happens once, not on every read, and cannot clobber a challenge + # written after connection (mid-session token expiry re-auth), + # because that is written while the session is already CONNECTED. + # The auto-reconnect in on_completion_handled also flows through + # connect() into this method, so this is the single hook for the + # transition into CONNECTED. + await self._clear_stale_auth_pending() except Exception as e: auth_challenge = await self.get_auth_challenge() @@ -489,17 +499,37 @@ async def get_auth_challenge(self) -> dict[str, str] | None: An auth challenge is created by the auth strategy when authentication is required but not yet complete (e.g., waiting for OAuth callback). - A connected session has no pending challenge by definition, so this - returns None when the session status is CONNECTED. Any record still - stored for this session at that point is stale (e.g. completion - cleanup did not run) and is cleared lazily. + This is a plain storage read with no side effects, so polling it + repeatedly is safe. A challenge can exist even while the session is + CONNECTED: a mid-session token expiry is handled at the transport + level and writes a fresh pending-auth record without changing the + session status, and that challenge must be surfaced so the user can + re-authorize. Stale records from completed authorizations are cleared + on the transition into CONNECTED, not here. Returns: Dict with challenge details (strategy-specific) or None if no pending challenge. For OAuth: {'authorization_url': str, 'state': str} For other strategies: may contain different fields """ - if self.status == SessionStatus.CONNECTED: + # Auth challenge is stored in coordinator + return await self.coordinator.get_auth_pending( + context_id=self.context.id, + server_name=self.server_name + ) + + async def _clear_stale_auth_pending(self) -> None: + """ + Best-effort clear of a lingering pending-auth record after connecting. + + Called on the transition into CONNECTED. A record still stored at + that point belongs to an already-completed authorization (e.g. the + completion cleanup never ran) and would otherwise keep being + advertised as an auth challenge. Storage errors are logged and + swallowed: a stale record must never fail an otherwise successful + connection. + """ + try: stale_challenge = await self.coordinator.get_auth_pending( context_id=self.context.id, server_name=self.server_name @@ -507,19 +537,17 @@ async def get_auth_challenge(self) -> dict[str, str] | None: if stale_challenge: logger.debug( f"Session {self.server_name}: Clearing stale auth challenge " - f"for connected session" + f"after successful connection" ) await self.coordinator.clear_auth_pending( context_id=self.context.id, server_name=self.server_name ) - return None - - # Auth challenge is stored in coordinator - return await self.coordinator.get_auth_pending( - context_id=self.context.id, - server_name=self.server_name - ) + except Exception as e: + logger.warning( + f"Session {self.server_name}: Failed to clear stale " + f"pending-auth record after connecting: {e}" + ) def __getattr__(self, name: str) -> Any: """ diff --git a/packages/mcp/tests/keycardai/mcp/client/test_client.py b/packages/mcp/tests/keycardai/mcp/client/test_client.py index 844bbc27..cb5102d9 100644 --- a/packages/mcp/tests/keycardai/mcp/client/test_client.py +++ b/packages/mcp/tests/keycardai/mcp/client/test_client.py @@ -4,6 +4,8 @@ connection handling, and coordination with sessions. """ +from unittest.mock import AsyncMock, MagicMock, patch + import pytest from mcp import Tool @@ -1017,24 +1019,21 @@ async def test_get_auth_challenges_includes_server_name(self): assert challenges[0]["other_field"] == "value" @pytest.mark.asyncio - async def test_get_auth_challenges_empty_for_connected_session_with_stale_record(self): - """Test that a connected session's stale pending-auth record is not surfaced. - - Simulates a completed authorization and auto-reconnect where the - completion cleanup task was dropped before running: the pending-auth - record is still in coordinator storage, but the session is CONNECTED. - get_auth_challenges() must return an empty list without any - consumer-side filtering. + async def test_get_auth_challenges_empty_after_completed_auth_and_reconnect(self): + """Completed authorization plus reconnect leaves no auth challenges. + + Simulates a completed authorization whose cleanup never ran: the + pending-auth record is still in coordinator storage when the session + reconnects. The transition into CONNECTED clears the stale record, so + get_auth_challenges() returns an empty list without any consumer-side + filtering. """ servers = {"test_server": {"url": "http://localhost:3000"}} mock_coordinator = MockAuthCoordinator() client = Client(servers=servers, auth_coordinator=mock_coordinator) - # Real Session (not a mock) so the CONNECTED short-circuit is exercised - session = client.sessions["test_server"] - - # Stale record left behind by a dropped cleanup task + # Stale record left behind by a cleanup that never ran await mock_coordinator.set_auth_pending( context_id=client.context.id, server_name="test_server", @@ -1043,18 +1042,75 @@ async def test_get_auth_challenges_empty_for_connected_session_with_stale_record "state": "state123" } ) - session.status = SessionStatus.CONNECTED + + # Real Session (not a mock) so connect() flows through the + # transition into CONNECTED, the same path the auto-reconnect + # after auth completion takes. + mock_connection = MagicMock() + mock_connection.start = AsyncMock(return_value=(MagicMock(), MagicMock())) + mock_connection.stop = AsyncMock() + + mock_client_session = MagicMock() + mock_client_session.__aenter__ = AsyncMock(return_value=mock_client_session) + mock_client_session.__aexit__ = AsyncMock() + mock_client_session.initialize = AsyncMock() + + with patch("keycardai.mcp.client.session.create_connection", return_value=mock_connection), \ + patch("keycardai.mcp.client.session.ClientSession", return_value=mock_client_session): + await client.connect() + + assert client.sessions["test_server"].status == SessionStatus.CONNECTED challenges = await client.get_auth_challenges() assert challenges == [] - # The stale record was cleared lazily + # The stale record was cleared at the transition into CONNECTED stored = await mock_coordinator.get_auth_pending( context_id=client.context.id, server_name="test_server" ) assert stored is None + @pytest.mark.asyncio + async def test_get_auth_challenges_surfaces_fresh_challenge_for_connected_session(self): + """A challenge written while a session is CONNECTED is surfaced. + + Regression guard for mid-session token expiry: transport-level 401 + handling writes a fresh pending-auth record without ever changing the + session status. get_auth_challenges() must include that challenge so + the user sees the re-auth URL, and reading it must not delete it. + """ + servers = {"test_server": {"url": "http://localhost:3000"}} + mock_coordinator = MockAuthCoordinator() + + client = Client(servers=servers, auth_coordinator=mock_coordinator) + + # Real Session (not a mock) so the real getter is exercised + session = client.sessions["test_server"] + session.status = SessionStatus.CONNECTED + + # Fresh record written by the transport while already CONNECTED + challenge = { + "authorization_url": "http://auth.example.com/reauth", + "state": "state456" + } + await mock_coordinator.set_auth_pending( + context_id=client.context.id, + server_name="test_server", + auth_metadata=challenge + ) + + challenges = await client.get_auth_challenges() + + assert challenges == [{**challenge, "server": "test_server"}] + # Reading did not delete the record: a second poll still surfaces it + assert await client.get_auth_challenges() == challenges + stored = await mock_coordinator.get_auth_pending( + context_id=client.context.id, + server_name="test_server" + ) + assert stored == challenge + class TestClientIdGeneration: """Test Client ID generation.""" diff --git a/packages/mcp/tests/keycardai/mcp/client/test_session.py b/packages/mcp/tests/keycardai/mcp/client/test_session.py index b2a5d36c..b043ce21 100644 --- a/packages/mcp/tests/keycardai/mcp/client/test_session.py +++ b/packages/mcp/tests/keycardai/mcp/client/test_session.py @@ -431,6 +431,78 @@ async def test_connect_cleans_up_on_error(self): assert session.status == SessionStatus.FAILED assert session.is_failed + @pytest.mark.asyncio + async def test_connect_clears_lingering_stale_auth_pending(self): + """Transitioning into CONNECTED clears a lingering pending-auth record. + + Simulates a completed authorization whose cleanup never ran: the + pending-auth record is still in storage when the session successfully + connects. After the transition into CONNECTED the record is gone and + no auth challenge is surfaced. + """ + storage = InMemoryBackend() + coordinator = MockAuthCoordinator(storage) + context = coordinator.create_context("user:alice") + + server_config = {"url": "http://localhost:3000"} + session = Session("test_server", server_config, context, coordinator) + + # Stale record left behind by a cleanup that never ran + await coordinator.set_auth_pending( + context_id=context.id, + server_name="test_server", + auth_metadata={ + "authorization_url": "http://auth.example.com", + "state": "state123" + } + ) + + mock_connection = MockConnection() + mock_client_session = MockClientSession() + + with patch("keycardai.mcp.client.session.create_connection", return_value=mock_connection): + with patch("keycardai.mcp.client.session.ClientSession", return_value=mock_client_session): + await session.connect() + + assert session.status == SessionStatus.CONNECTED + assert await session.get_auth_challenge() is None + stored = await coordinator.get_auth_pending( + context_id=context.id, + server_name="test_server" + ) + assert stored is None + + @pytest.mark.asyncio + async def test_connect_succeeds_when_stale_auth_pending_clear_fails(self): + """A storage error clearing the stale record must not fail the connection.""" + storage = InMemoryBackend() + coordinator = MockAuthCoordinator(storage) + context = coordinator.create_context("user:alice") + + server_config = {"url": "http://localhost:3000"} + session = Session("test_server", server_config, context, coordinator) + + await coordinator.set_auth_pending( + context_id=context.id, + server_name="test_server", + auth_metadata={ + "authorization_url": "http://auth.example.com", + "state": "state123" + } + ) + coordinator.clear_auth_pending = AsyncMock(side_effect=RuntimeError("storage down")) + + mock_connection = MockConnection() + mock_client_session = MockClientSession() + + with patch("keycardai.mcp.client.session.create_connection", return_value=mock_connection): + with patch("keycardai.mcp.client.session.ClientSession", return_value=mock_client_session): + await session.connect() + + assert session.status == SessionStatus.CONNECTED + assert session.is_operational + coordinator.clear_auth_pending.assert_awaited_once() + @pytest.mark.asyncio async def test_connect_disconnects_existing_session_if_connected_but_no_session(self): """Test that connect disconnects if marked connected but no session.""" @@ -893,13 +965,14 @@ async def test_get_auth_challenge_returns_challenge_when_auth_pending_status(sel assert result == challenge @pytest.mark.asyncio - async def test_get_auth_challenge_returns_none_when_connected(self): - """Test that a connected session never advertises an auth challenge. - - Simulates a completed authorization whose cleanup task was dropped or - cancelled: the pending-auth record is still in storage, but the session - is CONNECTED, so no challenge is returned and the stale record is - cleared lazily. + async def test_get_auth_challenge_surfaces_fresh_challenge_while_connected(self): + """A challenge written while the session is CONNECTED is surfaced, not deleted. + + Regression guard for mid-session token expiry: a live tool call gets a + 401, the transport's auth strategy writes a fresh pending-auth record, + but the session status never leaves CONNECTED. The challenge must be + returned so the user sees the re-auth URL, and reading it must not + delete the record (polling APIs call this getter repeatedly). """ storage = InMemoryBackend() coordinator = MockAuthCoordinator(storage) @@ -907,49 +980,51 @@ async def test_get_auth_challenge_returns_none_when_connected(self): server_config = {"url": "http://localhost:3000"} session = Session("test_server", server_config, context, coordinator) + session.status = SessionStatus.CONNECTED - # Stale record left behind by a dropped cleanup task + # Fresh challenge written by the transport while already CONNECTED + challenge = { + "authorization_url": "http://auth.example.com/reauth", + "state": "state456" + } await coordinator.set_auth_pending( context_id=context.id, server_name="test_server", - auth_metadata={ - "authorization_url": "http://auth.example.com", - "state": "state123" - } + auth_metadata=challenge ) - session.status = SessionStatus.CONNECTED result = await session.get_auth_challenge() + assert result == challenge - assert result is None - # The stale record was cleared from storage + # Reading is side-effect free: the record survives repeated polls + assert await session.get_auth_challenge() == challenge stored = await coordinator.get_auth_pending( context_id=context.id, server_name="test_server" ) - assert stored is None + assert stored == challenge @pytest.mark.asyncio - async def test_requires_auth_returns_false_when_connected_with_stale_record(self): - """Test that requires_auth is False for a connected session with a stale record.""" + async def test_requires_auth_true_when_connected_with_fresh_challenge(self): + """requires_auth reflects a fresh challenge even for a CONNECTED session.""" storage = InMemoryBackend() coordinator = MockAuthCoordinator(storage) context = coordinator.create_context("user:alice") server_config = {"url": "http://localhost:3000"} session = Session("test_server", server_config, context, coordinator) + session.status = SessionStatus.CONNECTED await coordinator.set_auth_pending( context_id=context.id, server_name="test_server", auth_metadata={ - "authorization_url": "http://auth.example.com", - "state": "state123" + "authorization_url": "http://auth.example.com/reauth", + "state": "state456" } ) - session.status = SessionStatus.CONNECTED - assert await session.requires_auth() is False + assert await session.requires_auth() is True class TestSessionStorageIsolation: From b4084940be2beab11f71d4eb1798447c9b017b81 Mon Sep 17 00:00:00 2001 From: GitHub Action Date: Tue, 21 Jul 2026 14:14:28 -0700 Subject: [PATCH 4/4] fix(keycardai-mcp): bound synchronous OAuth cleanup with a timeout Completion cleanup runs synchronously on the OAuth callback path, so a hung storage delete (remote backends such as Redis or DynamoDB) would stall the user-facing callback response indefinitely. Each cleanup storage call (PKCE state delete, pending-auth clear, completion-route delete) is now wrapped in asyncio.wait_for with a 5 second bound. A timeout falls through to the existing logged-and-swallowed per-step handling, so a timed-out step never skips the remaining steps. asyncio.wait_for raises TimeoutError (asyncio.TimeoutError on Python 3.10, an alias of the builtin on 3.11+), an Exception subclass on every supported Python, while outer cancellation (CancelledError, a BaseException) still propagates. Co-Authored-By: Claude Fable 5 --- .../src/keycardai/mcp/client/auth/handlers.py | 53 ++++-- .../mcp/client/auth/test_handlers.py | 176 ++++++++++++++++++ 2 files changed, 215 insertions(+), 14 deletions(-) diff --git a/packages/mcp/src/keycardai/mcp/client/auth/handlers.py b/packages/mcp/src/keycardai/mcp/client/auth/handlers.py index dfe3832d..76b5b63e 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/handlers.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/handlers.py @@ -1,5 +1,6 @@ """Auth completion routing and handlers.""" +import asyncio import warnings from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any @@ -15,6 +16,16 @@ logger = get_logger(__name__) +# Upper bound for each individual cleanup storage call. Cleanup runs +# synchronously on the OAuth callback path, so a hung remote backend +# (e.g. Redis, DynamoDB) must not stall the user-facing callback response +# indefinitely. asyncio.wait_for raises TimeoutError (asyncio.TimeoutError +# on Python 3.10, an alias of the builtin on 3.11+), which is an Exception +# subclass on every supported Python, so a timeout falls through to the +# same logged-and-swallowed handling as any other storage error while +# outer cancellation (CancelledError, a BaseException) still propagates. +_CLEANUP_TIMEOUT_SECONDS = 5.0 + class CompletionRouter: """ @@ -174,15 +185,20 @@ async def _cleanup_completion_route(self, state: str) -> None: """ Clean up completion routing metadata. - Runs before the completion result is returned. Storage errors are - logged and not propagated because the completion itself already - succeeded; cancellation propagates normally. + Runs before the completion result is returned. The delete is bounded + by _CLEANUP_TIMEOUT_SECONDS so a hung storage backend cannot stall + the callback response. Storage errors (including timeouts) are logged + and not propagated because the completion itself already succeeded; + cancellation propagates normally. Args: state: OAuth state parameter to clean up """ try: - await self.state_storage.delete_completion_route(state) + await asyncio.wait_for( + self.state_storage.delete_completion_route(state), + timeout=_CLEANUP_TIMEOUT_SECONDS, + ) logger.debug(f"Cleaned up completion route for state: {state[:8]}...") except Exception as e: # Log but don't propagate cleanup errors - the completion succeeded @@ -564,10 +580,11 @@ def my_factory(): namespace_parts = storage._namespace.split(":") context_id = namespace_parts[1] if len(namespace_parts) > 1 else "unknown" - # Cleanup runs synchronously before returning: both deletes are fast storage - # calls, and a fire-and-forget task can be garbage-collected (or cancelled at - # shutdown) before it runs, leaving a stale pending-auth record that sessions - # would keep advertising as an auth challenge. + # Cleanup runs synchronously before returning: a fire-and-forget task can be + # garbage-collected (or cancelled at shutdown) before it runs, leaving a + # stale pending-auth record that sessions would keep advertising as an auth + # challenge. Each storage call is bounded by _CLEANUP_TIMEOUT_SECONDS so a + # hung backend cannot stall the callback response. # Note: completion metadata cleanup is handled by coordinator await _cleanup_oauth_completion_state( storage, coordinator, state, context_id, server_name @@ -590,10 +607,12 @@ async def _cleanup_oauth_completion_state( Runs before the completion result is returned so a completed authorization never leaves a stale pending-auth record behind. Each step is attempted - independently: a failure deleting the PKCE state must not skip clearing - the pending-auth record. Storage errors are logged and not propagated - because the token exchange already succeeded; cancellation propagates - normally. + independently: a failure or timeout deleting the PKCE state must not skip + clearing the pending-auth record. Each storage call is bounded by + _CLEANUP_TIMEOUT_SECONDS so a hung backend cannot stall the callback + response. Storage errors (including timeouts) are logged and not + propagated because the token exchange already succeeded; cancellation + propagates normally. Args: storage: Strategy storage @@ -603,12 +622,18 @@ async def _cleanup_oauth_completion_state( server_name: Server name """ try: - await storage.delete(f"_pkce_state:{state}") + await asyncio.wait_for( + storage.delete(f"_pkce_state:{state}"), + timeout=_CLEANUP_TIMEOUT_SECONDS, + ) except Exception as e: logger.warning(f"Failed to delete PKCE state for {server_name}: {e}") try: - await coordinator.clear_auth_pending(context_id=context_id, server_name=server_name) + await asyncio.wait_for( + coordinator.clear_auth_pending(context_id=context_id, server_name=server_name), + timeout=_CLEANUP_TIMEOUT_SECONDS, + ) logger.debug(f"Cleaned up OAuth completion state for {server_name}") except Exception as e: logger.warning(f"Failed to clear pending auth for {server_name}: {e}") diff --git a/packages/mcp/tests/keycardai/mcp/client/auth/test_handlers.py b/packages/mcp/tests/keycardai/mcp/client/auth/test_handlers.py index 27aa717d..a48b50d1 100644 --- a/packages/mcp/tests/keycardai/mcp/client/auth/test_handlers.py +++ b/packages/mcp/tests/keycardai/mcp/client/auth/test_handlers.py @@ -1,11 +1,15 @@ """Unit tests for the CompletionHandlerRegistry.""" +import asyncio +import logging from unittest.mock import AsyncMock, MagicMock import pytest +import keycardai.mcp.client.auth.handlers as handlers_module from keycardai.mcp.client.auth.handlers import ( CompletionHandlerRegistry, + _cleanup_oauth_completion_state, get_default_handler_registry, oauth_completion_handler, register_completion_handler, @@ -290,6 +294,128 @@ async def test_run_cleanup_in_background_is_deprecated_and_ignored( assert await storage.get(f"_pkce_state:{state}") is None +class TestCleanupTimeout: + """Cleanup storage calls are bounded so a hung backend can't stall the callback.""" + + @pytest.mark.asyncio + async def test_hung_pkce_delete_times_out_and_still_clears_auth_pending( + self, monkeypatch, caplog + ): + """A hung PKCE delete times out, logs a warning, and the next step still runs.""" + monkeypatch.setattr(handlers_module, "_CLEANUP_TIMEOUT_SECONDS", 0.05) + + coordinator = AsyncMock() + backend = InMemoryBackend() + storage = NamespacedStorage( + backend, "client:test_user:server:test_server:connection:oauth" + ) + state = "state_hung_pkce_delete" + await _seed_pkce_state(storage, state) + + async def hung_delete(key: str) -> None: + await asyncio.Event().wait() # never set: hangs forever + + storage.delete = hung_delete + + with caplog.at_level(logging.WARNING, logger="keycardai.mcp.client"): + # Outer guard: if the timeout wrapper were broken, fail fast + # instead of hanging the test run. + result = await asyncio.wait_for( + oauth_completion_handler( + coordinator, + storage, + {"code": "auth_code", "state": state}, + client_factory=_token_client_factory, + ), + timeout=2.0, + ) + + assert result["success"] is True + # The timed-out PKCE delete did not skip clearing pending auth + coordinator.clear_auth_pending.assert_awaited_once_with( + context_id="test_user", server_name="test_server" + ) + assert any( + "Failed to delete PKCE state" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.asyncio + async def test_hung_clear_auth_pending_times_out_and_handler_returns( + self, monkeypatch, caplog + ): + """A hung pending-auth clear times out and the handler still returns success.""" + monkeypatch.setattr(handlers_module, "_CLEANUP_TIMEOUT_SECONDS", 0.05) + + coordinator = AsyncMock() + + async def hung_clear(context_id: str, server_name: str) -> None: + await asyncio.Event().wait() + + coordinator.clear_auth_pending = hung_clear + + backend = InMemoryBackend() + storage = NamespacedStorage( + backend, "client:test_user:server:test_server:connection:oauth" + ) + state = "state_hung_clear" + await _seed_pkce_state(storage, state) + + with caplog.at_level(logging.WARNING, logger="keycardai.mcp.client"): + result = await asyncio.wait_for( + oauth_completion_handler( + coordinator, + storage, + {"code": "auth_code", "state": state}, + client_factory=_token_client_factory, + ), + timeout=2.0, + ) + + assert result["success"] is True + # The PKCE delete step still ran to completion + assert await storage.get(f"_pkce_state:{state}") is None + assert any( + "Failed to clear pending auth" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.asyncio + async def test_outer_cancellation_propagates_through_cleanup(self): + """Cancelling the task running cleanup propagates CancelledError. + + The timeout wrapper converts inner timeouts to TimeoutError, but outer + cancellation must still propagate and abort the remaining steps. + """ + coordinator = AsyncMock() + backend = InMemoryBackend() + storage = NamespacedStorage( + backend, "client:test_user:server:test_server:connection:oauth" + ) + + started = asyncio.Event() + + async def hung_delete(key: str) -> None: + started.set() + await asyncio.Event().wait() + + storage.delete = hung_delete + + task = asyncio.create_task( + _cleanup_oauth_completion_state( + storage, coordinator, "state_cancelled", "test_user", "test_server" + ) + ) + await started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + # Cancellation aborts cleanup entirely: the next step never ran + coordinator.clear_auth_pending.assert_not_awaited() + + class TestCompletionRouterCleanup: """Test that CompletionRouter cleans up routing metadata synchronously.""" @@ -341,6 +467,56 @@ async def recording_delete(routing_key: str) -> None: assert deleted_routes == [state] assert await state_storage.get_completion_route(state) is None + @pytest.mark.asyncio + async def test_hung_route_delete_times_out_and_completion_returns( + self, monkeypatch, caplog + ): + """A hung completion-route delete times out instead of stalling the callback.""" + from keycardai.mcp.client.auth.coordinators.base import AuthCoordinator + + monkeypatch.setattr(handlers_module, "_CLEANUP_TIMEOUT_SECONDS", 0.05) + + class TestCoordinator(AuthCoordinator): + @property + def endpoint_type(self) -> str: + return "test" + + registry = CompletionHandlerRegistry() + + @registry.register("stub_completion") + async def stub_completion(coordinator, storage, params, **kwargs): + return {"success": True} + + coordinator = TestCoordinator( + backend=InMemoryBackend(), handler_registry=registry + ) + + state = "state_hung_route_delete" + await coordinator.register_completion_route( + routing_key=state, + handler_name="stub_completion", + storage_namespace="client:test_user:server:test_server:connection:oauth", + context_id="test_user", + server_name="test_server", + ) + + async def hung_delete(routing_key: str) -> None: + await asyncio.Event().wait() + + coordinator.completion_router.state_storage.delete_completion_route = hung_delete + + with caplog.at_level(logging.WARNING, logger="keycardai.mcp.client"): + result = await asyncio.wait_for( + coordinator.handle_completion({"code": "auth_code", "state": state}), + timeout=2.0, + ) + + assert result["success"] is True + assert any( + "Failed to cleanup completion route" in record.getMessage() + for record in caplog.records + ) + class TestClientFactoryManagement: """Test client factory management in the registry."""