"""Tests for server-side task support (handlers, capabilities, integration).""" from datetime import datetime, timezone from typing import Any import anyio import pytest import mcp.types as mcp_types from mcp.client.session import ClientSession from mcp.server import Server from mcp.server.lowlevel import NotificationOptions from mcp.server.models import InitializationOptions from mcp.server.session import ServerSession from mcp.shared.context import RequestContext from mcp.shared.exceptions import MCPError from mcp.shared.message import ServerMessageMetadata, SessionMessage from mcp.shared.response_router import ResponseRouter from mcp.shared.session import RequestResponder from mcp.types import ( INVALID_REQUEST, TASK_FORBIDDEN, TASK_OPTIONAL, TASK_REQUIRED, CallToolRequest, CallToolRequestParams, CallToolResult, CancelTaskRequest, CancelTaskRequestParams, CancelTaskResult, ClientResult, ErrorData, GetTaskPayloadRequest, GetTaskPayloadRequestParams, GetTaskPayloadResult, GetTaskRequest, GetTaskRequestParams, GetTaskResult, JSONRPCError, JSONRPCNotification, JSONRPCResponse, ListTasksRequest, ListTasksResult, ListToolsRequest, ListToolsResult, SamplingMessage, ServerCapabilities, ServerNotification, ServerRequest, ServerResult, Task, TaskMetadata, TextContent, Tool, ToolExecution, ) @pytest.mark.anyio async def test_list_tasks_handler() -> None: """Test that experimental list_tasks handler works.""" server = Server("test") now = datetime.now(timezone.utc) test_tasks = [ Task( task_id="task-1", status="working", created_at=now, last_updated_at=now, ttl=60000, poll_interval=1000, ), Task( task_id="task-2", status="completed", created_at=now, last_updated_at=now, ttl=60000, poll_interval=1000, ), ] @server.experimental.list_tasks() async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult: return ListTasksResult(tasks=test_tasks) handler = server.request_handlers[ListTasksRequest] request = ListTasksRequest(method="tasks/list") result = await handler(request) assert isinstance(result, ServerResult) assert isinstance(result, ListTasksResult) assert len(result.tasks) == 2 assert result.tasks[0].task_id == "task-1" assert result.tasks[1].task_id == "task-2" @pytest.mark.anyio async def test_get_task_handler() -> None: """Test that experimental get_task handler works.""" server = Server("test") @server.experimental.get_task() async def handle_get_task(request: GetTaskRequest) -> GetTaskResult: now = datetime.now(timezone.utc) return GetTaskResult( task_id=request.params.task_id, status="working", created_at=now, last_updated_at=now, ttl=60000, poll_interval=1000, ) handler = server.request_handlers[GetTaskRequest] request = GetTaskRequest( method="tasks/get", params=GetTaskRequestParams(task_id="test-task-123"), ) result = await handler(request) assert isinstance(result, ServerResult) assert isinstance(result, GetTaskResult) assert result.task_id == "test-task-123" assert result.status == "working" @pytest.mark.anyio async def test_get_task_result_handler() -> None: """Test that experimental get_task_result handler works.""" server = Server("test") @server.experimental.get_task_result() async def handle_get_task_result(request: GetTaskPayloadRequest) -> GetTaskPayloadResult: return GetTaskPayloadResult() handler = server.request_handlers[GetTaskPayloadRequest] request = GetTaskPayloadRequest( method="tasks/result", params=GetTaskPayloadRequestParams(task_id="test-task-123"), ) result = await handler(request) assert isinstance(result, ServerResult) assert isinstance(result, GetTaskPayloadResult) @pytest.mark.anyio async def test_cancel_task_handler() -> None: """Test that experimental cancel_task handler works.""" server = Server("test") @server.experimental.cancel_task() async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult: now = datetime.now(timezone.utc) return CancelTaskResult( task_id=request.params.task_id, status="cancelled", created_at=now, last_updated_at=now, ttl=60000, ) handler = server.request_handlers[CancelTaskRequest] request = CancelTaskRequest( method="tasks/cancel", params=CancelTaskRequestParams(task_id="test-task-123"), ) result = await handler(request) assert isinstance(result, ServerResult) assert isinstance(result, CancelTaskResult) assert result.task_id == "test-task-123" assert result.status == "cancelled" @pytest.mark.anyio async def test_server_capabilities_include_tasks() -> None: """Test that server capabilities include tasks when handlers are registered.""" server = Server("test") @server.experimental.list_tasks() async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult: raise NotImplementedError @server.experimental.cancel_task() async def handle_cancel_task(request: CancelTaskRequest) -> CancelTaskResult: raise NotImplementedError capabilities = server.get_capabilities( notification_options=NotificationOptions(), experimental_capabilities={}, ) assert capabilities.tasks is not None assert capabilities.tasks.list is not None assert capabilities.tasks.cancel is not None assert capabilities.tasks.requests is not None assert capabilities.tasks.requests.tools is not None @pytest.mark.anyio async def test_server_capabilities_partial_tasks() -> None: """Test capabilities with only some task handlers registered.""" server = Server("test") @server.experimental.list_tasks() async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult: raise NotImplementedError # Only list_tasks registered, not cancel_task capabilities = server.get_capabilities( notification_options=NotificationOptions(), experimental_capabilities={}, ) assert capabilities.tasks is not None assert capabilities.tasks.list is not None assert capabilities.tasks.cancel is None # Not registered @pytest.mark.anyio async def test_tool_with_task_execution_metadata() -> None: """Test that tools can declare task execution mode.""" from mcp import Client async def on_list_tools( ctx: RequestContext[ServerSession, Any, Any], params: mcp_types.PaginatedRequestParams | None, ) -> mcp_types.ListToolsResult: return mcp_types.ListToolsResult( tools=[ Tool( name="quick_tool", description="Fast tool", input_schema={"type": "object", "properties": {}}, execution=ToolExecution(task_support=TASK_FORBIDDEN), ), Tool( name="long_tool", description="Long running tool", input_schema={"type": "object", "properties": {}}, execution=ToolExecution(task_support=TASK_REQUIRED), ), Tool( name="flexible_tool", description="Can be either", input_schema={"type": "object", "properties": {}}, execution=ToolExecution(task_support=TASK_OPTIONAL), ), ] ) server = Server("test", on_list_tools=on_list_tools) async with Client(server) as client: result = await client.list_tools() tools = result.tools assert tools[0].execution is not None assert tools[0].execution.task_support == TASK_FORBIDDEN assert tools[1].execution is not None assert tools[1].execution.task_support == TASK_REQUIRED assert tools[2].execution is not None assert tools[2].execution.task_support == TASK_OPTIONAL @pytest.mark.anyio async def test_task_metadata_in_call_tool_request() -> None: """Test that task metadata is accessible via RequestContext when calling a tool.""" captured_task_metadata: TaskMetadata | None = None async def on_list_tools( ctx: RequestContext[ServerSession, Any, Any], params: mcp_types.PaginatedRequestParams | None, ) -> mcp_types.ListToolsResult: return mcp_types.ListToolsResult( tools=[ Tool( name="long_task", description="A long running task", input_schema={"type": "object", "properties": {}}, execution=ToolExecution(task_support="optional"), ) ] ) async def on_call_tool( ctx: RequestContext[ServerSession, Any, Any], params: mcp_types.CallToolRequestParams, ) -> mcp_types.CallToolResult: nonlocal captured_task_metadata captured_task_metadata = ctx.experimental.task_metadata return mcp_types.CallToolResult(content=[TextContent(type="text", text="done")]) server = Server("test", on_list_tools=on_list_tools, on_call_tool=on_call_tool) server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) async def message_handler( message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, ) -> None: ... # pragma: no branch async def run_server(): async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=server.get_capabilities( notification_options=NotificationOptions(), experimental_capabilities={}, ), ), ) as server_session: async with anyio.create_task_group() as tg: async def handle_messages(): async for message in server_session.incoming_messages: # pragma: no branch await server._handle_message(message, server_session, {}, False) tg.start_soon(handle_messages) await anyio.sleep_forever() async with anyio.create_task_group() as tg: tg.start_soon(run_server) async with ClientSession( server_to_client_receive, client_to_server_send, message_handler=message_handler, ) as client_session: await client_session.initialize() # Call tool with task metadata await client_session.send_request( CallToolRequest( params=CallToolRequestParams( name="long_task", arguments={}, task=TaskMetadata(ttl=60000), ), ), CallToolResult, ) tg.cancel_scope.cancel() assert captured_task_metadata is not None assert captured_task_metadata.ttl == 60000 @pytest.mark.anyio async def test_task_metadata_is_task_property() -> None: """Test that RequestContext.experimental.is_task works correctly.""" is_task_values: list[bool] = [] async def on_list_tools( ctx: RequestContext[ServerSession, Any, Any], params: mcp_types.PaginatedRequestParams | None, ) -> mcp_types.ListToolsResult: return mcp_types.ListToolsResult( tools=[ Tool( name="test_tool", description="Test tool", input_schema={"type": "object", "properties": {}}, ) ] ) async def on_call_tool( ctx: RequestContext[ServerSession, Any, Any], params: mcp_types.CallToolRequestParams, ) -> mcp_types.CallToolResult: is_task_values.append(ctx.experimental.is_task) return mcp_types.CallToolResult(content=[TextContent(type="text", text="done")]) server = Server("test", on_list_tools=on_list_tools, on_call_tool=on_call_tool) server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) async def message_handler( message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, ) -> None: ... # pragma: no branch async def run_server(): async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=server.get_capabilities( notification_options=NotificationOptions(), experimental_capabilities={}, ), ), ) as server_session: async with anyio.create_task_group() as tg: async def handle_messages(): async for message in server_session.incoming_messages: # pragma: no branch await server._handle_message(message, server_session, {}, False) tg.start_soon(handle_messages) await anyio.sleep_forever() async with anyio.create_task_group() as tg: tg.start_soon(run_server) async with ClientSession( server_to_client_receive, client_to_server_send, message_handler=message_handler, ) as client_session: await client_session.initialize() # Call without task metadata await client_session.send_request( CallToolRequest(params=CallToolRequestParams(name="test_tool", arguments={})), CallToolResult, ) # Call with task metadata await client_session.send_request( CallToolRequest( params=CallToolRequestParams(name="test_tool", arguments={}, task=TaskMetadata(ttl=60000)), ), CallToolResult, ) tg.cancel_scope.cancel() assert len(is_task_values) == 2 assert is_task_values[0] is False # First call without task assert is_task_values[1] is True # Second call with task @pytest.mark.anyio async def test_update_capabilities_no_handlers() -> None: """Test that update_capabilities returns early when no task handlers are registered.""" server = Server("test-no-handlers") # Access experimental to initialize it, but don't register any task handlers _ = server.experimental caps = server.get_capabilities(NotificationOptions(), {}) # Without any task handlers registered, tasks capability should be None assert caps.tasks is None @pytest.mark.anyio async def test_default_task_handlers_via_enable_tasks() -> None: """Test that enable_tasks() auto-registers working default handlers. This exercises the default handlers in lowlevel/experimental.py: - _default_get_task (task not found) - _default_get_task_result - _default_list_tasks - _default_cancel_task """ server = Server("test-default-handlers") # Enable tasks with default handlers (no custom handlers registered) task_support = server.experimental.enable_tasks() store = task_support.store server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) async def message_handler( message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception, ) -> None: ... # pragma: no branch async def run_server() -> None: async with task_support.run(): async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=server.get_capabilities( notification_options=NotificationOptions(), experimental_capabilities={}, ), ), ) as server_session: task_support.configure_session(server_session) async for message in server_session.incoming_messages: await server._handle_message(message, server_session, {}, False) async with anyio.create_task_group() as tg: tg.start_soon(run_server) async with ClientSession( server_to_client_receive, client_to_server_send, message_handler=message_handler, ) as client_session: await client_session.initialize() # Create a task directly in the store for testing task = await store.create_task(TaskMetadata(ttl=60000)) # Test list_tasks (default handler) list_result = await client_session.send_request(ListTasksRequest(), ListTasksResult) assert len(list_result.tasks) == 1 assert list_result.tasks[0].task_id == task.task_id # Test get_task (default handler - found) get_result = await client_session.send_request( GetTaskRequest(params=GetTaskRequestParams(task_id=task.task_id)), GetTaskResult, ) assert get_result.task_id == task.task_id assert get_result.status == "working" # Test get_task (default handler - not found path) with pytest.raises(MCPError, match="not found"): await client_session.send_request( GetTaskRequest(params=GetTaskRequestParams(task_id="nonexistent-task")), GetTaskResult, ) # Create a completed task to test get_task_result completed_task = await store.create_task(TaskMetadata(ttl=60000)) await store.store_result( completed_task.task_id, CallToolResult(content=[TextContent(type="text", text="Test result")]) ) await store.update_task(completed_task.task_id, status="completed") # Test get_task_result (default handler) payload_result = await client_session.send_request( GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(task_id=completed_task.task_id)), GetTaskPayloadResult, ) # The result should have the related-task metadata assert payload_result.meta is not None assert "io.modelcontextprotocol/related-task" in payload_result.meta # Test cancel_task (default handler) cancel_result = await client_session.send_request( CancelTaskRequest(params=CancelTaskRequestParams(task_id=task.task_id)), CancelTaskResult ) assert cancel_result.task_id == task.task_id assert cancel_result.status == "cancelled" tg.cancel_scope.cancel() @pytest.mark.anyio async def test_build_elicit_form_request() -> None: """Test that _build_elicit_form_request builds a proper elicitation request.""" server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) try: async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions(server_name="test-server", server_version="1.0.0", capabilities=ServerCapabilities()), ) as server_session: # Test without task_id request = server_session._build_elicit_form_request( message="Test message", requested_schema={"type": "object", "properties": {"answer": {"type": "string"}}}, ) assert request.method == "elicitation/create" assert request.params is not None assert request.params["message"] == "Test message" # Test with related_task_id (adds related-task metadata) request_with_task = server_session._build_elicit_form_request( message="Task message", requested_schema={"type": "object"}, related_task_id="test-task-123", ) assert request_with_task.method == "elicitation/create" assert request_with_task.params is not None assert "_meta" in request_with_task.params assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"] assert ( request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"] == "test-task-123" ) finally: await server_to_client_send.aclose() await server_to_client_receive.aclose() await client_to_server_send.aclose() await client_to_server_receive.aclose() @pytest.mark.anyio async def test_build_elicit_url_request() -> None: """Test that _build_elicit_url_request builds a proper URL mode elicitation request.""" server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) try: async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions(server_name="test-server", server_version="1.0.0", capabilities=ServerCapabilities()), ) as server_session: # Test without related_task_id request = server_session._build_elicit_url_request( message="Please authorize with GitHub", url="https://github.com/login/oauth/authorize", elicitation_id="oauth-123", ) assert request.method == "elicitation/create" assert request.params is not None assert request.params["message"] == "Please authorize with GitHub" assert request.params["url"] == "https://github.com/login/oauth/authorize" assert request.params["elicitationId"] == "oauth-123" assert request.params["mode"] == "url" # Test with related_task_id (adds related-task metadata) request_with_task = server_session._build_elicit_url_request( message="OAuth required", url="https://example.com/oauth", elicitation_id="oauth-456", related_task_id="test-task-789", ) assert request_with_task.method == "elicitation/create" assert request_with_task.params is not None assert "_meta" in request_with_task.params assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"] assert ( request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"] == "test-task-789" ) finally: await server_to_client_send.aclose() await server_to_client_receive.aclose() await client_to_server_send.aclose() await client_to_server_receive.aclose() @pytest.mark.anyio async def test_build_create_message_request() -> None: """Test that _build_create_message_request builds a proper sampling request.""" server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) try: async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=ServerCapabilities(), ), ) as server_session: messages = [ SamplingMessage(role="user", content=TextContent(type="text", text="Hello")), ] # Test without task_id request = server_session._build_create_message_request( messages=messages, max_tokens=100, system_prompt="You are helpful", ) assert request.method == "sampling/createMessage" assert request.params is not None assert request.params["maxTokens"] == 100 # Test with related_task_id (adds related-task metadata) request_with_task = server_session._build_create_message_request( messages=messages, max_tokens=50, related_task_id="sampling-task-456", ) assert request_with_task.method == "sampling/createMessage" assert request_with_task.params is not None assert "_meta" in request_with_task.params assert "io.modelcontextprotocol/related-task" in request_with_task.params["_meta"] assert ( request_with_task.params["_meta"]["io.modelcontextprotocol/related-task"]["taskId"] == "sampling-task-456" ) finally: await server_to_client_send.aclose() await server_to_client_receive.aclose() await client_to_server_send.aclose() await client_to_server_receive.aclose() @pytest.mark.anyio async def test_send_message() -> None: """Test that send_message sends a raw session message.""" server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) try: async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=ServerCapabilities(), ), ) as server_session: # Create a test message notification = JSONRPCNotification(jsonrpc="2.0", method="test/notification") message = SessionMessage( message=notification, metadata=ServerMessageMetadata(related_request_id="test-req-1"), ) # Send the message await server_session.send_message(message) # Verify it was sent to the stream received = await server_to_client_receive.receive() assert isinstance(received.message, JSONRPCNotification) assert received.message.method == "test/notification" finally: # pragma: lax no cover await server_to_client_send.aclose() await server_to_client_receive.aclose() await client_to_server_send.aclose() await client_to_server_receive.aclose() @pytest.mark.anyio async def test_response_routing_success() -> None: """Test that response routing works for success responses.""" server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) # Track routed responses with event for synchronization routed_responses: list[dict[str, Any]] = [] response_received = anyio.Event() class TestRouter(ResponseRouter): def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool: routed_responses.append({"id": request_id, "response": response}) response_received.set() return True # Handled def route_error(self, request_id: str | int, error: ErrorData) -> bool: raise NotImplementedError try: async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=ServerCapabilities(), ), ) as server_session: router = TestRouter() server_session.add_response_router(router) # Simulate receiving a response from client response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"}) message = SessionMessage(message=response) # Send from "client" side await client_to_server_send.send(message) # Wait for response to be routed with anyio.fail_after(5): await response_received.wait() # Verify response was routed assert len(routed_responses) == 1 assert routed_responses[0]["id"] == "test-req-1" assert routed_responses[0]["response"]["status"] == "ok" finally: # pragma: lax no cover await server_to_client_send.aclose() await server_to_client_receive.aclose() await client_to_server_send.aclose() await client_to_server_receive.aclose() @pytest.mark.anyio async def test_response_routing_error() -> None: """Test that error routing works for error responses.""" server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) # Track routed errors with event for synchronization routed_errors: list[dict[str, Any]] = [] error_received = anyio.Event() class TestRouter(ResponseRouter): def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool: raise NotImplementedError def route_error(self, request_id: str | int, error: ErrorData) -> bool: routed_errors.append({"id": request_id, "error": error}) error_received.set() return True # Handled try: async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=ServerCapabilities(), ), ) as server_session: router = TestRouter() server_session.add_response_router(router) # Simulate receiving an error response from client error_data = ErrorData(code=INVALID_REQUEST, message="Test error") error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data) message = SessionMessage(message=error_response) # Send from "client" side await client_to_server_send.send(message) # Wait for error to be routed with anyio.fail_after(5): await error_received.wait() # Verify error was routed assert len(routed_errors) == 1 assert routed_errors[0]["id"] == "test-req-2" assert routed_errors[0]["error"].message == "Test error" finally: # pragma: lax no cover await server_to_client_send.aclose() await server_to_client_receive.aclose() await client_to_server_send.aclose() await client_to_server_receive.aclose() @pytest.mark.anyio async def test_response_routing_skips_non_matching_routers() -> None: """Test that routing continues to next router when first doesn't match.""" server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) # Track which routers were called router_calls: list[str] = [] response_received = anyio.Event() class NonMatchingRouter(ResponseRouter): def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool: router_calls.append("non_matching_response") return False # Doesn't handle it def route_error(self, request_id: str | int, error: ErrorData) -> bool: raise NotImplementedError class MatchingRouter(ResponseRouter): def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool: router_calls.append("matching_response") response_received.set() return True # Handles it def route_error(self, request_id: str | int, error: ErrorData) -> bool: raise NotImplementedError try: async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=ServerCapabilities(), ), ) as server_session: # Add non-matching router first, then matching router server_session.add_response_router(NonMatchingRouter()) server_session.add_response_router(MatchingRouter()) # Send a response - should skip first router and be handled by second response = JSONRPCResponse(jsonrpc="2.0", id="test-req-1", result={"status": "ok"}) message = SessionMessage(message=response) await client_to_server_send.send(message) with anyio.fail_after(5): await response_received.wait() # Verify both routers were called (first returned False, second returned True) assert router_calls == ["non_matching_response", "matching_response"] finally: # pragma: lax no cover await server_to_client_send.aclose() await server_to_client_receive.aclose() await client_to_server_send.aclose() await client_to_server_receive.aclose() @pytest.mark.anyio async def test_error_routing_skips_non_matching_routers() -> None: """Test that error routing continues to next router when first doesn't match.""" server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10) client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10) # Track which routers were called router_calls: list[str] = [] error_received = anyio.Event() class NonMatchingRouter(ResponseRouter): def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool: raise NotImplementedError def route_error(self, request_id: str | int, error: ErrorData) -> bool: router_calls.append("non_matching_error") return False # Doesn't handle it class MatchingRouter(ResponseRouter): def route_response(self, request_id: str | int, response: dict[str, Any]) -> bool: raise NotImplementedError def route_error(self, request_id: str | int, error: ErrorData) -> bool: router_calls.append("matching_error") error_received.set() return True # Handles it try: async with ServerSession( client_to_server_receive, server_to_client_send, InitializationOptions( server_name="test-server", server_version="1.0.0", capabilities=ServerCapabilities(), ), ) as server_session: # Add non-matching router first, then matching router server_session.add_response_router(NonMatchingRouter()) server_session.add_response_router(MatchingRouter()) # Send an error - should skip first router and be handled by second error_data = ErrorData(code=INVALID_REQUEST, message="Test error") error_response = JSONRPCError(jsonrpc="2.0", id="test-req-2", error=error_data) message = SessionMessage(message=error_response) await client_to_server_send.send(message) with anyio.fail_after(5): await error_received.wait() # Verify both routers were called (first returned False, second returned True) assert router_calls == ["non_matching_error", "matching_error"] finally: # pragma: lax no cover await server_to_client_send.aclose() await server_to_client_receive.aclose() await client_to_server_send.aclose() await client_to_server_receive.aclose()