mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
test(mcp): cover configured SSE limits in the legacy harness
This commit is contained in:
parent
1e69fb070f
commit
5d70d0375e
1 changed files with 45 additions and 12 deletions
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue