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..76b5b63e 100644 --- a/packages/mcp/src/keycardai/mcp/client/auth/handlers.py +++ b/packages/mcp/src/keycardai/mcp/client/auth/handlers.py @@ -1,6 +1,7 @@ """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: """ @@ -55,8 +66,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 +94,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,22 +183,25 @@ 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. 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 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 +464,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 +477,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 +486,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 +514,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 +580,15 @@ 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: 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 - 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 +603,16 @@ 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 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 @@ -592,12 +622,19 @@ async def _cleanup_oauth_completion_state( server_name: Server name """ try: - await storage.delete(f"_pkce_state:{state}") - await coordinator.clear_auth_pending(context_id=context_id, server_name=server_name) + 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 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 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/src/keycardai/mcp/client/session.py b/packages/mcp/src/keycardai/mcp/client/session.py index 90a4b078..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,6 +499,14 @@ 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). + 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} @@ -500,6 +518,37 @@ async def get_auth_challenge(self) -> dict[str, str] | None: 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 + ) + if stale_challenge: + logger.debug( + f"Session {self.server_name}: Clearing stale auth challenge " + f"after successful connection" + ) + await self.coordinator.clear_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: """ Delegate unknown attributes to underlying MCP ClientSession. 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..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, @@ -169,6 +173,351 @@ 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 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.""" + + @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 + + @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.""" 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.""" diff --git a/packages/mcp/tests/keycardai/mcp/client/test_client.py b/packages/mcp/tests/keycardai/mcp/client/test_client.py index 6982707a..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 @@ -1016,6 +1018,99 @@ 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_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) + + # 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", + auth_metadata={ + "authorization_url": "http://auth.example.com", + "state": "state123" + } + ) + + # 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 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 71e36dcb..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.""" @@ -867,6 +939,93 @@ 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_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) + context = coordinator.create_context("user:alice") + + server_config = {"url": "http://localhost:3000"} + session = Session("test_server", server_config, context, coordinator) + session.status = SessionStatus.CONNECTED + + # 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=challenge + ) + + result = await session.get_auth_challenge() + assert result == challenge + + # 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 == challenge + + @pytest.mark.asyncio + 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/reauth", + "state": "state456" + } + ) + + assert await session.requires_auth() is True + class TestSessionStorageIsolation: """Test Session storage isolation between different sessions."""