Adapt proxy meta-forwarding and tests to SDK v2 request-context shim

This commit is contained in:
Jeremiah Lowin 2026-07-05 20:34:26 -04:00
commit 5a77eab277
No known key found for this signature in database
4 changed files with 40 additions and 46 deletions

View file

@ -185,27 +185,15 @@ class ProxyTool(Tool):
ctx.request_context,
ctx._fastmcp, # weakref to FastMCP, not the Context
)
# Build meta dict from request context
meta: dict[str, Any] | None = None
if hasattr(ctx, "request_context"):
req_ctx = ctx.request_context
# Start with existing meta if present
if hasattr(req_ctx, "meta") and req_ctx.meta:
meta = dict(req_ctx.meta)
# Add task metadata if this is a task request
if (
hasattr(req_ctx, "experimental")
and hasattr(req_ctx.experimental, "is_task")
and req_ctx.experimental.is_task
):
task_metadata = req_ctx.experimental.task_metadata
if task_metadata:
meta = meta or {}
meta["modelcontextprotocol.io/task"] = (
task_metadata.model_dump(
by_alias=True, exclude_none=True
)
)
# Forward the inbound request's `_meta` block (trace context,
# version, etc.) to the backend. In SDK v2 the request context
# exposes the lifted `_meta` dict directly; task submission is a
# first-class params field rather than context state, so there
# is no separate task-metadata injection here.
req_ctx = ctx.request_context
meta: dict[str, Any] | None = (
dict(req_ctx.meta) if req_ctx is not None and req_ctx.meta else None
)
result = await client.call_tool_mcp(
name=backend_name, arguments=arguments, meta=meta

View file

@ -4,6 +4,7 @@ from typing import Any
import mcp_types
import pytest
from mcp.server.context import ServerRequestContext
from fastmcp import Client, FastMCP
from fastmcp.server.context import Context
@ -152,15 +153,18 @@ def mcp_server(recording_middleware):
mcp.add_middleware(recording_middleware)
# Register progress handler
@mcp._mcp_server.progress_notification()
# Register a progress notification handler (v2 API: (ctx, params)).
async def handle_progress(
progress_token: str | int,
progress: float,
total: float | None,
message: str | None,
):
print("HI")
_ctx: ServerRequestContext,
_params: mcp_types.ProgressNotificationParams,
) -> None:
pass
mcp._mcp_server.add_notification_handler(
"notifications/progress",
mcp_types.ProgressNotificationParams,
handle_progress,
)
return mcp

View file

@ -4,6 +4,7 @@ from typing import Any
import mcp_types
import pytest
from mcp.server.context import ServerRequestContext
from fastmcp import Client, FastMCP
from fastmcp.exceptions import ToolError
@ -153,15 +154,18 @@ def mcp_server(recording_middleware):
mcp.add_middleware(recording_middleware)
# Register progress handler
@mcp._mcp_server.progress_notification()
# Register a progress notification handler (v2 API: (ctx, params)).
async def handle_progress(
progress_token: str | int,
progress: float,
total: float | None,
message: str | None,
):
print("HI")
_ctx: ServerRequestContext,
_params: mcp_types.ProgressNotificationParams,
) -> None:
pass
mcp._mcp_server.add_notification_handler(
"notifications/progress",
mcp_types.ProgressNotificationParams,
handle_progress,
)
return mcp

View file

@ -13,7 +13,7 @@ from fastmcp.client.logging import LogMessage
from fastmcp.client.transports import FastMCPTransport
from fastmcp.exceptions import ToolError
from fastmcp.server.context import _current_context
from fastmcp.server.dependencies import get_server, request_ctx
from fastmcp.server.dependencies import fastmcp_request_ctx, get_server
from fastmcp.server.elicitation import AcceptedElicitation
from fastmcp.server.providers.proxy import (
FastMCPProxy,
@ -220,7 +220,7 @@ class TestRestoreRequestContextCurrentServer:
async def _run_in_child_context(self, fn):
# Run in a child task so contextvar writes are isolated from the test
# task and `request_ctx` is genuinely unset (LookupError branch).
# task and `fastmcp_request_ctx` is genuinely unset (defaults to None).
return await asyncio.create_task(fn())
async def test_lookup_error_branch_restores_current_server(self):
@ -231,13 +231,12 @@ class TestRestoreRequestContextCurrentServer:
rc_ref: list = [(rc, weakref.ref(fastmcp))]
async def body():
with pytest.raises(LookupError):
request_ctx.get()
assert fastmcp_request_ctx.get() is None
_restore_request_context(rc_ref)
# The actual Bug 4 fix: get_server() now resolves.
assert get_server() is fastmcp
assert request_ctx.get() is rc
assert fastmcp_request_ctx.get() is rc
ctx = _current_context.get()
assert ctx is not None
@ -264,10 +263,10 @@ class TestRestoreRequestContextCurrentServer:
rc_ref: list = [(fresh_rc, weakref.ref(fastmcp))]
async def body():
request_ctx.set(stale_rc)
fastmcp_request_ctx.set(stale_rc)
_restore_request_context(rc_ref)
assert request_ctx.get() is fresh_rc
assert fastmcp_request_ctx.get() is fresh_rc
assert get_server() is fastmcp
await self._run_in_child_context(body)
@ -278,7 +277,6 @@ class TestRestoreRequestContextCurrentServer:
async def body():
# No stash: nothing restored, no error.
_restore_request_context(rc_ref)
with pytest.raises(LookupError):
request_ctx.get()
assert fastmcp_request_ctx.get() is None
await self._run_in_child_context(body)