diff --git a/src/fastmcp/client/auth/oauth.py b/src/fastmcp/client/auth/oauth.py index 224aeb45c..0795d6f76 100644 --- a/src/fastmcp/client/auth/oauth.py +++ b/src/fastmcp/client/auth/oauth.py @@ -1,9 +1,7 @@ from __future__ import annotations -import asyncio import time import webbrowser -from asyncio import Future from collections.abc import AsyncGenerator from typing import Any from urllib.parse import urlparse @@ -24,6 +22,7 @@ from typing_extensions import override from uvicorn.server import Server from fastmcp.client.oauth_callback import ( + OAuthCallbackResult, create_oauth_callback_server, ) from fastmcp.utilities.http import find_available_port @@ -247,14 +246,16 @@ class OAuth(OAuthClientProvider): async def callback_handler(self) -> tuple[str, str | None]: """Handle OAuth callback and return (auth_code, state).""" - # Create a future to capture the OAuth response - response_future: Future[Any] = asyncio.get_running_loop().create_future() + # Create result container and event to capture the OAuth response + result = OAuthCallbackResult() + result_ready = anyio.Event() - # Create server with the future + # Create server with result tracking server: Server = create_oauth_callback_server( port=self.redirect_port, server_url=self.server_base_url, - response_future=response_future, + result_container=result, + result_ready=result_ready, ) # Run server until response is received with timeout logic @@ -267,13 +268,15 @@ class OAuth(OAuthClientProvider): TIMEOUT = 300.0 # 5 minute timeout try: with anyio.fail_after(TIMEOUT): - auth_code, state = await response_future - return auth_code, state + await result_ready.wait() + if result.error: + raise result.error + return result.code, result.state # type: ignore except TimeoutError: raise TimeoutError(f"OAuth callback timed out after {TIMEOUT} seconds") finally: server.should_exit = True - await asyncio.sleep(0.1) # Allow server to shut down gracefully + await anyio.sleep(0.1) # Allow server to shut down gracefully tg.cancel_scope.cancel() raise RuntimeError("OAuth callback handler could not be started") diff --git a/src/fastmcp/client/oauth_callback.py b/src/fastmcp/client/oauth_callback.py index ced483ea2..cbe3c84ae 100644 --- a/src/fastmcp/client/oauth_callback.py +++ b/src/fastmcp/client/oauth_callback.py @@ -7,9 +7,9 @@ and display styled responses to users. from __future__ import annotations -import asyncio from dataclasses import dataclass +import anyio from starlette.applications import Starlette from starlette.requests import Request from starlette.routing import Route @@ -87,11 +87,21 @@ class CallbackResponse: return {k: v for k, v in self.__dict__.items() if v is not None} +@dataclass +class OAuthCallbackResult: + """Container for OAuth callback results, used with anyio.Event for async coordination.""" + + code: str | None = None + state: str | None = None + error: Exception | None = None + + def create_oauth_callback_server( port: int, callback_path: str = "/callback", server_url: str | None = None, - response_future: asyncio.Future | None = None, + result_container: OAuthCallbackResult | None = None, + result_ready: anyio.Event | None = None, ) -> Server: """ Create an OAuth callback server. @@ -100,7 +110,8 @@ def create_oauth_callback_server( port: The port to run the server on callback_path: The path to listen for OAuth redirects on server_url: Optional server URL to display in success messages - response_future: Optional future to resolve when OAuth callback is received + result_container: Optional container to store callback results + result_ready: Optional event to signal when callback is received Returns: Configured uvicorn Server instance (not yet running) @@ -120,9 +131,10 @@ def create_oauth_callback_server( else: user_message = f"Authorization failed: {error_desc}" - # Resolve future with exception if provided - if response_future and not response_future.done(): - response_future.set_exception(RuntimeError(user_message)) + # Store error and signal completion if result tracking provided + if result_container is not None and result_ready is not None: + result_container.error = RuntimeError(user_message) + result_ready.set() return create_secure_html_response( create_callback_html( @@ -135,9 +147,10 @@ def create_oauth_callback_server( if not callback_response.code: user_message = "No authorization code was received from the server." - # Resolve future with exception if provided - if response_future and not response_future.done(): - response_future.set_exception(RuntimeError(user_message)) + # Store error and signal completion if result tracking provided + if result_container is not None and result_ready is not None: + result_container.error = RuntimeError(user_message) + result_ready.set() return create_secure_html_response( create_callback_html( @@ -153,9 +166,10 @@ def create_oauth_callback_server( "The OAuth server did not return the expected state parameter." ) - # Resolve future with exception if provided - if response_future and not response_future.done(): - response_future.set_exception(RuntimeError(user_message)) + # Store error and signal completion if result tracking provided + if result_container is not None and result_ready is not None: + result_container.error = RuntimeError(user_message) + result_ready.set() return create_secure_html_response( create_callback_html( @@ -165,11 +179,11 @@ def create_oauth_callback_server( status_code=400, ) - # Success case - if response_future and not response_future.done(): - response_future.set_result( - (callback_response.code, callback_response.state) - ) + # Success case - store result and signal completion if result tracking provided + if result_container is not None and result_ready is not None: + result_container.code = callback_response.code + result_container.state = callback_response.state + result_ready.set() return create_secure_html_response( create_callback_html("", is_success=True, server_url=server_url) diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index 9c3feaf07..a0fcfc1c2 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -1,6 +1,5 @@ from __future__ import annotations -import asyncio import copy import inspect import logging @@ -636,30 +635,14 @@ class Context: def _queue_tool_list_changed(self) -> None: """Queue a tool list changed notification.""" self._notification_queue.add("notifications/tools/list_changed") - self._try_flush_notifications() def _queue_resource_list_changed(self) -> None: """Queue a resource list changed notification.""" self._notification_queue.add("notifications/resources/list_changed") - self._try_flush_notifications() def _queue_prompt_list_changed(self) -> None: """Queue a prompt list changed notification.""" self._notification_queue.add("notifications/prompts/list_changed") - self._try_flush_notifications() - - def _try_flush_notifications(self) -> None: - """Synchronous method that attempts to flush notifications if we're in an async context.""" - try: - # Check if we're in an async context - loop = asyncio.get_running_loop() - if loop and not loop.is_running(): - return - # Schedule flush as a task (fire-and-forget) - asyncio.create_task(self._flush_notifications()) - except RuntimeError: - # No event loop - will flush later - pass async def _flush_notifications(self) -> None: """Send all queued notifications.""" diff --git a/src/fastmcp/server/middleware/error_handling.py b/src/fastmcp/server/middleware/error_handling.py index 176240584..0e222cd7b 100644 --- a/src/fastmcp/server/middleware/error_handling.py +++ b/src/fastmcp/server/middleware/error_handling.py @@ -6,6 +6,7 @@ import traceback from collections.abc import Callable from typing import Any +import anyio from mcp import McpError from mcp.types import ErrorData @@ -100,6 +101,7 @@ class ErrorHandlingMiddleware(Middleware): return McpError( ErrorData(code=-32000, message=f"Permission denied: {str(error)}") ) + # asyncio.TimeoutError is a subclass of TimeoutError in Python 3.10, alias in 3.11+ elif error_type in (TimeoutError, asyncio.TimeoutError): return McpError( ErrorData(code=-32000, message=f"Request timeout: {str(error)}") @@ -201,7 +203,7 @@ class RetryMiddleware(Middleware): f"{type(error).__name__}: {str(error)}. Retrying in {delay:.1f}s..." ) - await asyncio.sleep(delay) + await anyio.sleep(delay) # Re-raise the last error if all retries failed if last_error: