refactor: split RequestContext between server and client (#1987)
This commit is contained in:
parent
050aeb6456
commit
b1f7eec3cd
14 changed files with 210 additions and 212 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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...")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue