fix(keycardai-mcp): synchronous OAuth completion cleanup, no stale auth challenges

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):

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:
         """

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.

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}")
 

packages/mcp/src/keycardai/mcp/client/auth/strategies/oauth.py

@@ -272,23 +272,15 @@ 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",
                 storage_namespace=self.oauth_storage._storage._namespace,
                 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)
             )

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.