Repository navigation
test: fix batches metadata 400 and MCP cancellation deadline flake - #1
Merged
jibanez-staticduo merged 1 commit intoOct 1, 2026
Conversation
- batches _raw_batches_request now provides the ASGI scope real starlette requests carry, so the OTLP-trace branch of _read_request_body no longer collapses the body and the metadata 400 names the offending field again - test_outer_deadline_delivers_session_termination gives the initialize/ tools-call handshake its own 2s budget and keeps the original 0.2s fail_after as the real outer cancellation deadline; the in-flight task is cancelled and awaited in finally so session-termination DELETE is delivered before the caller resumes Validated: batches 155 passed; test_mcp_client 409 passed; 14 cancel/teardown variants green under 8-way CPU load.
Owner
Author
|
Incorporated into BerriAI#41050 via final tip 0b630f4 (fast-forward push to litellm_forward_reasoning_content). This PR is MERGED; batches ASGI-scope mock and MCP cancellation-handshake fixes are on the published tip. No further work here. |
jibanez-staticduo
pushed a commit
that referenced
this pull request
Oct 8, 2026
…adata (BerriAI#43688) * fix(proxy): traceparent/baggage fallback must not override caller metadata The W3C traceparent/baggage fallback in add_litellm_metadata_from_request_headers documents itself as last-resort: Lower priority than everything above - only fires when neither the explicit litellm headers nor the Anthropic-metadata path found anything But it guards on the top-level `litellm_trace_id` / `litellm_session_id` body keys and never checks `metadata`, which is the documented way callers set trace_id / session_id (`metadata: {"trace_id": ...}` on /chat/completions). A caller that explicitly sets metadata.trace_id has it silently replaced by the header value, so the implementation contradicts its own stated precedence. This is not a corner case on managed platforms: GCP's front end injects a traceparent into every inbound request, so the fallback fires on traffic whose caller never sent the header. The request still returns 200 and the trace still reaches the logging backend, just under an id the caller never chose, so any caller correlating by its own id silently fails to find its trace. Guard both fallbacks on the caller's request-body metadata as well. Deliberately narrow: - x-litellm-trace-id still outranks the body (documented priority #1) - a request that steers neither field still adopts traceparent/baggage exactly as before - steering is per-field: setting only trace_id still lets session_id come from baggage - litellm_metadata is checked too, for LITELLM_METADATA_ROUTES (/v1/responses, /v1/messages, batches, files) Also corrects the comment, which understated the guard. 4 new tests; each fails without the source change. The existing traceparent and baggage tests are unchanged and still pass. * fix(proxy): check only the active metadata container and ignore empty values Addresses review. The first version treated any value in either `metadata` or `litellm_metadata` as caller steering. On LITELLM_METADATA_ROUTES the body `metadata` is provider-facing and is only promoted later, so a session_id there suppressed baggage before apply_missing_session_id_policy ran, and a request with a usable baggage session id got a 400 under `missing_session_id: reject`. An empty session_id did the same, since the policy treats "" as absent Now only the active metadata container counts, and only a truthy value, which matches how apply_missing_session_id_policy decides a session id is present. The helper takes a typed `object` rather than a bare dict Adds regression tests for both cases, including the end to end reject path, and drops the test docstrings * fix(proxy): count promoted caller trace ids on litellm_metadata routes Addresses review. On LITELLM_METADATA_ROUTES the proxy promotes the caller's trace control fields (trace_id, session_id, ...) from `metadata` into `litellm_metadata`, but only after the header fallback runs. Checking only `litellm_metadata` let traceparent and baggage fill those fields first, and the promotion then skipped them because they were already set, so the trace was recorded under the header ids The check now also reads `metadata` for fields in LITELLM_TRACE_CONTROL_METADATA_FIELDS on those routes, and apply_missing_session_id_policy uses the same check, so `reject` no longer refuses a request whose session id is about to be promoted The earlier test asserting that a session_id in `metadata` must not block baggage on /v1/responses had the premise wrong, since that value is promoted. It is replaced by tests that assert the promoted caller ids win on /v1/responses and /v1/messages, with no policy, reject, and generate * fix(proxy): an empty trace field in litellm_metadata shadows the promoted one Addresses review. Promotion copies a trace control field from `metadata` into `litellm_metadata` only when the key is absent, so a key that is present but empty in `litellm_metadata` wins and the `metadata` value is never used. The check still looked at `metadata` in that case, which let `missing_session_id: reject` accept a request that ends up with an empty session id, and let an empty trace_id or session_id block the headers The check now follows the same rule as promotion. If the active container has the key, its value decides. Only when the key is absent does the promoted `metadata` value count. This makes the conflicting-bucket cases behave exactly as on main * fix(proxy): stay within the LIT006 cast budget The type-discipline gate failed: this branch added 4 unsuppressed `cast()` calls and LIT006 was already at its limit. Each cast now carries a `# cast-ok` reason. The two in the helper follow an `isinstance` check that proves the Mapping. The two at the call sites stay because the method's `data` parameter is a bare `dict`, and dropping those casts adds basedpyright unknown-type errors instead. No behavior change * test(proxy): cover caller metadata precedence over W3C trace headers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): generated session uses promoted caller trace id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): move test context into assertion messages Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): audit cells for w3c fallback precedence Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tighten audit chaos cells Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(test): drop stray blank lines Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): ignore body-less model-info probes in langfuse precedence upstreams Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep root trace/session fields on equal ids and usable caller values The _caller_trace_field gating added for caller-metadata precedence was presence-aware, which regressed three pre-call behaviors versus main: - Equal ids in W3C headers and body metadata suppressed the traceparent/baggage branch entirely, leaving the root litellm_trace_id/litellm_session_id unset so router fallbacks minted a fresh uuid4 per request. The header branch now also fires when the caller value equals the header-derived id, so equal ids stamp the root fields exactly like main. - A truthy but non-string metadata value (e.g. session_id 4815162342) counted as caller-supplied and suppressed the baggage fallback, so downstream str-only consumers (code interpreter sandbox reuse) got a new sandbox per turn. _caller_trace_field now counts only non-empty string values, and the generate policy falls through to generation when the caller value is not usable. - On litellm_metadata routes a usable caller session id suppressed the missing_session_id policies while never landing on the root field, so the root session stayed unset. The policy now promotes the caller's usable session id to litellm_session_id instead of leaving it stranded. * test(integration): non-string caller ids fall back to W3C like empty ids The audit cell compared a headers leg against a headerless leg for equality, which only held while a truthy non-string id suppressed the W3C fallback on both legs. With unusable values ignored again, the headers leg resolves to the header ids while the headerless leg cannot, so assert the headers leg's concrete outcome instead, mirroring the empty/null cell. * style(proxy): trim the session promotion comment to the non-obvious why --------- Co-authored-by: Filipe Andujar <filipeandujar@gmail.com> Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Two CI-only test failures on the PR41050 branch (BerriAI#41050):
test_create__non_object_metadata_is_400[metadata|litellm_metadata]returned a bogus 400Missing required parameter: 'input_file_id'instead of the 400 naming the bad metadata field. The_raw_batches_requesthelper mocksRequestwith a spec'd MagicMock that lacks the ASGIscope; after upstream feat(tracing): port OTLP ingestion to current trace foundation BerriAI/litellm#43915 the body reader inspectsrequest.scope.get("path") == "/v1/traces"for OTLP gzip bodies, so the mock raisedAttributeError, the generic handler swallowed it and_read_request_bodyreturned{}— the endpoint then never saw the body and the required-param guard masked the real validation error. Real starlette requests always carryscope, so this is mock under-specification, not a runtime defect.test_outer_deadline_delivers_session_termination[failure-True-True]intermittently assertedstarted.is_set()False on contended runners. The outeranyio.fail_after(0.2)budget started before the session handshake (initialize + notifications + tools/call), so a slow handshake consumed the cancellation window and the tool call never began.Fix
scope(path/v1/batches) so the body reader takes the same path as production.startedevent (tools/call actually in flight), then the original 0.2 sfail_afterremains the real outer cancellation deadline applied to the in-flight call. The call runs as a task so it progresses during the handshake;finallycancels it and awaits under suppression so session-termination (DELETE) completes before the test assertsdeleted.is_set(), exactly as the direct-await version guaranteed. Assertions, the 6.5 s bound and the cancellation semantics are unchanged.Validation
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py: 155 passed.tests/unit/experimental_mcp_client/test_mcp_client.py: 409 passed; 12test_outer_deadlinevariants plus the twotest_task_cancellation_during_cleanupvariants: 14 passed under an 8-way CPU load generator; variants were also green x3 unloaded after the change.