|
| 1 | +import json |
1 | 2 | import multiprocessing |
2 | 3 | import socket |
3 | 4 | import time |
@@ -318,3 +319,200 @@ async def test_sse_client_basic_connection_mounted_app( |
318 | 319 | # Test ping |
319 | 320 | ping_result = await session.send_ping() |
320 | 321 | assert isinstance(ping_result, EmptyResult) |
| 322 | + |
| 323 | + |
| 324 | +# Test server with request context that returns headers in the response |
| 325 | +class RequestContextServer(Server): |
| 326 | + def __init__(self): |
| 327 | + super().__init__("request_context_server") |
| 328 | + |
| 329 | + @self.call_tool() |
| 330 | + async def handle_call_tool(name: str, args: dict) -> list[TextContent]: |
| 331 | + # Capture request context if available and return it |
| 332 | + headers_info = {} |
| 333 | + try: |
| 334 | + context = self.request_context |
| 335 | + if context.request: |
| 336 | + headers_info = context.request.get("headers", {}) |
| 337 | + except LookupError: |
| 338 | + pass # No request context available |
| 339 | + |
| 340 | + if name == "echo_headers": |
| 341 | + # Return the headers as JSON in the response |
| 342 | + import json |
| 343 | + |
| 344 | + return [TextContent(type="text", text=json.dumps(headers_info))] |
| 345 | + elif name == "echo_context": |
| 346 | + # Return context info with request ID |
| 347 | + import json |
| 348 | + |
| 349 | + context_data = { |
| 350 | + "request_id": args.get("request_id"), |
| 351 | + "headers": headers_info, |
| 352 | + } |
| 353 | + return [TextContent(type="text", text=json.dumps(context_data))] |
| 354 | + |
| 355 | + return [TextContent(type="text", text=f"Called {name}")] |
| 356 | + |
| 357 | + @self.list_tools() |
| 358 | + async def handle_list_tools() -> list[Tool]: |
| 359 | + return [ |
| 360 | + Tool( |
| 361 | + name="echo_headers", |
| 362 | + description="Echoes request headers", |
| 363 | + inputSchema={"type": "object", "properties": {}}, |
| 364 | + ), |
| 365 | + Tool( |
| 366 | + name="echo_context", |
| 367 | + description="Echoes request context", |
| 368 | + inputSchema={ |
| 369 | + "type": "object", |
| 370 | + "properties": {"request_id": {"type": "string"}}, |
| 371 | + "required": ["request_id"], |
| 372 | + }, |
| 373 | + ), |
| 374 | + ] |
| 375 | + |
| 376 | + |
| 377 | +def run_context_server(server_port: int) -> None: |
| 378 | + """Run a server that captures request context""" |
| 379 | + sse = SseServerTransport("/messages/") |
| 380 | + context_server = RequestContextServer() |
| 381 | + |
| 382 | + async def handle_sse(request: Request) -> Response: |
| 383 | + async with sse.connect_sse( |
| 384 | + request.scope, request.receive, request._send |
| 385 | + ) as streams: |
| 386 | + await context_server.run( |
| 387 | + streams[0], streams[1], context_server.create_initialization_options() |
| 388 | + ) |
| 389 | + return Response() |
| 390 | + |
| 391 | + app = Starlette( |
| 392 | + routes=[ |
| 393 | + Route("/sse", endpoint=handle_sse), |
| 394 | + Mount("/messages/", app=sse.handle_post_message), |
| 395 | + ] |
| 396 | + ) |
| 397 | + |
| 398 | + server = uvicorn.Server( |
| 399 | + config=uvicorn.Config( |
| 400 | + app=app, host="127.0.0.1", port=server_port, log_level="error" |
| 401 | + ) |
| 402 | + ) |
| 403 | + print(f"starting context server on {server_port}") |
| 404 | + server.run() |
| 405 | + |
| 406 | + |
| 407 | +@pytest.fixture() |
| 408 | +def context_server(server_port: int) -> Generator[None, None, None]: |
| 409 | + """Fixture that provides a server with request context capture""" |
| 410 | + proc = multiprocessing.Process( |
| 411 | + target=run_context_server, kwargs={"server_port": server_port}, daemon=True |
| 412 | + ) |
| 413 | + print("starting context server process") |
| 414 | + proc.start() |
| 415 | + |
| 416 | + # Wait for server to be running |
| 417 | + max_attempts = 20 |
| 418 | + attempt = 0 |
| 419 | + print("waiting for context server to start") |
| 420 | + while attempt < max_attempts: |
| 421 | + try: |
| 422 | + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 423 | + s.connect(("127.0.0.1", server_port)) |
| 424 | + break |
| 425 | + except ConnectionRefusedError: |
| 426 | + time.sleep(0.1) |
| 427 | + attempt += 1 |
| 428 | + else: |
| 429 | + raise RuntimeError( |
| 430 | + f"Context server failed to start after {max_attempts} attempts" |
| 431 | + ) |
| 432 | + |
| 433 | + yield |
| 434 | + |
| 435 | + print("killing context server") |
| 436 | + proc.kill() |
| 437 | + proc.join(timeout=2) |
| 438 | + if proc.is_alive(): |
| 439 | + print("context server process failed to terminate") |
| 440 | + |
| 441 | + |
| 442 | +@pytest.mark.anyio |
| 443 | +async def test_request_context_propagation( |
| 444 | + context_server: None, server_url: str |
| 445 | +) -> None: |
| 446 | + """Test that request context is properly propagated through SSE transport.""" |
| 447 | + # Test with custom headers |
| 448 | + custom_headers = { |
| 449 | + "Authorization": "Bearer test-token", |
| 450 | + "X-Custom-Header": "test-value", |
| 451 | + "X-Trace-Id": "trace-123", |
| 452 | + } |
| 453 | + |
| 454 | + async with sse_client(server_url + "/sse", headers=custom_headers) as ( |
| 455 | + read_stream, |
| 456 | + write_stream, |
| 457 | + ): |
| 458 | + async with ClientSession(read_stream, write_stream) as session: |
| 459 | + # Initialize the session |
| 460 | + result = await session.initialize() |
| 461 | + assert isinstance(result, InitializeResult) |
| 462 | + |
| 463 | + # Call the tool that echoes headers back |
| 464 | + tool_result = await session.call_tool("echo_headers", {}) |
| 465 | + |
| 466 | + # Parse the JSON response |
| 467 | + |
| 468 | + assert len(tool_result.content) == 1 |
| 469 | + headers_data = json.loads( |
| 470 | + tool_result.content[0].text |
| 471 | + if tool_result.content[0].type == "text" |
| 472 | + else "{}" |
| 473 | + ) |
| 474 | + |
| 475 | + # Verify headers were propagated |
| 476 | + assert headers_data.get("authorization") == "Bearer test-token" |
| 477 | + assert headers_data.get("x-custom-header") == "test-value" |
| 478 | + assert headers_data.get("x-trace-id") == "trace-123" |
| 479 | + |
| 480 | + |
| 481 | +@pytest.mark.anyio |
| 482 | +async def test_request_context_isolation(context_server: None, server_url: str) -> None: |
| 483 | + """Test that request contexts are isolated between different SSE clients.""" |
| 484 | + contexts = [] |
| 485 | + |
| 486 | + # Create multiple clients with different headers |
| 487 | + for i in range(3): |
| 488 | + headers = {"X-Request-Id": f"request-{i}", "X-Custom-Value": f"value-{i}"} |
| 489 | + |
| 490 | + async with sse_client(server_url + "/sse", headers=headers) as ( |
| 491 | + read_stream, |
| 492 | + write_stream, |
| 493 | + ): |
| 494 | + async with ClientSession(read_stream, write_stream) as session: |
| 495 | + await session.initialize() |
| 496 | + |
| 497 | + # Call the tool that echoes context |
| 498 | + tool_result = await session.call_tool( |
| 499 | + "echo_context", {"request_id": f"request-{i}"} |
| 500 | + ) |
| 501 | + |
| 502 | + # Parse and store the result |
| 503 | + import json |
| 504 | + |
| 505 | + assert len(tool_result.content) == 1 |
| 506 | + context_data = json.loads( |
| 507 | + tool_result.content[0].text |
| 508 | + if tool_result.content[0].type == "text" |
| 509 | + else "{}" |
| 510 | + ) |
| 511 | + contexts.append(context_data) |
| 512 | + |
| 513 | + # Verify each request had its own context |
| 514 | + assert len(contexts) == 3 |
| 515 | + for i, ctx in enumerate(contexts): |
| 516 | + assert ctx["request_id"] == f"request-{i}" |
| 517 | + assert ctx["headers"].get("x-request-id") == f"request-{i}" |
| 518 | + assert ctx["headers"].get("x-custom-value") == f"value-{i}" |
0 commit comments