88import json
99import multiprocessing
1010import socket
11+ import threading
1112import time
1213import traceback
1314from collections.abc import AsyncIterator, Generator
@@ -1462,7 +1463,7 @@ async def sampling_callback(
14621463
14631464
14641465# Context-aware server implementation for testing request context propagation
1465- async def _handle_context_list_tools( # pragma: no cover
1466+ async def _handle_context_list_tools(
14661467 ctx: ServerRequestContext, params: PaginatedRequestParams | None
14671468) -> ListToolsResult:
14681469 return ListToolsResult(
@@ -1487,9 +1488,7 @@ async def _handle_context_list_tools( # pragma: no cover
14871488 )
14881489
14891490
1490- async def _handle_context_call_tool( # pragma: no cover
1491- ctx: ServerRequestContext, params: CallToolRequestParams
1492- ) -> CallToolResult:
1491+ async def _handle_context_call_tool(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
14931492 name = params.name
14941493 args = params.arguments or {}
14951494
@@ -1517,8 +1516,8 @@ async def _handle_context_call_tool( # pragma: no cover
15171516
15181517
15191518# Server runner for context-aware testing
1520- def run_context_aware_server (port: int): # pragma: no cover
1521- """Run the context-aware test server."""
1519+ def _create_context_aware_app (port: int) -> tuple[uvicorn.Server, Starlette]:
1520+ """Create the context-aware test server app and uvicorn.Server ."""
15221521 server = Server(
15231522 "ContextAwareServer",
15241523 on_list_tools=_handle_context_list_tools,
@@ -1547,24 +1546,22 @@ def run_context_aware_server(port: int): # pragma: no cover
15471546 log_level="error",
15481547 )
15491548 )
1550- server_instance.run()
1549+ return server_instance, app
15511550
15521551
15531552@pytest.fixture
15541553def context_aware_server(basic_server_port: int) -> Generator[None, None, None]:
1555- """Start the context-aware server in a separate process."""
1556- proc = multiprocessing.Process(target=run_context_aware_server, args=(basic_server_port,), daemon=True)
1557- proc.start()
1554+ """Start the context-aware server on a background thread (in-process for coverage)."""
1555+ server_instance, _ = _create_context_aware_app(basic_server_port)
1556+ thread = threading.Thread(target=server_instance.run, daemon=True)
1557+ thread.start()
15581558
1559- # Wait for server to be running
15601559 wait_for_server(basic_server_port)
15611560
15621561 yield
15631562
1564- proc.kill()
1565- proc.join(timeout=2)
1566- if proc.is_alive(): # pragma: no cover
1567- print("Context-aware server process failed to terminate")
1563+ server_instance.should_exit = True
1564+ thread.join(timeout=5)
15681565
15691566
15701567@pytest.mark.anyio
0 commit comments