2 commits

Author	SHA1	Message	Date
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
yucheng-berri	dd511a2b22	fix(logging)!: fire one CustomLogger failure hook per async call outside the async flag list (#47124) ... * fix(logging): fire one CustomLogger hook per async call outside the async flag list aspeech, atranscription, aimage_generation, aimage_edit, aocr and asearch are missing from is_sync_litellm_request's flag list, so the sync success and failure handlers called a CustomLogger's sync hook on top of the async hook the async handler had already called for the same request. The sync handlers now skip a CustomLogger's log_success_event and log_failure_event when that run is paired with an async handler run (for_async_request) and the async handler gives the logger its own async hook. The async client wrapper, dispatch_failure_handlers, the proxy-only failure dispatch and the native bridge's sync failure step for an async call mark their sync failure run as paired. Resolves LIT-9486 * fix(logging): keep the sync CustomLogger hook when no async success runs The async client wrapper skips async success logging for calls carrying a fallbacks kwarg, and the native bridge skips it for internal and fallbacks calls, yet the paired sync success run still dropped a both-hooks logger's sync hook, so that logger got no success event at all. The sync success run now suppresses the sync hook only when the async success handler was actually scheduled for the call: the client wrapper passes async_handler_runs=not is_completion_with_fallbacks, and the native host records per call whether it enqueued or deferred the async success. Tests that pinned the doubled failure hook on native aocr, messages and responses calls and on proxy-rejected requests now expect the async hook alone. The integration test adds /v1/messages and /v1/responses cells and runs the key-auth rejection on every route. * fix(logging): type dispatch_success_handlers kwargs The new async_handler_runs keyword on success_handler made the untyped **kwargs that dispatch_success_handlers forwards an unknown argument for it, which grew the reportUnknownArgumentType count past its merge-base ceiling. Typing the forwarded kwargs as Any, as _success_handler_body already does, removes those diagnostics without changing what is forwarded. * test(integration): cover CustomLogger hook counts across modes, chaos and router strategies Moves the hook count rig into tests/integration/_support/hook_counts.py and extends the cells it drives: - streamed chat, Messages and Responses requests, success and failure - loggers the async handler never reaches (sync-only overrides, sync event lists only) keep one sync hook per request on every surface outside the async flag list, key-auth rejections included - a concurrent burst while the logger's destination stalls and then fails, and a burst while a worker is killed, never give a request two hooks - least-busy keeps counting a request in flight while async speech requests finish around it, and usage-based routing v2 counts an async image request's tokens once * test(integration): Datadog and LangSmith deliver one record per async request Points the built-in Datadog, LangSmith and OpenTelemetry loggers at a local sink and checks every async surface outside the flag list: one Datadog log entry and one LangSmith run per request (two on the merge base), and never two OpenTelemetry spans, which its own emit-once guard already ensured. * fix(logging): give a both-hooks logger its sync hook when the async run never reaches it The sync hook was skipped whenever the logger sat in the async callback list, so a logger the async run never reached got no hook at all: a sibling stalled past the logging worker timeout, a batch poll that skips default logging, or an async failure run cut off. The sync run now holds the hook and the async run hands it back unless it reached the logger * fix(logging): run the held sync success hooks when the proxy drops a blocked response's success A post-call guardrail block drops the parked async success, which left the sync success hook held for it unrun. Base gave a both-hooks logger that hook on these call types, so the drop now ends the pairing and the held hook runs as before * fix(logging): run held sync hooks only after every sync run of the request has closed A held hook released by the async run could run on another thread while the sync run was still walking its other callbacks. The pairing now waits until the async run has ended and every sync run has closed, and the last sync run to close runs the held hooks itself * fix(logging): end the async run when its coroutine is dropped before it starts A logging queue that drops the async success coroutine, or a deferred success that fails to schedule, left the sync hooks held for it unrun. The enqueued async run now ends when its coroutine is discarded, and the proxy ends it when firing the deferred success raises * fix(logging): leave a still-queued async run to the exit flush The discard finalizer ran at interpreter exit before the logging worker's exit flush, so a queued request could get its held sync hook and then its async hook * test: keep the router and file-stream tests in step with main Main gives each owned proxy a private Redis by default, so the router cells share the worker's Redis to read its counters, and the file-stream stub takes the hook pairing the stream now passes * fix(logging): run held sync hooks when a during_call guardrail raises or the exit flush gives up A during_call guardrail that raised after the LLM call parked its deferred success skipped the proxy's flush, and async runs the exit flush never started were never ended, so a both-hooks logger lost its sync hook in both cases * fix(logging): skip the sync failure hook when the async failure run logs to the logger An async call outside is_sync_litellm_request's flag list ran both failure handlers, so a CustomLogger with both failure hooks got both. The sync run of an async request now skips log_failure_event for a logger that overrides async_log_failure_event and is registered for async failures. The four sync runs of async failures (dispatch_failure_handlers, the @client async wrapper, the proxy-only failure thread and the Rust host) pass for_async_request=True. Drops the HookPairing fallback. Successes are left to #47192, which already gives such a logger only async_log_success_event. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>	2026-10-20 15:19:14 -07:00