test(mcp): cover configured SSE limits in the legacy harness

This commit is contained in:
Joshua Valluru 2026-10-20 16:08:04 -07:00
parent 1e69fb070f
commit 5d70d0375e

View file

@ -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]