From 5d70d0375ea681fc7be57ecef50cd69997d7ef4d Mon Sep 17 00:00:00 2001
From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Date: Tue, 20 Oct 2026 16:08:04 -0700
Subject: [PATCH] test(mcp): cover configured SSE limits in the legacy harness

---
 .../mcp/test_mcp_protocol_errors.py           | 57 +++++++++++++++----
 1 file changed, 45 insertions(+), 12 deletions(-)

diff --git a/tests/integration/mcp/test_mcp_protocol_errors.py b/tests/integration/mcp/test_mcp_protocol_errors.py
index 35aa0f0709e..8d29034bb3d 100644
--- a/tests/integration/mcp/test_mcp_protocol_errors.py
+++ b/tests/integration/mcp/test_mcp_protocol_errors.py
@@ -2,6 +2,11 @@ import json
 import queue
 import uuid
 from typing import Final
+from pathlib import Path
+
+import httpx
+import yaml
+from integration._support.process import owned_proxy
 
 import pytest
 from integration._support.client import Gateway
@@ -87,7 +92,8 @@ def test_jsonrpc_error_and_malformed_tool_result_remain_errors(gateway: Gateway)
 
 
 @pytest.mark.covers("other.mcp.transport.large_sse_tool_result_reaches_caller")
-def test_tool_result_larger_than_one_mib_in_a_single_sse_event_reaches_the_caller(gateway: Gateway) -> None:
+@pytest.mark.parametrize("limit", (64 * 1024 * 1024, 4096))
+def test_tool_result_larger_than_one_mib_in_a_single_sse_event_reaches_the_caller(tmp_path: Path, limit: int) -> None:
     payload: Final = "x" * (3 * 1024 * 1024)
     request_json: Final = TypeAdapter(dict[str, JsonValue])
 
@@ -126,22 +132,49 @@ def test_tool_result_larger_than_one_mib_in_a_single_sse_event_reaches_the_calle
                 ).encode()
             )
         assert method == "tools/call"
+        params: Final = request_json.validate_python(body["params"])
+        arguments: Final = request_json.validate_python(params["arguments"])
+        size: Final = TypeAdapter(int).validate_python(arguments["size"])
         response: Final = json.dumps(
             {
                 "jsonrpc": "2.0",
                 "id": body["id"],
-                "result": {"content": [{"type": "text", "text": payload}], "isError": False},
+                "result": {"content": [{"type": "text", "text": payload[:size]}], "isError": False},
             }
         )
         return Reply(content_type="text/event-stream", body=f"event: message\ndata: {response}\n\n".encode())
 
-    with wire_server(provider) as wire, gateway.scenario() as scenario:
-        peer: Final = McpPeer(wire.url + "/mcp", queue.Queue())
-        identity: Final = register_mcp(scenario, peer, "bigsse" + uuid.uuid4().hex[:8])
-        key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
-        names: Final = tool_names(gateway, key, identity)
-        response: Final = call_tool(gateway, key, identity, names["add"], {})
-        assert response.status_code == 200, response.text[:500]
-        body: Final = response.json()
-        assert body["isError"] is False, response.text[:500]
-        assert body["content"][0]["text"] == payload
+    key: Final = "sk-local-mcp-sse-limit-regression-key"
+    with wire_server(provider) as wire, httpx.Client() as client:
+        config: Final = tmp_path / "sse-limit.yaml"
+        config.write_text(
+            yaml.safe_dump(
+                {
+                    "general_settings": {"master_key": key, "store_model_in_db": False},
+                    "mcp_servers": {"big": {"url": wire.url + "/mcp", "transport": "http"}},
+                }
+            )
+        )
+        seed: Final = Gateway(client, key, wire.url)
+        with owned_proxy(
+            seed,
+            tmp_path / "proxy",
+            {
+                "STORE_MODEL_IN_DB": "False",
+                "DISABLE_SCHEMA_UPDATE": "true",
+                "LITELLM_LOCAL_MODEL_COST_MAP": "True",
+                "LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES": str(limit),
+            },
+            config=config,
+            remove_environment=("DATABASE_URL", "DATABASE_URL_READ_REPLICA", "DIRECT_URL"),
+        ) as gateway:
+            names: Final = tool_names(gateway, key, "big")
+            for size in (512, len(payload), 512):
+                response: Final = call_tool(gateway, key, "big", names["add"], {"size": size})
+                assert response.status_code == 200, response.text[:500]
+                body: Final = response.json()
+                assert body["isError"] is (size > limit), response.text[:500]
+                if size > limit:
+                    assert str(limit) in body["content"][0]["text"], response.text[:500]
+                else:
+                    assert body["content"][0]["text"] == payload[:size]