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