scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

test(responses): audit #47445 (#48902)

Browse source 
* test(responses): audit #47445

Integration cells for billing every round of a streamed MCP auto-execute on one spend row, on a native Responses upstream and the chat bridge, across tool rounds, a rejected or broken follow-up, a post-call guardrail block, a client hang-up, the response cache, and the next bridged turn

* test(responses): check a broken follow-up bills its own nonzero usage

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27 2026-10-28 18:05:08 -07:00 • committed by GitHub
parent 77cad6e903
commit 07d314beb5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 351 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

365
tests/integration/mcp/test_mcp_llm_endpoints.py
View file

@ -6,7 +6,7 @@ from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from hashlib import sha256
from itertools import chain
from itertools import chain, count
from pathlib import Path
from typing import Final, Literal, TypeVar
@ -50,6 +50,7 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat
directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway")
stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = directory / "stored-responses.yaml"
(directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE)
config.write_text(
yaml.safe_dump(
{
@ -57,8 +58,17 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat
"general_settings": {**object_value(stock["general_settings"]), "store_prompts_in_spend_logs": True},
"litellm_settings": {
**object_value(stock["litellm_settings"]),
"cost_margin_config": {"hosted_vllm": {"fixed_amount": FIXED_MARGIN}},
"cost_margin_config": {
"hosted_vllm": {"fixed_amount": FIXED_MARGIN},
"deepseek": {"fixed_amount": FIXED_MARGIN},
},
},
"guardrails": [
{
"guardrail_name": STREAM_BLOCK,
"litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"},
}
],
}
)
)
@ -79,7 +89,7 @@ def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway
yield private
Surface = Literal["chat", "responses", "messages", "messages_bridge"]
Surface = Literal["chat", "responses", "messages", "messages_bridge", "responses_bridge"]
SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge")
ADD: Final = {"a": 2, "b": 3}
ANSWER: Final = "the sum is 5"
@ -88,6 +98,19 @@ AUTO: Final = {**GATEWAY_REF, "require_approval": "never"}
OUTAGE: Final = "bridge-outage"
DIRECT: Final = "answer-without-tools"
FIXED_MARGIN: Final = 0.25
STREAM_BLOCK: Final = "mcp-stream-block"
STREAM_BLOCK_MODULE: Final = "mcp_stream_block"
STREAM_BLOCK_DETAIL: Final = "blocked by output policy"
STREAM_BLOCK_SOURCE: Final = (
"from fastapi import HTTPException\n\n"
"from litellm.integrations.custom_guardrail import CustomGuardrail\n\n\n"
"class BlockGuard(CustomGuardrail):\n"
" async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):\n"
" async for _ in response:\n"
" pass\n"
f' raise HTTPException(status_code=400, detail="{STREAM_BLOCK_DETAIL}")\n'
" yield\n"
)
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
@ -128,6 +151,8 @@ class Billing:
first: tuple[int, int] = (10, 5)
follow_up: tuple[int, int] = (10, 5)
follow_up_fails: bool = False
follow_up_breaks: bool = False
tool_rounds: int = 1
EVEN_ROUNDS: Final = Billing()
@ -138,6 +163,10 @@ def _fixed_turn(tool: str) -> Callable[[Mapping[str, JsonValue]], Turn]:
return lambda _: Turn(tool, json.dumps(ADD), ANSWER)
def _slow_turn(tool: str, seconds: float) -> Callable[[Mapping[str, JsonValue]], Turn]:
return lambda _: Turn(tool, json.dumps({"seconds": seconds}), ANSWER)
def _echoing_turn(body: Mapping[str, JsonValue]) -> Turn:
names: Final = _tool_names(body)
return Turn(names[0] if names else "", json.dumps({"query": _prompt(body)}), _tool_result_text(body) or "")
@ -183,6 +212,50 @@ def _item_events(index: int, item: Mapping[str, JsonValue]) -> tuple[dict[str, J
)
def _event_stream(events: Sequence[Mapping[str, JsonValue]]) -> Reply:
return Reply(
content_type="text/event-stream",
chunks=tuple(
f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode()
for number, event in enumerate(events)
),
)
def _broken_follow_up(target: str, identity: str, tokens: tuple[int, int]) -> Reply:
outage: Final = {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None}
if target.endswith("/chat/completions"):
delta: Final = {"role": "assistant", "content": "the sum"}
started: Final = {
"id": f"chatcmpl-{identity}",
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "delta": delta, "finish_reason": None}],
}
return Reply(
content_type="text/event-stream",
chunks=tuple(f"data: {json.dumps(frame)}\n\n".encode() for frame in (started, {"error": outage})),
)
input_tokens, output_tokens = tokens
failed: Final[dict[str, JsonValue]] = {
"id": f"resp_{identity}",
"object": "response",
"created_at": 1,
"status": "failed",
"model": "gpt-4o-mini",
"output": [],
"error": {"code": "server_error", "message": FOLLOW_UP_OUTAGE},
"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens, "total_tokens": sum(tokens)},
}
return _event_stream(
(
{"type": "response.created", "response": {**failed, "status": "in_progress", "error": None, "usage": None}},
{"type": "response.failed", "response": failed},
)
)
def _responses_stream(response: Mapping[str, JsonValue], items: Sequence[Mapping[str, JsonValue]]) -> Reply:
item_events: Final = tuple(chain.from_iterable(_item_events(index, item) for index, item in enumerate(items)))
events: Final = (
@ -191,13 +264,7 @@ def _responses_stream(response: Mapping[str, JsonValue], items: Sequence[Mapping
*item_events,
{"type": "response.completed", "response": response},
)
return Reply(
content_type="text/event-stream",
chunks=tuple(
f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode()
for number, event in enumerate(events)
),
)
return _event_stream(events)
def _block_ids(message: JsonValue, kind: str, field: str) -> frozenset[str]:
@ -252,6 +319,8 @@ def _anthropic_tool_result_rule(body: Mapping[str, JsonValue]) -> Reply | None:
def _model_double(
plan: Callable[[Mapping[str, JsonValue]], Turn], billing: Billing = EVEN_ROUNDS, parallel: int = 1
) -> Callable[[Request], Reply]:
served: Final = count()
def respond(request: Request) -> Reply:
if request.method == "GET" and request.target.endswith("/models"):
return _json({"object": "list", "data": []})
@ -260,11 +329,12 @@ def _model_double(
outage: Final = {"error": {"message": "scripted provider outage", "type": "server_error", "code": None}}
return Reply(status=500, body=json.dumps(outage).encode())
turn: Final = plan(body)
call_number: Final = next(served)
answered: Final = _has_tool_result(body)
if answered and billing.follow_up_fails:
failure: Final = {"error": {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None}}
return Reply(status=500, body=json.dumps(failure).encode())
done: Final = answered or DIRECT in _prompt(body)
done: Final = DIRECT in _prompt(body) or (answered and call_number >= billing.tool_rounds)
calls: Final = range(1, parallel + 1)
prompt_tokens, completion_tokens = billing.follow_up if answered else billing.first
identity: Final = uuid.uuid4().hex[:12]
@ -273,9 +343,15 @@ def _model_double(
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
if answered and billing.follow_up_breaks:
return _broken_follow_up(request.target, identity, (prompt_tokens, completion_tokens))
if request.target.endswith("/chat/completions"):
requested: Final = tuple(
{"id": f"call_{call}", "type": "function", "function": {"name": turn.tool, "arguments": turn.arguments}}
{
"id": f"call_{identity}_{call}",
"type": "function",
"function": {"name": turn.tool, "arguments": turn.arguments},
}
for call in calls
)
message: Final = (
@ -460,6 +536,7 @@ def _upstream_model(surface: Surface) -> str:
"responses": "openai/responses/gpt-4o-mini",
"messages": "anthropic/claude-sonnet-4-5",
"messages_bridge": "hosted_vllm/gpt-4o-mini",
"responses_bridge": "deepseek/gpt-4o-mini",
}[surface]
@ -475,11 +552,15 @@ def _rig(
billing: Billing = EVEN_ROUNDS,
prices: Mapping[str, JsonValue] | None = None,
parallel: int = 1,
slow_tool_seconds: float | None = None,
) -> Iterator[Rig]:
alias: Final = "llm" + uuid.uuid4().hex[:8]
plan: Final = (
_fixed_turn(f"{alias}-add") if slow_tool_seconds is None else _slow_turn(f"{alias}-slow", slow_tool_seconds)
)
with (
mcp_peer() as peer,
wire_server(_model_double(_fixed_turn(f"{alias}-add"), billing, parallel)) as wire,
mcp_peer(rich=slow_tool_seconds is not None) as peer,
wire_server(_model_double(plan, billing, parallel)) as wire,
gateway.scenario() as scenario,
):
server_id: Final = register_mcp(scenario, peer, alias)
@ -1048,6 +1129,262 @@ def test_bridged_next_turn_replays_the_mcp_turn_answer(gateway: Gateway) -> None
assert ANSWER in _assistant_texts(opening), opening
StreamedUpstream = Literal["messages_bridge", "responses_bridge"]
STREAMED_UPSTREAMS: Final[tuple[StreamedUpstream, ...]] = ("messages_bridge", "responses_bridge")
STREAMED_IDS: Final = ("native", "bridged")
SLOW_TOOL_SECONDS: Final = 5.0
def _events(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]:
return tuple(
JSON_OBJECT.validate_json(line.removeprefix("data: "))
for line in response.text.splitlines()
if line.startswith("data: ") and line != "data: [DONE]"
)
def _completed(response: httpx.Response) -> dict[str, JsonValue]:
finished: Final = tuple(
object_value(event["response"]) for event in _events(response) if event["type"] == "response.completed"
)
assert len(finished) == 1, response.text
return finished[0]
def _model_rows(call_id: str, settled: Callable[[list[dict[str, JsonValue]]], bool]) -> list[dict[str, JsonValue]]:
return eventually(
lambda: read_rows(
"SELECT request_id, status, spend, prompt_tokens, completion_tokens "
'FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type NOT IN (%s, %s)',
(call_id, "call_mcp_tool", "list_mcp_tools"),
),
settled,
seconds=70,
return_last_on_timeout=True,
)
def _billed(rows: Sequence[Mapping[str, JsonValue]]) -> float:
return sum(float(str(row["spend"])) for row in rows)
def _bills_at_least(expected: float) -> Callable[[list[dict[str, JsonValue]]], bool]:
return lambda rows: bool(rows) and _billed(rows) >= expected - 1e-9
def _statuses(rows: Sequence[Mapping[str, JsonValue]]) -> list[str]:
return sorted(str(row["status"]) for row in rows)
def _logs_a_failure_billing(expected: float) -> Callable[[list[dict[str, JsonValue]]], bool]:
return lambda rows: "failure" in _statuses(rows) and _billed(rows) >= expected - 1e-9
def _limited_key(rig: Rig) -> str:
return rig.scenario.key(object_permission={"mcp_servers": [rig.server_id]}, max_parallel_requests=1)
def _answers_directly(rig: Rig, key: str) -> int:
return rig.send(key, [AUTO], prompt=f"{DIRECT} {rig.alias}").status_code
STREAMED_ROUNDS: Final = (
pytest.param("messages_bridge", 1, id="native"),
pytest.param("responses_bridge", 1, id="bridged"),
pytest.param("messages_bridge", 2, id="native-two-tool-rounds"),
pytest.param("responses_bridge", 2, id="bridged-two-tool-rounds"),
pytest.param("messages_bridge", 0, id="native-no-tool-call"),
pytest.param("responses_bridge", 0, id="bridged-no-tool-call"),
)
@pytest.mark.parametrize(("upstream", "tool_rounds"), STREAMED_ROUNDS)
def test_streaming_responses_auto_execute_bills_every_round_once_on_one_success_row(
gateway: Gateway, upstream: StreamedUpstream, tool_rounds: int
) -> None:
rounds: Final = (FIRST_ROUND, *(FOLLOW_UP_ROUND,) * tool_rounds)
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, tool_rounds=max(tool_rounds, 1))
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
prompt: Final = None if tool_rounds else f"{DIRECT} {rig.alias}"
response: Final = rig.send(_granted_key(rig), [AUTO], prompt=prompt, stream=True)
assert response.status_code == 200, response.text
completed: Final = _completed(response)
assert rig.final_text(completed) == ANSWER, response.text
assert len(rig.upstream_tools()) == len(rounds)
assert len(_peer_add_calls(rig.peer)) == tool_rounds
expected: Final = _price(*rounds) + FIXED_MARGIN
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _bills_at_least(expected))
assert [row["status"] for row in rows] == ["success"], rows
row: Final = rows[0]
assert (row["prompt_tokens"], row["completion_tokens"]) == (
sum(prompt_tokens for prompt_tokens, _ in rounds),
sum(completion_tokens for _, completion_tokens in rounds),
), row
usage: Final = object_value(completed["usage"])
assert (usage["input_tokens"], usage["output_tokens"]) == (row["prompt_tokens"], row["completion_tokens"])
assert float(str(row["spend"])) == pytest.approx(expected), row
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
def test_streaming_responses_rejected_follow_up_bills_round_one_once(
gateway: Gateway, upstream: StreamedUpstream
) -> None:
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True)
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
response: Final = rig.send(_granted_key(rig), [AUTO], stream=True)
assert FOLLOW_UP_OUTAGE in response.text, response.text
assert ANSWER not in response.text, response.text
assert len(rig.upstream_tools()) == 2
expected: Final = _price(FIRST_ROUND) + FIXED_MARGIN
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _bills_at_least(expected))
assert _billed(rows) == pytest.approx(expected), rows
BROKEN_FOLLOW_UPS: Final = (
pytest.param("messages_bridge", FOLLOW_UP_ROUND, id="native"),
pytest.param("responses_bridge", None, id="bridged"),
)
@pytest.mark.parametrize(("upstream", "reported_usage"), BROKEN_FOLLOW_UPS)
def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_and_itself_on_a_failure_row(
gateway: Gateway, upstream: StreamedUpstream, reported_usage: tuple[int, int] | None
) -> None:
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True)
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
response: Final = rig.send(_granted_key(rig), [AUTO], stream=True)
assert ANSWER not in response.text, response.text
assert len(rig.upstream_tools()) == 2
round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one))
assert _statuses(rows) == ["failure", "success"], rows
(success,) = (row for row in rows if row["status"] == "success")
(failure,) = (row for row in rows if row["status"] == "failure")
assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows
assert float(str(success["spend"])) == pytest.approx(round_one), rows
failed_round: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"])))
assert min(failed_round) > 0, rows
assert reported_usage in (None, failed_round), rows
assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once(
gateway: Gateway, upstream: StreamedUpstream
) -> None:
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND)
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
key: Final = _granted_key(rig)
one_round: Final = rig.send(key, [], prompt=f"{DIRECT} {rig.alias}", stream=True, guardrails=[STREAM_BLOCK])
assert STREAM_BLOCK_DETAIL in one_round.text, one_round.text
one_round_rows: Final = _model_rows(
one_round.headers["x-litellm-call-id"], _logs_a_failure_billing(_price(FIRST_ROUND) + FIXED_MARGIN)
)
rig.wire.drain()
response: Final = rig.send(key, [AUTO], stream=True, guardrails=[STREAM_BLOCK])
assert STREAM_BLOCK_DETAIL in response.text, response.text
assert len(rig.upstream_tools()) == 2 and len(_peer_add_calls(rig.peer)) == 1
expected: Final = _price(FIRST_ROUND, FOLLOW_UP_ROUND) + FIXED_MARGIN
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(expected))
assert _statuses(rows) == _statuses(one_round_rows), (rows, one_round_rows)
assert _billed(rows) == pytest.approx(expected), rows
@pytest.mark.parametrize(
("surface", "stream"),
(("responses", True), ("responses", False), ("chat", False)),
ids=("responses-stream", "responses", "chat"),
)
def test_auto_executed_tool_call_runs_inside_the_requests_one_parallel_slot(
gateway: Gateway, surface: Surface, stream: bool
) -> None:
with _rig(gateway, surface, slow_tool_seconds=SLOW_TOOL_SECONDS) as rig:
key: Final = _limited_key(rig)
with ThreadPoolExecutor(max_workers=1) as pool:
running: Final = pool.submit(rig.send, key, [AUTO], None, stream=stream)
eventually(lambda: _peer_add_calls(rig.peer), lambda found: len(found) == 1, seconds=30)
during: Final = _answers_directly(rig, key)
served: Final = running.result(timeout=90)
assert during == 429
assert served.status_code == 200 and ANSWER in served.text, served.text
eventually(lambda: _answers_directly(rig, key), lambda status: status == 200, seconds=30)
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
def test_streaming_responses_hung_up_during_the_tool_call_bills_each_served_round_once_and_frees_the_slot(
gateway: Gateway, upstream: StreamedUpstream
) -> None:
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND)
with _rig(
gateway,
"responses",
upstream=upstream,
billing=billing,
prices=PRICES,
slow_tool_seconds=SLOW_TOOL_SECONDS,
) as rig:
key: Final = _limited_key(rig)
with rig.gateway.client.stream(
"POST",
"/v1/responses",
headers={"Authorization": f"Bearer {key}"},
json={"model": rig.model, "input": f"add {rig.alias}", "tools": [AUTO], "stream": True},
timeout=90,
) as hung_up:
assert hung_up.status_code == 200
call_id: Final = hung_up.headers["x-litellm-call-id"]
eventually(lambda: _peer_add_calls(rig.peer), lambda found: len(found) == 1, seconds=30)
eventually(lambda: _answers_directly(rig, key), lambda status: status == 200, seconds=60)
posted: Final = tuple(
JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST"
)
served: Final = tuple(
FOLLOW_UP_ROUND if _has_tool_result(body) else FIRST_ROUND for body in posted if DIRECT not in _prompt(body)
)
expected: Final = _price(*served) + FIXED_MARGIN
rows: Final = _model_rows(call_id, _bills_at_least(expected))
assert served and rows, (served, rows)
assert _billed(rows) == pytest.approx(expected), (served, rows)
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
def test_streaming_responses_repeated_from_the_response_cache_bills_nothing(
gateway: Gateway, upstream: StreamedUpstream
) -> None:
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND)
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
key: Final = _granted_key(rig)
paid: Final = rig.send(key, [AUTO], stream=True)
assert paid.status_code == 200 and rig.final_text(_completed(paid)) == ANSWER, paid.text
paid_price: Final = _price(FIRST_ROUND, FOLLOW_UP_ROUND) + FIXED_MARGIN
_model_rows(paid.headers["x-litellm-call-id"], _bills_at_least(paid_price))
assert len(rig.upstream_tools()) == 2
cached: Final = rig.send(key, [AUTO], stream=True)
assert cached.status_code == 200 and rig.final_text(_completed(cached)) == ANSWER, cached.text
assert rig.upstream_tools() == ()
rows: Final = _model_rows(
cached.headers["x-litellm-call-id"], lambda found: "success" in {row["status"] for row in found}
)
assert "success" in {row["status"] for row in rows} and _billed(rows) == 0.0, rows
def test_bridged_next_turn_after_a_streamed_mcp_turn_replays_its_answer(gateway: Gateway) -> None:
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND)
with _rig(gateway, "responses", upstream="responses_bridge", billing=billing, prices=PRICES) as rig:
key: Final = _granted_key(rig)
first: Final = rig.send(key, [AUTO], stream=True)
assert first.status_code == 200, first.text
first_id: Final = string_value(_completed(first)["id"])
_model_rows(first.headers["x-litellm-call-id"], _bills_at_least(_price(FIRST_ROUND, FOLLOW_UP_ROUND)))
rig.wire.drain()
second: Final = rig.send(key, [AUTO], previous_response_id=first_id, cache={"no-cache": True})
assert second.status_code == 200, second.text
opening: Final = next(
JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST"
)
assert ANSWER in _assistant_texts(opening), opening
Bridge = Literal["chat", "responses", "messages"]
BRIDGES: Final[tuple[Bridge, ...]] = ("chat", "responses")
Client = Literal["sync", "async"]