chore: import upstream snapshot with attribution

tests/server/middleware/test_initialization_middleware.py

@@ -0,0 +1,428 @@
+"""Tests for middleware support during initialization."""
+
+from collections.abc import Sequence
+from typing import Any
+
+import mcp_types as mt
+import pytest
+from mcp import MCPError
+from mcp_types import TextContent
+
+from fastmcp import Client, FastMCP
+from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
+from fastmcp.tools.base import Tool
+
+
+class InitializationMiddleware(Middleware):
+    """Middleware that captures initialization details.
+
+    Note: Session state is NOT available during on_initialize because
+    the MCP session has not been established yet. Use instance variables
+    to store data that needs to persist across the session.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.initialized = False
+        self.client_info = None
+        self.session_data = {}
+
+    async def on_initialize(
+        self,
+        context: MiddlewareContext[mt.InitializeRequest],
+        call_next: CallNext[mt.InitializeRequest, mt.InitializeResult | None],
+    ) -> mt.InitializeResult | None:
+        """Capture initialization details."""
+        self.initialized = True
+
+        # Extract client info from the initialize params
+        if hasattr(context.message, "params") and hasattr(
+            context.message.params, "clientInfo"
+        ):
+            self.client_info = context.message.params.client_info
+
+        # Store in instance for cross-request access
+        # (session state is not available during on_initialize)
+        self.session_data["client_initialized"] = True
+        if self.client_info:
+            self.session_data["client_name"] = getattr(
+                self.client_info, "name", "unknown"
+            )
+
+        return await call_next(context)
+
+
+class ClientDetectionMiddleware(Middleware):
+    """Middleware that detects specific clients and modifies behavior.
+
+    This demonstrates storing data in the middleware instance itself
+    for cross-request access, since context state is request-scoped.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.is_test_client = False
+        self.tools_modified = False
+        self.initialization_called = False
+
+    async def on_initialize(
+        self,
+        context: MiddlewareContext[mt.InitializeRequest],
+        call_next: CallNext[mt.InitializeRequest, mt.InitializeResult | None],
+    ) -> mt.InitializeResult | None:
+        """Detect test client during initialization."""
+        self.initialization_called = True
+
+        # For testing purposes, always set it to true
+        # Store in instance variable for cross-request access
+        self.is_test_client = True
+
+        return await call_next(context)
+
+    async def on_list_tools(
+        self,
+        context: MiddlewareContext[mt.ListToolsRequest],
+        call_next: CallNext[mt.ListToolsRequest, Sequence[Tool]],
+    ) -> Sequence[Tool]:
+        """Modify tools based on client detection."""
+        tools = await call_next(context)
+
+        # Use the instance variable set during initialization
+        if self.is_test_client:
+            # Add a special annotation to tools for test clients
+            for tool in tools:
+                if not hasattr(tool, "annotations"):
+                    tool.annotations = mt.ToolAnnotations()
+                if tool.annotations is None:
+                    tool.annotations = mt.ToolAnnotations()
+                # Mark as read-only for test clients
+                tool.annotations.read_only_hint = True
+            self.tools_modified = True
+
+        return tools
+
+
+async def test_simple_initialization_hook():
+    """Test that the on_initialize hook is called."""
+    server = FastMCP("TestServer")
+
+    class SimpleInitMiddleware(Middleware):
+        def __init__(self):
+            super().__init__()
+            self.called = False
+
+        async def on_initialize(
+            self,
+            context: MiddlewareContext[mt.InitializeRequest],
+            call_next: CallNext[mt.InitializeRequest, mt.InitializeResult | None],
+        ) -> mt.InitializeResult | None:
+            self.called = True
+            return await call_next(context)
+
+    middleware = SimpleInitMiddleware()
+    server.add_middleware(middleware)
+
+    # Connect client
+    async with Client(server):
+        # Middleware should have been called
+        assert middleware.called is True, "on_initialize was not called"
+
+
+async def test_middleware_receives_initialization():
+    """Test that middleware can intercept initialization requests."""
+    server = FastMCP("TestServer")
+    middleware = InitializationMiddleware()
+    server.add_middleware(middleware)
+
+    @server.tool
+    def test_tool(x: int) -> str:
+        return f"Result: {x}"
+
+    # Connect client
+    async with Client(server) as client:
+        # Middleware should have been called during initialization
+        assert middleware.initialized is True
+
+        # Test that the tool still works
+        result = await client.call_tool("test_tool", {"x": 42})
+        assert isinstance(result.content[0], TextContent)
+        assert result.content[0].text == "Result: 42"
+
+
+async def test_client_detection_middleware():
+    """Test middleware that detects specific clients and modifies behavior."""
+    server = FastMCP("TestServer")
+    middleware = ClientDetectionMiddleware()
+    server.add_middleware(middleware)
+
+    @server.tool
+    def example_tool() -> str:
+        return "example"
+
+    # Connect with a client
+    async with Client(server) as client:
+        # Middleware should have been called during initialization
+        assert middleware.initialization_called is True
+        assert middleware.is_test_client is True
+
+        # List tools to trigger modification
+        tools = await client.list_tools()
+        assert len(tools) == 1
+        assert middleware.tools_modified is True
+
+        # Check that the tool has the modified annotation
+        tool = tools[0]
+        assert tool.annotations is not None
+        assert tool.annotations.read_only_hint is True
+
+
+async def test_multiple_middleware_initialization():
+    """Test that multiple middleware can handle initialization."""
+    server = FastMCP("TestServer")
+
+    init_mw = InitializationMiddleware()
+    detect_mw = ClientDetectionMiddleware()
+
+    server.add_middleware(init_mw)
+    server.add_middleware(detect_mw)
+
+    @server.tool
+    def test_tool() -> str:
+        return "test"
+
+    async with Client(server) as client:
+        # Both middleware should have processed initialization
+        assert init_mw.initialized is True
+        assert detect_mw.initialization_called is True
+        assert detect_mw.is_test_client is True
+
+        # List tools to check detection worked
+        await client.list_tools()
+        assert detect_mw.tools_modified is True
+
+
+async def test_session_state_persists_across_tool_calls():
+    """Test that session-scoped state persists across multiple tool calls.
+
+    Session state is only available after the session is established,
+    so it can't be set during on_initialize. This test shows state set
+    during one tool call is accessible in subsequent tool calls.
+    """
+    server = FastMCP("TestServer")
+
+    class StateTrackingMiddleware(Middleware):
+        def __init__(self):
+            super().__init__()
+            self.call_count = 0
+            self.state_values = []
+
+        async def on_call_tool(
+            self,
+            context: MiddlewareContext[mt.CallToolRequestParams],
+            call_next: CallNext[mt.CallToolRequestParams, Any],
+        ) -> Any:
+            self.call_count += 1
+
+            if context.fastmcp_context:
+                # Read existing state
+                counter = await context.fastmcp_context.get_state("call_counter")
+                self.state_values.append(counter)
+
+                # Increment and save
+                new_counter = (counter or 0) + 1
+                await context.fastmcp_context.set_state("call_counter", new_counter)
+
+            return await call_next(context)
+
+    middleware = StateTrackingMiddleware()
+    server.add_middleware(middleware)
+
+    @server.tool
+    def test_tool() -> str:
+        return "success"
+
+    async with Client(server) as client:
+        # First call - state should be None initially
+        result = await client.call_tool("test_tool", {})
+        assert isinstance(result.content[0], TextContent)
+        assert result.content[0].text == "success"
+
+        # Second call - state should show previous value (1)
+        result = await client.call_tool("test_tool", {})
+        assert isinstance(result.content[0], TextContent)
+
+        # Third call - state should show previous value (2)
+        result = await client.call_tool("test_tool", {})
+        assert isinstance(result.content[0], TextContent)
+
+        # Verify state persisted across calls within the session
+        assert middleware.call_count == 3
+        # First call saw None, second saw 1, third saw 2
+        assert middleware.state_values == [None, 1, 2]
+
+
+async def test_middleware_can_access_initialize_result():
+    """Test that middleware can access the InitializeResult from call_next().
+
+    This verifies that the initialize response is returned through the middleware
+    chain, not just sent directly via the responder (fixes #2504).
+    """
+    server = FastMCP("TestServer")
+
+    class ResponseCapturingMiddleware(Middleware):
+        def __init__(self):
+            super().__init__()
+            self.initialize_result: mt.InitializeResult | None = None
+
+        async def on_initialize(
+            self,
+            context: MiddlewareContext[mt.InitializeRequest],
+            call_next: CallNext[mt.InitializeRequest, mt.InitializeResult | None],
+        ) -> mt.InitializeResult | None:
+            # Call next and capture the result
+            result = await call_next(context)
+            self.initialize_result = result
+            return result
+
+    middleware = ResponseCapturingMiddleware()
+    server.add_middleware(middleware)
+
+    async with Client(server):
+        # Middleware should have captured the InitializeResult
+        assert middleware.initialize_result is not None
+        assert isinstance(middleware.initialize_result, mt.InitializeResult)
+
+        # Verify the result contains expected server info
+        assert middleware.initialize_result.server_info.name == "TestServer"
+        assert middleware.initialize_result.protocol_version is not None
+        assert middleware.initialize_result.capabilities is not None
+
+
+async def test_middleware_mcp_error_during_initialization():
+    """Test that MCPError raised in middleware during initialization is sent to client."""
+    server = FastMCP("TestServer")
+
+    class ErrorThrowingMiddleware(Middleware):
+        async def on_initialize(
+            self,
+            context: MiddlewareContext[mt.InitializeRequest],
+            call_next: CallNext[mt.InitializeRequest, mt.InitializeResult | None],
+        ) -> mt.InitializeResult | None:
+            raise MCPError(
+                code=mt.INVALID_PARAMS, message="Invalid initialization parameters"
+            )
+
+    server.add_middleware(ErrorThrowingMiddleware())
+
+    with pytest.raises(MCPError) as exc_info:
+        async with Client(server):
+            pass
+
+    assert exc_info.value.error.message == "Invalid initialization parameters"
+    assert exc_info.value.error.code == mt.INVALID_PARAMS
+
+
+async def test_middleware_mcp_error_before_call_next():
+    """Test MCPError raised before calling next middleware."""
+    server = FastMCP("TestServer")
+
+    class EarlyErrorMiddleware(Middleware):
+        async def on_initialize(
+            self,
+            context: MiddlewareContext[mt.InitializeRequest],
+            call_next: CallNext[mt.InitializeRequest, mt.InitializeResult | None],
+        ) -> mt.InitializeResult | None:
+            raise MCPError(code=mt.INVALID_REQUEST, message="Request validation failed")
+
+    server.add_middleware(EarlyErrorMiddleware())
+
+    with pytest.raises(MCPError) as exc_info:
+        async with Client(server):
+            pass
+
+    assert exc_info.value.error.message == "Request validation failed"
+    assert exc_info.value.error.code == mt.INVALID_REQUEST
+
+
+async def test_middleware_mcp_error_after_call_next():
+    """Test that MCPError raised after call_next doesn't break the connection.
+
+    When an error is raised after call_next, the responder has already completed,
+    so the error is caught but not sent to the responder (checked via _completed flag).
+    """
+    server = FastMCP("TestServer")
+
+    class PostProcessingErrorMiddleware(Middleware):
+        def __init__(self):
+            super().__init__()
+            self.error_raised = False
+
+        async def on_initialize(
+            self,
+            context: MiddlewareContext[mt.InitializeRequest],
+            call_next: CallNext[mt.InitializeRequest, mt.InitializeResult | None],
+        ) -> mt.InitializeResult | None:
+            await call_next(context)
+            self.error_raised = True
+            raise MCPError(code=mt.INTERNAL_ERROR, message="Post-processing failed")
+
+    middleware = PostProcessingErrorMiddleware()
+    server.add_middleware(middleware)
+
+    # Error is logged but not re-raised to prevent duplicate response
+    async with Client(server):
+        pass
+
+    assert middleware.error_raised is True
+
+
+async def test_state_isolation_between_streamable_http_clients():
+    """Test that different HTTP clients have isolated session state.
+
+    Each client should have its own session ID and isolated state.
+    """
+    from fastmcp.client.transports import StreamableHttpTransport
+    from fastmcp.server.context import Context
+    from fastmcp.utilities.tests import run_server_async
+
+    server = FastMCP("TestServer")
+
+    @server.tool
+    async def store_and_read(value: str, ctx: Context) -> dict:
+        """Store a value and return session info."""
+        existing = await ctx.get_state("client_value")
+        await ctx.set_state("client_value", value)
+        return {
+            "existing": existing,
+            "stored": value,
+            "session_id": ctx.session_id,
+        }
+
+    async with run_server_async(server, transport="streamable-http") as url:
+        import json
+
+        # Client 1 stores its value
+        transport1 = StreamableHttpTransport(url=url)
+        async with Client(transport=transport1) as client1:
+            result1 = await client1.call_tool(
+                "store_and_read", {"value": "client1-value"}
+            )
+            data1 = json.loads(result1.content[0].text)
+            assert data1["existing"] is None
+            assert data1["stored"] == "client1-value"
+            session_id_1 = data1["session_id"]
+
+        # Client 2 should have completely isolated state
+        transport2 = StreamableHttpTransport(url=url)
+        async with Client(transport=transport2) as client2:
+            result2 = await client2.call_tool(
+                "store_and_read", {"value": "client2-value"}
+            )
+            data2 = json.loads(result2.content[0].text)
+            # Should NOT see client1's value
+            assert data2["existing"] is None
+            assert data2["stored"] == "client2-value"
+            session_id_2 = data2["session_id"]
+
+        # Session IDs should be different
+        assert session_id_1 != session_id_2