Merge branch 'ihrpr/client-resumability' into ihrpr/server-closing-streams

Author: ihrpr

.gitignore

@@ -166,4 +166,5 @@ cython_debug/
 
 # vscode
 .vscode/
+.windsurfrules
 **/CLAUDE.local.md

CLAUDE.md

@@ -19,7 +19,7 @@ This document contains critical information about working with this codebase. Fo
    - Line length: 88 chars maximum
 
 3. Testing Requirements
-   - Framework: `uv run pytest`
+   - Framework: `uv run --frozen pytest`
    - Async testing: use anyio, not asyncio
    - Coverage: test edge cases and errors
    - New features require tests
@@ -54,9 +54,9 @@ This document contains critical information about working with this codebase. Fo
 ## Code Formatting
 
 1. Ruff
-   - Format: `uv run ruff format .`
-   - Check: `uv run ruff check .`
-   - Fix: `uv run ruff check . --fix`
+   - Format: `uv run --frozen ruff format .`
+   - Check: `uv run --frozen ruff check .`
+   - Fix: `uv run --frozen ruff check . --fix`
    - Critical issues:
      - Line length (88 chars)
      - Import sorting (I001)
@@ -67,7 +67,7 @@ This document contains critical information about working with this codebase. Fo
      - Imports: split into multiple lines
 
 2. Type Checking
-   - Tool: `uv run pyright`
+   - Tool: `uv run --frozen pyright`
    - Requirements:
      - Explicit None checks for Optional
      - Type narrowing for strings
@@ -104,6 +104,10 @@ This document contains critical information about working with this codebase. Fo
      - Add None checks
      - Narrow string types
      - Match existing patterns
+   - Pytest:
+     - If the tests aren't finding the anyio pytest mark, try adding PYTEST_DISABLE_PLUGIN_AUTOLOAD=""
+       to the start of the pytest run command eg:
+       `PYTEST_DISABLE_PLUGIN_AUTOLOAD="" uv run --frozen pytest`
 
 3. Best Practices
    - Check git status before commits

README.md

@@ -309,6 +309,33 @@ async def long_task(files: list[str], ctx: Context) -> str:
     return "Processing complete"
 ```
 
+### Authentication
+
+Authentication can be used by servers that want to expose tools accessing protected resources.
+
+`mcp.server.auth` implements an OAuth 2.0 server interface, which servers can use by
+providing an implementation of the `OAuthServerProvider` protocol.
+
+```
+mcp = FastMCP("My App",
+        auth_provider=MyOAuthServerProvider(),
+        auth=AuthSettings(
+            issuer_url="https://myapp.com",
+            revocation_options=RevocationOptions(
+                enabled=True,
+            ),
+            client_registration_options=ClientRegistrationOptions(
+                enabled=True,
+                valid_scopes=["myscope", "myotherscope"],
+                default_scopes=["myscope"],
+            ),
+            required_scopes=["myscope"],
+        ),
+)
+```
+
+See [OAuthServerProvider](mcp/server/auth/provider.py) for more details.
+
 ## Running Your Server
 
 ### Development Mode

examples/clients/simple-chatbot/mcp_simple_chatbot/main.py

@@ -323,8 +323,7 @@ async def process_llm_response(self, llm_response: str) -> str:
                                 total = result["total"]
                                 percentage = (progress / total) * 100
                                 logging.info(
-                                    f"Progress: {progress}/{total} "
-                                    f"({percentage:.1f}%)"
+                                    f"Progress: {progress}/{total} ({percentage:.1f}%)"
                                 )
 
                             return f"Tool execution result: {result}"

pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
     "httpx-sse>=0.4",
     "pydantic>=2.7.2,<3.0.0",
     "starlette>=0.27",
+    "python-multipart>=0.0.9",
     "sse-starlette>=1.6.1",
     "pydantic-settings>=2.5.2",
     "uvicorn>=0.23.1; sys_platform != 'emscripten'",
@@ -53,6 +54,7 @@ dev = [
     "pytest-flakefinder>=1.1.0",
     "pytest-xdist>=3.6.1",
     "pytest-examples>=0.0.14",
+    "pytest-pretty>=1.2.0",
 ]
 docs = [
     "mkdocs>=1.6.1",

src/mcp/client/session.py

@@ -8,9 +8,6 @@
 import mcp.types as types
 from mcp.shared.context import RequestContext
 from mcp.shared.message import (
-    ClientMessageMetadata,
-    ResumptionToken,
-    ResumptionTokenUpdateCallback,
     SessionMessage,
 )
 from mcp.shared.session import BaseSession, RequestResponder
@@ -263,16 +260,9 @@ async def call_tool(
         self,
         name: str,
         arguments: dict[str, Any] | None = None,
-        on_resumption_token_update: ResumptionTokenUpdateCallback | None = None,
-        resumption_token: ResumptionToken | None = None,
+        read_timeout_seconds: timedelta | None = None,
     ) -> types.CallToolResult:
         """Send a tools/call request."""
-        metadata = None
-        if on_resumption_token_update or resumption_token:
-            metadata = ClientMessageMetadata(
-                on_resumption_token_update=on_resumption_token_update,
-                resumption_token=resumption_token,
-            )
 
         return await self.send_request(
             types.ClientRequest(
@@ -282,7 +272,7 @@ async def call_tool(
                 )
             ),
             types.CallToolResult,
-            metadata=metadata,
+            request_read_timeout_seconds=read_timeout_seconds,
         )
 
     async def list_prompts(self) -> types.ListPromptsResult:

src/mcp/client/streamable_http.py

@@ -7,7 +7,7 @@
 """
 
 import logging
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncGenerator, Awaitable, Callable
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import timedelta
@@ -16,7 +16,7 @@
 import anyio
 import httpx
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
-from httpx_sse import EventSource, aconnect_sse
+from httpx_sse import EventSource, ServerSentEvent, aconnect_sse
 
 from mcp.shared.message import ClientMessageMetadata, SessionMessage
 from mcp.types import (
@@ -26,15 +26,16 @@
     JSONRPCNotification,
     JSONRPCRequest,
     JSONRPCResponse,
+    RequestId,
 )
 
 logger = logging.getLogger(__name__)
 
 
-MessageOrError = SessionMessage | Exception
-StreamWriter = MemoryObjectSendStream[MessageOrError]
+SessionMessageOrError = SessionMessage | Exception
+StreamWriter = MemoryObjectSendStream[SessionMessageOrError]
 StreamReader = MemoryObjectReceiveStream[SessionMessage]
-
+GetSessionIdCallback = Callable[[], str | None]
 
 MCP_SESSION_ID = "mcp-session-id"
 LAST_EVENT_ID = "last-event-id"
@@ -123,23 +124,21 @@ def _is_initialized_notification(self, message: JSONRPCMessage) -> bool:
             and message.root.method == "notifications/initialized"
         )
 
-    def _extract_session_id_from_response(
+    def _maybe_extract_session_id_from_response(
         self,
         response: httpx.Response,
-        is_initialization: bool,
     ) -> None:
         """Extract and store session ID from response headers."""
-        if is_initialization:
-            new_session_id = response.headers.get(MCP_SESSION_ID)
-            if new_session_id:
-                self.session_id = new_session_id
-                logger.info(f"Received session ID: {self.session_id}")
+        new_session_id = response.headers.get(MCP_SESSION_ID)
+        if new_session_id:
+            self.session_id = new_session_id
+            logger.info(f"Received session ID: {self.session_id}")
 
     async def _handle_sse_event(
         self,
-        sse: Any,
+        sse: ServerSentEvent,
         read_stream_writer: StreamWriter,
-        original_request_id: Any | None = None,
+        original_request_id: RequestId | None = None,
         resumption_callback: Callable[[str], Awaitable[None]] | None = None,
     ) -> bool:
         """Handle an SSE event, returning True if the response is complete."""
@@ -161,7 +160,8 @@ async def _handle_sse_event(
                 if sse.id and resumption_callback:
                     await resumption_callback(sse.id)
 
-                # If this is a response or error, we're done
+                # If this is a response or error return True indicating completion
+                # Otherwise, return False to continue listening
                 return isinstance(message.root, JSONRPCResponse | JSONRPCError)
 
             except Exception as exc:
@@ -262,7 +262,8 @@ async def _handle_post_request(self, ctx: RequestContext) -> None:
                 return
 
             response.raise_for_status()
-            self._extract_session_id_from_response(response, is_initialization)
+            if is_initialization:
+                self._maybe_extract_session_id_from_response(response)
 
             content_type = response.headers.get(CONTENT_TYPE, "").lower()
 
@@ -324,7 +325,7 @@ async def _handle_unexpected_content_type(
     async def _send_session_terminated_error(
         self,
         read_stream_writer: StreamWriter,
-        request_id: Any,
+        request_id: RequestId,
     ) -> None:
         """Send a session terminated error response."""
         jsonrpc_error = JSONRPCError(
@@ -411,16 +412,26 @@ async def streamablehttp_client(
     headers: dict[str, Any] | None = None,
     timeout: timedelta = timedelta(seconds=30),
     sse_read_timeout: timedelta = timedelta(seconds=60 * 5),
-):
+    terminate_on_close: bool = True,
+) -> AsyncGenerator[
+    tuple[
+        MemoryObjectReceiveStream[SessionMessage | Exception],
+        MemoryObjectSendStream[SessionMessage],
+        GetSessionIdCallback,
+    ],
+    None,
+]:
     """
     Client transport for StreamableHTTP.
 
     `sse_read_timeout` determines how long (in seconds) the client will wait for a new
     event before disconnecting. All other HTTP operations are controlled by `timeout`.
 
     Yields:
-        Tuple of (read_stream, write_stream, terminate_callback,
-                  get_session_id_callback)
+        Tuple containing:
+            - read_stream: Stream for reading messages from the server
+            - write_stream: Stream for sending messages to the server
+            - get_session_id_callback: Function to retrieve the current session ID
     """
     transport = StreamableHTTPTransport(url, headers, timeout, sse_read_timeout)
 
@@ -448,9 +459,6 @@ def start_get_stream() -> None:
                         transport.handle_get_stream, client, read_stream_writer
                     )
 
-                async def terminate_session() -> None:
-                    await transport.terminate_session(client)
-
                 tg.start_soon(
                     transport.post_writer,
                     client,
@@ -464,10 +472,11 @@ async def terminate_session() -> None:
                     yield (
                         read_stream,
                         write_stream,
-                        terminate_session,
                         transport.get_session_id,
                     )
                 finally:
+                    if transport.session_id and terminate_on_close:
+                        await transport.terminate_session(client)
                     tg.cancel_scope.cancel()
         finally:
             await read_stream_writer.aclose()

src/mcp/server/auth/__init__.py

@@ -0,0 +1,3 @@
+"""
+MCP OAuth server authorization components.
+"""

src/mcp/server/auth/errors.py

@@ -0,0 +1,8 @@
+from pydantic import ValidationError
+
+
+def stringify_pydantic_error(validation_error: ValidationError) -> str:
+    return "\n".join(
+        f"{'.'.join(str(loc) for loc in e['loc'])}: {e['msg']}"
+        for e in validation_error.errors()
+    )

src/mcp/server/auth/handlers/__init__.py

@@ -0,0 +1,3 @@
+"""
+Request handlers for MCP authorization endpoints.
+"""