Expose InitializeResult to middleware (#2516)

Wrap responder.respond() to capture the InitializeResult before it's
sent to the write stream, then return it through the middleware chain.
This allows middleware (e.g., logging) to access the server's initialize
response, not just the client's request.
This commit is contained in:
Jeremiah Lowin 2025-12-01 20:55:29 -05:00 committed by GitHub
commit 54692c361e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 67 additions and 6 deletions

View file

@ -59,13 +59,37 @@ class MiddlewareServerSession(ServerSession):
from fastmcp.server.middleware.middleware import MiddlewareContext
if isinstance(responder.request.root, mcp.types.InitializeRequest):
# The MCP SDK's ServerSession._received_request() handles the
# initialize request internally by calling responder.respond()
# to send the InitializeResult directly to the write stream, then
# returning None. This bypasses the middleware return path entirely,
# so middleware would only see the request, never the response.
#
# To expose the response to middleware (e.g., for logging server
# capabilities), we wrap responder.respond() to capture the
# InitializeResult before it's sent, then return it from
# call_original_handler so it flows back through the middleware chain.
captured_response: mcp.types.ServerResult | None = None
original_respond = responder.respond
async def capturing_respond(
response: mcp.types.ServerResult,
) -> None:
nonlocal captured_response
captured_response = response
return await original_respond(response)
responder.respond = capturing_respond # type: ignore[method-assign]
async def call_original_handler(
ctx: MiddlewareContext,
) -> None:
return await super(MiddlewareServerSession, self)._received_request(
responder
)
) -> mcp.types.InitializeResult | None:
await super(MiddlewareServerSession, self)._received_request(responder)
if captured_response is not None and isinstance(
captured_response.root, mcp.types.InitializeResult
):
return captured_response.root
return None
async with fastmcp.server.context.Context(
fastmcp=self.fastmcp

View file

@ -150,8 +150,8 @@ class Middleware:
async def on_initialize(
self,
context: MiddlewareContext[mt.InitializeRequest],
call_next: CallNext[mt.InitializeRequest, None],
) -> None:
call_next: CallNext[mt.InitializeRequest, mt.InitializeResult | None],
) -> mt.InitializeResult | None:
return await call_next(context)
async def on_call_tool(

View file

@ -249,3 +249,40 @@ async def test_initialization_middleware_with_state_sharing():
# This test shows the pattern, but actual cross-request state would need
# external storage (Redis, DB, etc.)
# The middleware.tool_state might be None if state doesn't persist
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.serverInfo.name == "TestServer"
assert middleware.initialize_result.protocolVersion is not None
assert middleware.initialize_result.capabilities is not None