litellm/tests/integration/_support
simonp9 160a0b72cf
fix(mcp): bill both calls of a streamed chat completions MCP auto-execute (#48899)
* fix(mcp): bill both rounds of a streamed chat completions MCP auto-execute

A streamed /v1/chat/completions request with an MCP tool set to require_approval "never" runs two model calls on one logging object, and the final-stream dedupe kept only the first round's spend row. The final usage chunk also showed only the second round.

The MCP stream iterator now holds the first round's assembled response, prices it, and logs one success row that carries the usage and cost of both rounds, with the fixed margin charged once. The final usage chunk shows the same summed usage. A single-round stream still logs its one row when the answer finishes.

When the follow-up round fails, the first round is logged from a copy of the logging object taken before the follow-up opened, and the failure row carries only the follow-up's partial usage with no second fixed margin. With proxy post-call guardrails armed, only (assembled, cache_hit) tuples are parked for the proxy, and the deferred log is armed only while the iterator reads, so the proxy's orphaned-callback cleanup never sees it after the Router's first-chunk read

* fix(mcp): send the usage chunk of a streamed chat MCP request that ends in one round

When a streamed chat completions MCP request ends in its first round (the model answers without a tool call, or no tool results come back), the iterator drained the inner stream to fire its spend log and dropped the usage chunk that stream produced, so a client that asked for include_usage got no usage chunk at all.

The drained usage chunk is now held and sent to the client after the final content chunk and before [DONE], the same place a plain stream sends it, so the last usage chunk matches the billed row. A two-round request still drops the first round's usage chunk and sends the merged one from the follow-up, and a stream without include_usage still sends none

* fix(mcp): hold round one's log for the guardrail, bill the chunk cost once, and keep round one's id

When the follow-up call of a streamed chat MCP auto-execute fails, round one's
success log now goes through the proxy's deferred guardrail release instead of
straight to the handlers, so a post-call guardrail runs on it and the row
carries it. Without a proxy release it logs directly as before

The last usage chunk now takes its cost from the combined row's billed total,
so a fixed cost margin is counted once in the chunk, as it is on the row. The
follow-up's usage chunk is held until the follow-up ends to make that possible

The combined row now keeps round one's id, the stream's first chunk id, which
matches how a single-round stream is logged

* fix(mcp): bill every chat MCP stream round once on failure, hang-up, and sync drop

When the follow-up call fails, round 1's success row is now logged on a copy of the
logging object that does not share the request's guardrail results, so a post-call
guardrail that reports on the failed request counts once instead of on both rows

A client hang-up during the follow-up now bills round 1 plus the follow-up's partial
answer on one success row with round 1's id and the fixed margin counted once, the
same way a plain stream's hang-up bills its partial answer

The last usage chunk now carries round 1's usage and cost when only round 1 is billed,
and a synchronous reader that drops the stream mid-answer bills round 1 on the loop it
read with instead of discarding it

The deferred release's logged_on argument is typed and every release closure accepts
it, and the parked dispatch constants and protocols both MCP billing paths use now live
in one module. The unit tests wait on a logging signal instead of a fixed delay and
check the billed row through the proxy stream hook instead of private callback state

* test(mcp): drive the unassembled follow-up through an upstream audio delta that is not base64

The test replaced litellm.stream_chunk_builder with a stub that raises. Round 2 now streams an audio delta whose data is not valid base64, so the real builder raises, the stream parks only best-effort usage, and the client's last usage chunk must carry round 1's tokens and cost

* fix(mcp): count the fixed margin once in the combined stream row's usage cost

* fix(mcp): bill round 2's partial answer when a client hangs up on a routed stream

On the proxy the Router wraps the MCP stream, and when a client hangs up during
round 2 the router's own cleanup closes that stream before the proxy's disconnect
cleanup runs. The close logged round 1 alone, so round 2's partial answer was
never billed. The MCP stream now bills round 1 plus round 2's partial answer as
soon as the hang-up's cancellation reaches its read, before anything closes it

* fix(mcp): count the follow-up's own messages in round 2's partial prompt

The follow-up stream shares the request's logger, so it took the client's
original messages as its prompt. When round 2 ended early, from a hang-up or a
broken upstream, its partial usage counted only those messages and left out the
assistant tool call and the tool results the follow-up actually sent. The
follow-up stream now counts the messages the follow-up call sent

* fix(mcp): bill round 2's partial when a routed stream's client hangs up between reads

On a proxy the stream the disconnect cleanup holds is the router's fallback wrapper, so the
cleanup never found the MCP stream's in-flight billing and released round 1 alone when the client
hung up while the proxy sat at a yield. The round billing now registers its in-flight biller on
the request's logging object, next to the parked release, and the cleanup looks it up there, so
any wrapper chain bills round 1 plus round 2's partial once

* fix(mcp): send the MCP stream usage chunk under the billed row id and count the partial off the loop

A streamed chat completion that runs a LiteLLM-hosted MCP tool bills both model calls on one row under round 1's
chatcmpl id, but the usage chunk it sent the client was round 2's held chunk and kept round 2's id, so a caller could
not join the usage chunk to its row. The usage chunk now carries round 1's id in every path that sends it, with or
without a deferred post-call guardrail, while its usage and cost and every content chunk's id stay as they were

When the client hangs up during round 2, assembling round 2's partial answer counts the follow-up messages, tool
result included, with the local tokenizer. That ran on the event loop and stalled every other request on the worker
for about a second on a large tool result. The assembly now runs on a worker thread through asyncify from both the
reader's cancellation and the proxy's disconnect cleanup, shielded from cancellation, and round 1 is checked again
after the await so a release that happened meanwhile is never billed twice

* test(mcp): wait on the loop with a hang guard instead of a 300 ms deadline

* fix(mcp): give a proxy-rejected follow-up's failure row round 1's own cost breakdown

* refactor(streaming): fold the partial response's empty guards into one check

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-10-30 07:34:10 -07:00
..
__init__.py test: add CircleCI integration contract foundation 2026-09-14 03:30:52 -07:00
actors.py test(integration): give each owned proxy a private database and Redis by default (#47213) 2026-10-18 03:37:41 +00:00
agentic_probe.py fix(responses): honor caller stream flag when provider forces SSE (internal copy of #34095) (#41235) 2026-10-06 03:30:47 +00:00
anthropic_refusal.py fix(anthropic): keep refusal stop_reason and stop_details on the logged response (#46104) 2026-10-14 14:42:49 -07:00
anthropic_sse.py fix(router): retry a /v1/messages stream the provider drops before the first content chunk (#44276) 2026-10-05 22:53:39 +00:00
anthropic_thinking.py test: expect the list-position index on streamed thinking blocks in the integration suite (#46618) 2026-10-15 05:52:56 +00:00
asgi.py fix(ci): isolate integration workers and repair test fixtures (#45863) 2026-10-11 01:56:02 -07:00
async_server.py fix(bedrock): forward Nova Sonic realtime tool calls with their arguments (#46214) 2026-10-16 01:10:53 +00:00
aws_control_plane.py fix(bedrock): list the account's invocable models behind bedrock/* when check_provider_endpoint is on (#43981) 2026-10-09 14:14:17 -07:00
bedrock_discovery.py test(integration): give each owned proxy a private database and Redis by default (#47213) 2026-10-18 03:37:41 +00:00
bedrock_realtime_peer.py fix(bedrock)!: block realtime replay after an input-first start and reject malformed Nova Sonic session settings (#46936) 2026-10-17 00:29:50 +00:00
bedrock_runtime_peer.py test(ci): fix six CircleCI test regressions on main (#44429) 2026-10-03 20:55:01 -07:00
browser_state.py test: add extension and browser integration contracts 2026-09-16 13:15:53 -07:00
budget.py test(integration): give chaos and timing tests a harness API and gate raw signals (#46031) 2026-10-12 05:27:33 -07:00
claude_code.py test(anthropic): native /v1/messages reasoning integration tests built on a captured Claude Code request (#43361) 2026-10-01 09:12:53 -07:00
client.py fix(caching): read the cache key's model group from the router's metadata for the call (#48917) 2026-10-29 05:35:23 -07:00
codex_vendor.py fix(responses): honor caller stream flag when provider forces SSE (internal copy of #34095) (#41235) 2026-10-06 03:30:47 +00:00
copilot_session_prometheus_worker.py fix(prometheus): count provider-side 401s in LLM failure and deployment metrics (#46668) 2026-10-26 17:57:06 -07:00
daily_activity.py fix(proxy): look up hashed key names with two spend log rows per key (#43656) 2026-09-29 20:21:55 -07:00
daily_spend_rows.py fix(ui): group cost optimization cache leakage by model group (#43008) 2026-09-24 23:51:51 -07:00
database.py test(integration): isolate every multi-worker integration suite by default (#46027) 2026-10-12 03:46:06 -07:00
database_relay.py test(integration): hand every harness server a port from one per-worker block (#46761) 2026-10-15 16:56:40 -07:00
deepseek_tokens.py fix(token_counter): count DeepSeek V4 text with DeepSeek's own tokenizer (#47353) 2026-10-20 03:43:35 +00:00
device_login.py fix(github_copilot): treat a malformed api-key.json like a missing one (#44723) 2026-10-29 03:46:44 -07:00
forward_proxy.py feat(strands_decider): call a Strands Decider on an AgentCore runtime ARN (#45678) 2026-10-09 22:51:19 -07:00
generation.py fix(tests): resolve the integration support package without run.py's PYTHONPATH 2026-09-15 21:10:08 -07:00
hook_count_probe.py fix(mcp): bill both calls of a streamed chat completions MCP auto-execute (#48899) 2026-10-30 07:34:10 -07:00
hook_counts.py fix(logging)!: fire one CustomLogger failure hook per async call outside the async flag list (#47124) 2026-10-20 15:19:14 -07:00
idle_gateway.py fix(http): restore the 15s aiohttp keep-alive default in the Python and Rust clients (#46355) 2026-10-14 10:50:11 -07:00
interception_vendor.py fix(websearch_interception): drop stream_options when forcing stream=False and replay it on the final stream (#45433) 2026-10-09 22:16:02 -07:00
mail.py feat(proxy): email alerts at configured percentages of a team member budget (#42665) 2026-09-26 02:42:45 +00:00
manifest.py fix(caching): never store or serve a chat completion with no choices (#44709) 2026-10-06 20:07:30 +00:00
mcp.py fix(mcp): upgrade to MCP SDK 2.3.0 and raise the per-event SSE size cap (#47034) 2026-10-20 16:00:00 -07:00
mcp_e2e_upstream_server.py test: migrate proxy security, schema migration and MCP legacy tests (#45937) 2026-10-11 16:17:24 -07:00
mcp_grants.py fix(mcp): resolve team-granted toolsets for non-admin keys and dashboard sessions (#43908) 2026-10-01 14:15:06 -07:00
mcp_stdio_peer.py test(integration): add MCP gateway coverage wave 1 with a dedicated mcp shard and proxy coverage artifact (#42711) 2026-09-23 07:48:46 -07:00
oauth_server.py test(integration): credential canary slots for MCP and pass-through credentials (#43308) 2026-09-28 17:10:51 -07:00
openai_wire.py test: fix shared-provider discovery, Codex catalog size and generated master key mismatches in CircleCI suites (#44905) 2026-10-07 00:11:25 +00:00
ordering.py fix(ci): isolate integration workers and repair test fixtures (#45863) 2026-10-11 01:56:02 -07:00
otlp_sink.py test(integration): hand every harness server a port from one per-worker block (#46761) 2026-10-15 16:56:40 -07:00
pdf_document.py fix(token_counter): price a base64 PDF document per page instead of as one image (#45301) 2026-10-08 21:20:46 -07:00
ports.py test(integration): hand every harness server a port from one per-worker block (#46761) 2026-10-15 16:56:40 -07:00
prisma_clients.py test: migrate tests/proxy_behavior into integration and unit tiers (#45939) 2026-10-11 08:28:42 -07:00
process.py fix(proxy): keep every startup line JSON when json_logs is on (#47035) 2026-10-27 14:49:07 -07:00
prometheus_series.py test(integration): give each owned proxy a private database and Redis by default (#47213) 2026-10-18 03:37:41 +00:00
prompt_cache_breakpoint.py fix(responses): map input_audio blocks in the chat-to-Responses bridge (#45224) 2026-10-07 22:25:35 -07:00
provider.py fix(ci): isolate integration workers and repair test fixtures (#45863) 2026-10-11 01:56:02 -07:00
proxy.py perf(ci): isolate and shard integration suites (#44146) 2026-10-10 11:15:19 -07:00
redaction_probe.py fix(logging): deliver completed batch success callbacks consistently (#47192) 2026-10-18 09:19:04 -07:00
redis_breaker_worker.py fix(caching): report Redis circuit-breaker metrics through the Prometheus services logger (#46484) 2026-10-14 14:29:54 -07:00
redis_process.py test(integration): hand every harness server a port from one per-worker block (#46761) 2026-10-15 16:56:40 -07:00
responses_replay.py fix(responses): replay the original Responses item order, phase, and boundaries on the next chat turn (#48912) 2026-10-29 00:56:02 -07:00
responses_stream.py fix(responses): keep the provider code and message on top-level stream error events (#46438) 2026-10-14 16:16:46 -07:00
responses_vendor.py fix(responses): drop bridge-minted reasoning items from OpenAI replays (#44132) 2026-10-02 19:39:48 -07:00
responses_ws_peer.py fix(responses): write one spend-log row per native WebSocket turn (#46768) 2026-10-15 19:05:22 -07:00
riva.py test(nvidia_riva): audit #41021 (#48801) 2026-10-28 05:24:08 -07:00
routing.py test(integration): isolate every multi-worker integration suite by default (#46027) 2026-10-12 03:46:06 -07:00
runtime.py test(integration): wait for every worker to report the websocket session cap instead of sleeping (#47412) 2026-10-19 15:35:29 -07:00
safeguards_rig.py test(integration): give each owned proxy a private database and Redis by default (#47213) 2026-10-18 03:37:41 +00:00
sagemaker_embedding.py feat(sagemaker)!: inference components and OpenAI-compatible embeddings (#44985) 2026-10-12 22:25:52 +00:00
scim.py fix(scim)!: delete a group through the team delete's bulk path instead of one member at a time (#47544) 2026-10-21 00:59:23 +00:00
sdk_models.py test(integration): build OpenAI and Anthropic SDK models before tests start threads (#46797) 2026-10-15 18:28:32 -07:00
sharding.py perf(ci): use native timing splits with bounded integration concurrency (#45837) 2026-10-10 13:35:11 -07:00
sigv4.py fix(bedrock): list the account's invocable models behind bedrock/* when check_provider_endpoint is on (#43981) 2026-10-09 14:14:17 -07:00
timing.py test(integration): give chaos and timing tests a harness API and gate raw signals (#46031) 2026-10-12 05:27:33 -07:00
tls.py fix(azure_storage): name Data Lake objects without base64 padding or slashes (#43914) 2026-09-30 18:17:13 -07:00
tool_rows.py feat(tool-policies): show the user who owns the key that discovered a tool (#43892) 2026-10-01 20:34:57 +00:00
tool_search.py fix(responses bridge): round-trip tool_search_call and tool_search_output items between turns (#46444) 2026-10-14 21:40:14 -07:00
upstream.py test(integration): stop GET /__observations from draining the upstream log (#47256) 2026-10-18 09:58:05 +00:00
vertex.py fix(vertex_ai): forward system and tools to partner model count_tokens (#43900) 2026-10-03 16:36:14 -07:00
wire.py fix(ollama): forward embedding credentials, headers and timeout (#48707) 2026-10-28 09:32:47 -07:00
worker_readiness.py test(integration): cover responses endpoint contracts (#44733) 2026-10-14 06:59:15 +00:00
workers.py test(integration): give each owned proxy a private database and Redis by default (#47213) 2026-10-18 03:37:41 +00:00