refactor: split RequestContext between server and client (#1987)

This commit is contained in:
Marcelo Trylesinski 2026-02-03 14:35:07 +01:00 • committed by GitHub
commit b1f7eec3cd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 210 additions and 212 deletions

View file

@ -5,7 +5,6 @@ from mcp import Client
from mcp.client.session import ClientSession
from mcp.server.mcpserver import MCPServer
from mcp.server.mcpserver.server import Context
from mcp.server.session import ServerSession
from mcp.shared.context import RequestContext
from mcp.types import ListRootsResult, Root, TextContent
@ -16,24 +15,18 @@ async def test_list_roots_callback():
callback_return = ListRootsResult(
roots=[
Root(
uri=FileUrl("file://users/fake/test"),
name="Test Root 1",
),
Root(
uri=FileUrl("file://users/fake/test/2"),
name="Test Root 2",
),
Root(uri=FileUrl("file://users/fake/test"), name="Test Root 1"),
Root(uri=FileUrl("file://users/fake/test/2"), name="Test Root 2"),
]
)
async def list_roots_callback(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
) -> ListRootsResult:
return callback_return
@server.tool("test_list_roots")
async def test_list_roots(context: Context[ServerSession, None], message: str):
async def test_list_roots(context: Context[None], message: str):
roots = await context.session.list_roots()
assert roots == callback_return
return True

View file

@ -26,7 +26,7 @@ async def test_sampling_callback():
)
async def sampling_callback(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: CreateMessageRequestParams,
) -> CreateMessageResult:
return callback_return
@ -71,7 +71,7 @@ async def test_create_message_backwards_compat_single_content():
)
async def sampling_callback(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: CreateMessageRequestParams,
) -> CreateMessageResult:
return callback_return

View file

@ -1,4 +1,4 @@
from typing import Any
from __future__ import annotations
import anyio
import pytest
@ -390,7 +390,7 @@ async def test_client_capabilities_with_custom_callbacks():
received_capabilities = None
async def custom_sampling_callback( # pragma: no cover
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult | types.ErrorData:
return types.CreateMessageResult(
@ -400,7 +400,7 @@ async def test_client_capabilities_with_custom_callbacks():
)
async def custom_list_roots_callback( # pragma: no cover
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
) -> types.ListRootsResult | types.ErrorData:
return types.ListRootsResult(roots=[])
@ -474,7 +474,7 @@ async def test_client_capabilities_with_sampling_tools():
received_capabilities = None
async def custom_sampling_callback( # pragma: no cover
context: RequestContext["ClientSession", Any],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult | types.ErrorData:
return types.CreateMessageResult(

View file

@ -88,13 +88,13 @@ async def test_client_capabilities_with_tasks():
# Define custom handlers to trigger capability building (never actually called)
async def my_list_tasks_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.PaginatedRequestParams | None,
) -> types.ListTasksResult | types.ErrorData:
raise NotImplementedError
async def my_cancel_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.CancelTaskRequestParams,
) -> types.CancelTaskResult | types.ErrorData:
raise NotImplementedError
@ -168,13 +168,13 @@ async def test_client_capabilities_auto_built_from_handlers():
# Define custom handlers (not defaults)
async def my_list_tasks_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.PaginatedRequestParams | None,
) -> types.ListTasksResult | types.ErrorData:
raise NotImplementedError
async def my_cancel_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.CancelTaskRequestParams,
) -> types.CancelTaskResult | types.ErrorData:
raise NotImplementedError
@ -249,7 +249,7 @@ async def test_client_capabilities_with_task_augmented_handlers():
# Define task-augmented handler
async def my_augmented_sampling_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
task_metadata: types.TaskMetadata,
) -> types.CreateTaskResult | types.ErrorData:

View file

@ -113,7 +113,7 @@ async def test_client_handles_get_task_request(client_streams: ClientTestStreams
received_task_id: str | None = None
async def get_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskRequestParams,
) -> GetTaskResult | ErrorData:
nonlocal received_task_id
@ -176,7 +176,7 @@ async def test_client_handles_get_task_result_request(client_streams: ClientTest
store = InMemoryTaskStore()
async def get_task_result_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskPayloadRequestParams,
) -> GetTaskPayloadResult | ErrorData:
result = await store.get_result(params.task_id)
@ -239,7 +239,7 @@ async def test_client_handles_list_tasks_request(client_streams: ClientTestStrea
store = InMemoryTaskStore()
async def list_tasks_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: types.PaginatedRequestParams | None,
) -> ListTasksResult | ErrorData:
cursor = params.cursor if params else None
@ -294,7 +294,7 @@ async def test_client_handles_cancel_task_request(client_streams: ClientTestStre
store = InMemoryTaskStore()
async def cancel_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: CancelTaskRequestParams,
) -> CancelTaskResult | ErrorData:
task = await store.get_task(params.task_id)
@ -361,7 +361,7 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
background_tg: list[TaskGroup | None] = [None]
async def task_augmented_sampling_callback(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: CreateMessageRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult:
@ -384,7 +384,7 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
return CreateTaskResult(task=task)
async def get_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskRequestParams,
) -> GetTaskResult | ErrorData:
task = await store.get_task(params.task_id)
@ -400,7 +400,7 @@ async def test_client_task_augmented_sampling(client_streams: ClientTestStreams)
)
async def get_task_result_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskPayloadRequestParams,
) -> GetTaskPayloadResult | ErrorData:
result = await store.get_result(params.task_id)
@ -505,7 +505,7 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
background_tg: list[TaskGroup | None] = [None]
async def task_augmented_elicitation_callback(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: ElicitRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult | ErrorData:
@ -524,7 +524,7 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
return CreateTaskResult(task=task)
async def get_task_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskRequestParams,
) -> GetTaskResult | ErrorData:
task = await store.get_task(params.task_id)
@ -540,7 +540,7 @@ async def test_client_task_augmented_elicitation(client_streams: ClientTestStrea
)
async def get_task_result_handler(
context: RequestContext[ClientSession, None],
context: RequestContext[ClientSession],
params: GetTaskPayloadRequestParams,
) -> GetTaskPayloadResult | ErrorData:
result = await store.get_result(params.task_id)

View file

@ -53,7 +53,7 @@ def create_client_task_handlers(
task_complete_events: dict[str, Event] = {}
async def handle_augmented_elicitation(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: ElicitRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult:
@ -72,7 +72,7 @@ def create_client_task_handlers(
return CreateTaskResult(task=task)
async def handle_get_task(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: Any,
) -> GetTaskResult:
"""Handle tasks/get from server."""
@ -89,7 +89,7 @@ def create_client_task_handlers(
)
async def handle_get_task_result(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: Any,
) -> GetTaskPayloadResult | ErrorData:
"""Handle tasks/result from server."""
@ -121,7 +121,7 @@ def create_sampling_task_handlers(
task_complete_events: dict[str, Event] = {}
async def handle_augmented_sampling(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: CreateMessageRequestParams,
task_metadata: TaskMetadata,
) -> CreateTaskResult:
@ -140,7 +140,7 @@ def create_sampling_task_handlers(
return CreateTaskResult(task=task)
async def handle_get_task(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: Any,
) -> GetTaskResult:
"""Handle tasks/get from server."""
@ -157,7 +157,7 @@ def create_sampling_task_handlers(
)
async def handle_get_task_result(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: Any,
) -> GetTaskPayloadResult | ErrorData:
"""Handle tasks/result from server."""
@ -211,7 +211,7 @@ async def test_scenario1_normal_tool_normal_elicitation() -> None:
# Elicitation callback for client
async def elicitation_callback(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: ElicitRequestParams,
) -> ElicitResult:
elicit_received.set()
@ -379,7 +379,7 @@ async def test_scenario3_task_augmented_tool_normal_elicitation() -> None:
# Elicitation callback for client
async def elicitation_callback(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: ElicitRequestParams,
) -> ElicitResult:
elicit_received.set()

View file

@ -2,8 +2,9 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from mcp.server.context import ServerRequestContext
from mcp.server.experimental.request_context import Experimental
from mcp.server.mcpserver import Context
from mcp.shared.context import RequestContext
pytestmark = pytest.mark.anyio
@ -16,11 +17,12 @@ async def test_progress_token_zero_first_call():
mock_session.send_progress_notification = AsyncMock()
# Create request context with progress token 0
request_context = RequestContext(
request_context = ServerRequestContext(
request_id="test-request",
session=mock_session,
meta={"progress_token": 0},
lifespan_context=None,
experimental=Experimental(),
)
# Create context with our mocks

View file

@ -3,7 +3,6 @@ from contextlib import asynccontextmanager
from dataclasses import dataclass
from mcp.server.mcpserver import Context, MCPServer
from mcp.server.session import ServerSession
class Database: # Replace with your actual DB type
@ -45,7 +44,7 @@ mcp = MCPServer("My App", lifespan=app_lifespan)
# Access type-safe lifespan context in tools
@mcp.tool()
def query_db(ctx: Context[ServerSession, AppContext]) -> str: # pragma: no cover
def query_db(ctx: Context[AppContext]) -> str: # pragma: no cover
"""Tool that uses initialized resources"""
db = ctx.request_context.lifespan_context.db
return db.query()

View file

@ -65,7 +65,7 @@ async def test_stdio_elicitation():
create_ask_user_tool(mcp)
# Create a custom handler for elicitation requests
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
if params.message == "Tool wants to ask: What is your name?":
return ElicitResult(action="accept", content={"answer": "Test User"})
else: # pragma: no cover
@ -82,7 +82,7 @@ async def test_stdio_elicitation_decline():
mcp = MCPServer(name="StdioElicitationDeclineServer")
create_ask_user_tool(mcp)
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
return ElicitResult(action="decline")
await call_tool_and_assert(
@ -121,7 +121,7 @@ async def test_elicitation_schema_validation():
# Dummy callback (won't be called due to validation failure)
async def elicitation_callback(
context: RequestContext[ClientSession, None], params: ElicitRequestParams
context: RequestContext[ClientSession], params: ElicitRequestParams
): # pragma: no cover
return ElicitResult(action="accept", content={})
@ -177,7 +177,7 @@ async def test_elicitation_with_optional_fields():
for content, expected in test_cases:
async def callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
return ElicitResult(action="accept", content=content)
await call_tool_and_assert(mcp, callback, "optional_tool", {}, expected)
@ -196,7 +196,7 @@ async def test_elicitation_with_optional_fields():
return f"Validation failed: {str(e)}"
async def elicitation_callback(
context: RequestContext[ClientSession, None], params: ElicitRequestParams
context: RequestContext[ClientSession], params: ElicitRequestParams
): # pragma: no cover
return ElicitResult(action="accept", content={})
@ -220,7 +220,7 @@ async def test_elicitation_with_optional_fields():
return f"Name: {result.data.name}, Tags: {', '.join(result.data.tags)}"
return f"User {result.action}" # pragma: no cover
async def multiselect_callback(context: RequestContext[ClientSession, Any], params: ElicitRequestParams):
async def multiselect_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
if "Please provide tags" in params.message:
return ElicitResult(action="accept", content={"name": "Test", "tags": ["tag1", "tag2"]})
return ElicitResult(action="decline") # pragma: no cover
@ -240,7 +240,7 @@ async def test_elicitation_with_optional_fields():
return f"Name: {result.data.name}, Tags: {tags_str}"
return f"User {result.action}" # pragma: no cover
async def optional_multiselect_callback(context: RequestContext[ClientSession, Any], params: ElicitRequestParams):
async def optional_multiselect_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
if "Please provide optional tags" in params.message:
return ElicitResult(action="accept", content={"name": "Test", "tags": ["tag1", "tag2"]})
return ElicitResult(action="decline") # pragma: no cover
@ -274,7 +274,7 @@ async def test_elicitation_with_default_values():
return f"User {result.action}"
# First verify that defaults are present in the JSON schema sent to clients
async def callback_schema_verify(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def callback_schema_verify(context: RequestContext[ClientSession], params: ElicitRequestParams):
# Verify the schema includes defaults
assert isinstance(params, types.ElicitRequestFormParams), "Expected form mode elicitation"
schema = params.requested_schema
@ -296,7 +296,7 @@ async def test_elicitation_with_default_values():
)
# Test overriding defaults
async def callback_override(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def callback_override(context: RequestContext[ClientSession], params: ElicitRequestParams):
return ElicitResult(
action="accept", content={"email": "john@example.com", "name": "John", "age": 25, "subscribe": False}
)
@ -372,7 +372,7 @@ async def test_elicitation_with_enum_titles():
return f"User: {result.data.user_name}, Color: {result.data.color}"
return f"User {result.action}" # pragma: no cover
async def enum_callback(context: RequestContext[ClientSession, Any], params: ElicitRequestParams):
async def enum_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
if "colors" in params.message and "legacy" not in params.message:
return ElicitResult(action="accept", content={"user_name": "Bob", "favorite_colors": ["red", "green"]})
elif "color" in params.message:

View file

@ -185,7 +185,7 @@ def create_client_for_transport(transport: str, server_url: str):
# Callback functions for testing
async def sampling_callback(
context: RequestContext[ClientSession, None], params: CreateMessageRequestParams
context: RequestContext[ClientSession], params: CreateMessageRequestParams
) -> CreateMessageResult:
"""Sampling callback for tests."""
return CreateMessageResult(
@ -198,7 +198,7 @@ async def sampling_callback(
)
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
"""Elicitation callback for tests."""
# For restaurant booking test
if "No tables available" in params.message:

View file

@ -6,12 +6,12 @@ from typing import Any, TypedDict
import pytest
from pydantic import BaseModel
from mcp.server.context import LifespanContextT, RequestT
from mcp.server.mcpserver import Context, MCPServer
from mcp.server.mcpserver.exceptions import ToolError
from mcp.server.mcpserver.tools import Tool, ToolManager
from mcp.server.mcpserver.utilities.func_metadata import ArgModelBase, FuncMetadata
from mcp.server.session import ServerSessionT
from mcp.shared.context import LifespanContextT, RequestT
from mcp.types import TextContent, ToolAnnotations
@ -347,9 +347,7 @@ class TestContextHandling:
tool = manager.add_tool(tool_without_context)
assert tool.context_kwarg is None
def tool_with_parametrized_context(
x: int, ctx: Context[ServerSessionT, LifespanContextT, RequestT]
) -> str: # pragma: no cover
def tool_with_parametrized_context(x: int, ctx: Context[LifespanContextT, RequestT]) -> str: # pragma: no cover
return str(x)
tool = manager.add_tool(tool_with_parametrized_context)

View file

@ -29,7 +29,7 @@ async def test_url_elicitation_accept():
return f"User {result.action}"
# Create elicitation callback that accepts URL mode
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
assert params.mode == "url"
assert params.url == "https://example.com/api_key_setup"
assert params.elicitation_id == "test-elicitation-001"
@ -58,7 +58,7 @@ async def test_url_elicitation_decline():
# Test only checks decline path
return f"User {result.action} authorization"
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
assert params.mode == "url"
return ElicitResult(action="decline")
@ -84,7 +84,7 @@ async def test_url_elicitation_cancel():
# Test only checks cancel path
return f"User {result.action} payment"
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
assert params.mode == "url"
return ElicitResult(action="cancel")
@ -111,7 +111,7 @@ async def test_url_elicitation_helper_function():
# Test only checks accept path - return the type name
return type(result).__name__
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
return ElicitResult(action="accept")
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
@ -138,7 +138,7 @@ async def test_url_no_content_in_response():
assert result.content is None
return f"Action: {result.action}, Content: {result.content}"
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
# Verify that this is URL mode
assert params.mode == "url"
assert isinstance(params, types.ElicitRequestURLParams)
@ -171,7 +171,7 @@ async def test_form_mode_still_works():
assert result.data is not None
return f"Hello, {result.data.name}!"
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
# Verify form mode parameters
assert params.mode == "form"
assert isinstance(params, types.ElicitRequestFormParams)
@ -207,7 +207,7 @@ async def test_elicit_complete_notification():
return "Elicitation completed"
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
return ElicitResult(action="accept") # pragma: no cover
async with Client(mcp, elicitation_callback=elicitation_callback) as client:
@ -264,7 +264,7 @@ async def test_elicit_url_typed_results():
return "Not cancelled" # pragma: no cover
# Test declined result
async def decline_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def decline_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
return ElicitResult(action="decline")
async with Client(mcp, elicitation_callback=decline_callback) as client:
@ -274,7 +274,7 @@ async def test_elicit_url_typed_results():
assert result.content[0].text == "Declined"
# Test cancelled result
async def cancel_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def cancel_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
return ElicitResult(action="cancel")
async with Client(mcp, elicitation_callback=cancel_callback) as client:
@ -304,7 +304,7 @@ async def test_deprecated_elicit_method():
return f"Email: {result.content.get('email', 'none')}"
return "No email provided" # pragma: no cover
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
# Verify this is form mode
assert params.mode == "form"
assert params.requested_schema is not None
@ -332,7 +332,7 @@ async def test_ctx_elicit_url_convenience_method():
)
return f"Result: {result.action}"
async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
assert params.mode == "url"
assert params.elicitation_id == "ctx-test-001"
return ElicitResult(action="accept")

View file

@ -1,4 +1,4 @@
from typing import Any, cast
from typing import Any
from unittest.mock import patch
import anyio
@ -14,7 +14,7 @@ from mcp.server.session import ServerSession
from mcp.shared.context import RequestContext
from mcp.shared.message import SessionMessage
from mcp.shared.progress import progress
from mcp.shared.session import BaseSession, RequestResponder
from mcp.shared.session import RequestResponder
@pytest.mark.anyio
@ -279,14 +279,10 @@ async def test_progress_context_manager():
request_id="test-request",
session=client_session,
meta={"progress_token": progress_token},
lifespan_context=None,
)
# cast for type checker
typed_context = cast(RequestContext[BaseSession[Any, Any, Any, Any, Any], Any], request_context)
# Utilize progress context manager
with progress(typed_context, total=100) as p:
with progress(request_context, total=100) as p:
await p.progress(10, message="Loading configuration...")
await p.progress(30, message="Connecting to database...")
await p.progress(40, message="Fetching data...")

View file

@ -1344,7 +1344,7 @@ async def test_streamablehttp_server_sampling(basic_server: None, basic_server_u
# Define sampling callback that returns a mock response
async def sampling_callback(
context: RequestContext[ClientSession, Any],
context: RequestContext[ClientSession],
params: types.CreateMessageRequestParams,
) -> types.CreateMessageResult:
nonlocal sampling_callback_invoked, captured_message_params