mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-20 04:24:17 +02:00
Progress replacing asyncio with anyio (#2143)
* Replace asyncio.sleep() with anyio.sleep() - Replace asyncio.sleep() in error_handling.py retry middleware - Replace asyncio.sleep() in oauth.py callback shutdown - Keep asyncio.TimeoutError check for Python 3.10 compatibility - Add anyio import to error_handling.py All core library sleep calls now use anyio primitives. Tests and example code still use asyncio where appropriate. * Replace OAuth asyncio.Future with anyio.Event pattern - Create OAuthCallbackResult dataclass for result storage - Replace Future with Event + result container pattern - Update oauth_callback.py to use anyio.Event coordination - Update auth/oauth.py callback_handler to use Event pattern - Remove asyncio imports from OAuth flow OAuth callback now uses anyio primitives for async coordination instead of asyncio.Future. * Remove asyncio fire-and-forget task hack from Context - Remove _try_flush_notifications() method entirely - Update _queue_*_list_changed() to only queue notifications - Remove asyncio import from context.py - Keep _flush_notifications() for deferred sending on context exit Notifications now flush reliably on request completion (__aexit__) instead of attempting immediate delivery with asyncio.create_task(). Slight delay is acceptable - all notifications are deduplicated and sent when the MCP request handler completes.
This commit is contained in:
Jeremiah Lowin
GitHub
parent
3321644ad3
commit
96aa150cf0
4 changed files with 46 additions and 44 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue