mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-28 02:10:38 +02:00
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:
parent
d5ef4132b7
commit
54692c361e
3 changed files with 67 additions and 6 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue