From 7b4df9a706e755c091a880ad6182e5b44f07dae1 Mon Sep 17 00:00:00 2001
From: jibanez-staticduo <staticduo@gmail.com>
Date: Thu, 1 Oct 2026 02:38:47 +0200
Subject: [PATCH] test: harden batches mock scope and MCP cancellation deadline
 handshake

- batches _raw_batches_request now provides the ASGI scope real starlette
  requests carry, so the OTLP-trace branch of _read_request_body no longer
  collapses the body and the metadata 400 names the offending field again
- test_outer_deadline_delivers_session_termination gives the initialize/
  tools-call handshake its own 2s budget and keeps the original 0.2s
  fail_after as the real outer cancellation deadline; the in-flight task is
  cancelled and awaited in finally so session-termination DELETE is
  delivered before the caller resumes

Validated: batches 155 passed; test_mcp_client 409 passed; 14 cancel/teardown
variants green under 8-way CPU load.
---
 .../proxy/batches_endpoints/test_endpoints.py | 14 ++++++++++++++
 .../test_mcp_client.py                        | 19 ++++++++++++++-----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
index 2d597abf3b8..8e3d25552e0 100644
--- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
+++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
@@ -1038,6 +1038,20 @@ def _raw_batches_request(body: Dict[str, Any]) -> MagicMock:
     request.headers = {"Content-Type": "application/json"}
     request.client = MagicMock()
     request.client.host = "127.0.0.1"
+    request.scope = {
+        "type": "http",
+        "asgi": {"version": "3.0", "spec_version": "2.3"},
+        "http_version": "1.1",
+        "method": "POST",
+        "scheme": "http",
+        "path": "/v1/batches",
+        "raw_path": b"/v1/batches",
+        "query_string": b"",
+        "root_path": "",
+        "headers": [(b"content-type", b"application/json"), (b"host", b"localhost")],
+        "client": ("127.0.0.1", 54321),
+        "server": ("localhost", 8000),
+    }
     request.body = AsyncMock(return_value=json.dumps(body).encode())
     return request
 
diff --git a/tests/unit/experimental_mcp_client/test_mcp_client.py b/tests/unit/experimental_mcp_client/test_mcp_client.py
index 1a56227b008..3860f0f44bb 100644
--- a/tests/unit/experimental_mcp_client/test_mcp_client.py
+++ b/tests/unit/experimental_mcp_client/test_mcp_client.py
@@ -2644,12 +2644,21 @@ async def test_outer_deadline_delivers_session_termination(termination: str, gro
     client: Final = _MockTransportClient(respond, server_url="https://example.com/mcp", timeout=30)
 
     async def invoke():
-        with anyio.fail_after(0.2):
-            pending: Final = client.call_tool(CallToolRequestParams(name="slow", arguments={}), raise_on_error=raise_on_error)
-            if grouped:
-                await asyncio.gather(pending)
-            else:
+        pending: Final = asyncio.ensure_future(client.call_tool(CallToolRequestParams(name="slow", arguments={}), raise_on_error=raise_on_error))
+        try:
+            with anyio.fail_after(2.0):
+                await started.wait()
+            with anyio.fail_after(0.2):
+                if grouped:
+                    await asyncio.gather(pending)
+                else:
+                    await pending
+        finally:
+            pending.cancel()
+            try:
                 await pending
+            except BaseException:
+                pass
 
     before: Final = anyio.current_time()
     with pytest.raises(TimeoutError):