Skip to content

fix(spend-logs): preserve error_message on ProxyException failures - #2

Merged
songkuan-zheng merged 1 commit into
ship/v1.83.10from
fix/error-information-message-populated
May 15, 2026
Merged

fix(spend-logs): preserve error_message on ProxyException failures#2
songkuan-zheng merged 1 commit into
ship/v1.83.10from
fix/error-information-message-populated

Conversation

@songkuan-zheng

Copy link
Copy Markdown
Collaborator

Summary

StandardLoggingPayloadSetup.get_error_information stored an empty
error_message field in spend_logs.metadata.error_information for
every ProxyException, even though both the HTTP response body and
the captured traceback contained the full human-readable text.

Result: dashboard "LLM Failure" rows became untriageable — operators
had to manually unpack metadata.error_information.traceback JSON
via psql to distinguish a 401 from a 429 or a budget overrun.

Root cause

ProxyException (litellm/proxy/_types.py:3453):

class ProxyException(Exception):
    def __init__(self, message, type, param, code=None, ...):
        self.message = str(message)        # sets attribute
        # NOTE: does NOT call super().__init__(message)
        # NOTE: does NOT define __str__

Therefore str(ProxyException(...)) returns the empty string, and
get_error_information was calling str(original_exception) to
populate error_message. Every auth/budget/quota rejection landed
with error_message="".

Fix

litellm/litellm_core_utils/litellm_logging.py:5138-5148 — prefer the
.message attribute, fall back to str(exc) for non-litellm exception
types:

message_attr = getattr(original_exception, "message", None)
if message_attr:
    error_message = str(message_attr)
else:
    error_message = str(original_exception) if original_exception else ""

.message is set uniformly by ProxyException and every
litellm.exceptions.* class (see litellm/exceptions.py), so the
change is backward-compatible. Plain Exception / ValueError /
third-party exceptions still resolve via the str() fallback.

What this is NOT

  • Not a request-correctness bug — the 401/429/etc. responses to clients
    were always correct
  • Not a billing bug — the failures never reached the cost calculator
  • Purely an observability regression — only spend_logs was affected

Test plan

  • New unit test: test_get_error_information_prefers_message_attribute_over_str
  • New unit test: test_get_error_information_falls_back_to_str_when_no_message_attr
  • Existing test test_get_error_information_error_code_priority still passes
  • Existing test test_get_error_information still passes
  • End-to-end verified against the proxy: bad-key 401 now stores the
    full "Authentication Error, Invalid proxy server token passed..."
    message in spend_logs.metadata.error_information.error_message
  • Black 24.10.0 formatting clean

Out of scope

  • Direct fix at ProxyException (adding super().__init__(message))
    was considered but rejected: that class is marked DO NOT MODIFY
    for OpenAI-compatibility mapping reasons; the .message-preference
    fix at the logging layer is local and risk-bounded.

`StandardLoggingPayloadSetup.get_error_information` used
`str(original_exception)` to populate the human-readable error message
stored in `spend_logs.metadata.error_information.error_message`.

`ProxyException` (litellm/proxy/_types.py:3453) sets `self.message` in
its constructor but does NOT call `super().__init__(message)` and does
NOT define `__str__`. As a result, `str(ProxyException(...))` returns
the empty string, and every auth/budget/quota rejection was landing
in spend_logs with `error_message=""` despite a fully populated
traceback.

Operator impact: dashboard "LLM Failure" rows became untriageable —
the only way to tell a 401 from a 429 was to manually unpack the
traceback JSON via psql. Burst failure patterns (e.g. a UI session
polling with a stale token) produced 20-30 indistinguishable
`error_code=401` rows per second.

Fix: prefer the `.message` attribute (set by ProxyException and every
litellm.exceptions.* class) over `str(exc)`. The `str(exc)` fallback
is retained for non-litellm exception types, preserving prior behavior.

Test plan:
  - 2 new unit tests in tests/test_litellm/litellm_core_utils/
    test_litellm_logging.py:
    * test_get_error_information_prefers_message_attribute_over_str
    * test_get_error_information_falls_back_to_str_when_no_message_attr
  - Existing test_get_error_information_error_code_priority still passes
  - End-to-end verified: bad-key 401 now stores full
    "Authentication Error, Invalid proxy server token passed..."
    message in spend_logs.metadata.error_information.error_message
@songkuan-zheng
songkuan-zheng merged commit 7c0fa97 into ship/v1.83.10 May 15, 2026
1 check passed
songkuan-zheng added a commit that referenced this pull request May 18, 2026
Pass-through streaming requests (/v1/messages, /vertex_ai/*, /gemini/*,
/cohere/*, /assemblyai/*, /openai/*, /cursor/*) all share
PassThroughStreamingHandler.chunk_processor, which had two timing bugs
that interacted to collapse spend_logs.completionStartTime onto
spend_logs.endTime (off by ~1ms of clock resolution, not literally
identical) — making the streaming phase (endTime - completionStartTime)
round to roughly zero and TTFT effectively soak up the entire request
duration for every pass-through streaming row.

Root cause

1. `start_time` arg too late. The caller's start_time originates in
   BaseAnthropicMessagesStreamingIterator.__init__, which runs AFTER
   the upstream HTTP response has already been received. SpendLogs.
   startTime therefore reflects "moment we started reading the
   stream", not "moment the client request entered the proxy" — the
   real TTFT window is silently subtracted from Duration.

2. First-chunk arrival never recorded. The chunk loop yielded bytes
   to the client and collected them for logging, but never noted when
   the first byte arrived. With litellm_logging_obj.completion_start_time
   left as None, the fallback at litellm_logging.py:1834-1837 sets it
   to end_time — completionStartTime lands within ~1ms of endTime and
   streaming_phase rounds to 0.

Both bugs hide each other. Fixing only #2 gives TTFT close to 0 with
Duration deflated by ~TTFT. Fixing only #1 leaves completionStartTime
still pinned to endTime. Both must be fixed for the math to be correct.

Fix

In chunk_processor, at the top of the try block:
  - Override start_time with litellm_logging_obj.start_time when the
    latter is an earlier datetime — that's the true request-entry
    timestamp set in common_request_processing.base_process_llm_request.
  - On the first chunk yielded by response.aiter_bytes(), call
    litellm_logging_obj._update_completion_start_time(datetime.now())
    to populate the field that downstream payload builders look for.

Verified end-to-end (Anthropic claude-sonnet-4-6, 200-word stream):

  Before fix:  Duration=6528ms TTFT=6527ms streaming_phase=1ms
  After fix:   Duration=8496ms TTFT=2373ms streaming_phase=6123ms
  Control:     /v1/chat/completions (same model+prompt)
               Duration=8143ms TTFT=2071ms streaming_phase=6072ms

Test plan
  - New e2e case 13 (`13_passthrough_streaming_ttft.md` + data/13_*.sh)
    sends a real ~200-word streamed completion through /v1/messages,
    polls spend_logs for up to 30s, asserts:
      * streaming_phase_ms > 1000   (catches bug #2 regression)
      * ttft_ms > 300              (catches bug #1 regression)
      * ttft_ms < duration_ms / 2  (catches either regression)
    Plus a soft parity check vs /v1/chat/completions for the same
    upstream model.
  - Cost ~$0.005 per case run.
  - GREEN against the fix; was RED before (streaming_phase=1ms,
    ttft=6527ms) — assertions 1 and 3 fired.

Blast radius: all pass-through endpoints — they all flow through this
single chunk_processor. User noticed the bug via Anthropic; the same
fix improves TTFT observability for Vertex AI, Gemini, Cohere, and
the other pass-through providers in the same release.
songkuan-zheng added a commit that referenced this pull request Jun 4, 2026
* chore(e2e): add Claude-driven real-provider test harness

Adds a top-level e2e/ test harness designed to be driven by Claude
Code: scripts are single-purpose Unix tools, scenarios live as
markdown runbooks (cases/*.md), no pytest framework lock-in.

Why this exists separate from tests/:
  - tests/test_litellm/ uses mocks and never makes real provider calls,
    so it cannot catch integration bugs where mock-based unit tests pass
    but the real /metrics HTTP path or success-callback wiring is broken
  - Real-provider tests cost money per run and depend on external API
    availability — they must NEVER auto-run in `make test-unit` or CI
  - Long-running development on this fork needs a reusable harness, not
    one-off smoke scripts

Layout:
  e2e/
  ├── README.md                ← how to run, env contract
  ├── _config/
  │   └── docker-compose.yml   ← litellm built from local source + Postgres
  ├── tools/                   ← Unix-style single-purpose CLIs
  │   ├── proxy                ← lifecycle: start|stop|status|logs|rebuild
  │   ├── call                 ← one chat-completions request → JSON
  │   ├── metrics              ← /metrics: snapshot|diff|get
  │   ├── keys                 ← virtual-key CRUD + hash
  │   └── teams                ← team CRUD
  └── cases/
      ├── README.md            ← case index
      └── 01..12_*.md          ← runbooks Claude executes

Configuration via e2e/.env (gitignored). Postgres is ephemeral by
design — every `proxy stop` wipes data so test runs are reproducible
and stale virtual keys can't poison later cases.

Initial case set (12 cases):
  01-04: Anthropic prompt-cache metrics (5m/1h TTL, read, no-cache baseline)
  05-06: OpenAI prompt-cache metrics (cached_tokens, no creation metric)
  07:    /metrics endpoint smoke
  08-09: virtual-key + per-team label isolation (needs Postgres)
  10:    cost_breakdown must include cache_read_cost/cache_creation_cost
         when the static model_cost entry has cache rates (GREEN guard)
  11:    spend_logs.error_information.error_message must not be silently
         empty on auth failures — companion regression case for PR #2
  12:    custom_pricing path must not drop cache pricing when a
         dashboard-added deployment lacks cache_*_input_token_cost in
         litellm_params — currently RED, guards the future fix to
         router.py:7237 (deployment-UUID entry merge)

Also documents the fork's branching strategy in CLAUDE.md: ship/v1.83.10
is the long-term ship branch (TAG + accumulated fix/* merges);
internal/v1.83.10-stable is the upstream-sync working branch (1700+
upstream commits, NOT to be used as fix base); litellm_internal_staging
is pure upstream tracking. All internal fix PRs target ship/v1.83.10.

.gitignore additions cover the rendered config (e2e/_config/
.litellm.rendered.yaml, produced by `proxy start` from .env values)
and tool pycache.

Test plan:
  - Cases 01-07 verified GREEN against real Anthropic + OpenAI/GLM
    providers (Bug #1 metrics fix)
  - Case 11 GREEN after PR #2 merge (Bug #3 error_message fix)
  - Case 10 GREEN against current ship state
  - Case 12 deliberately RED — guards future router.py fix
  - No changes to tests/test_litellm/ or any other ci-managed paths

* chore(e2e): add --user-id sticky routing + master test runner

Two related improvements to the e2e test harness, both driven by
real-prod observations of the corp Anthropic gateway behavior:

1. `e2e/tools/call` learns `--user-id`. Forwarded to the proxy as the
   OpenAI `user` field, which litellm in turn maps to Anthropic's
   `metadata.user_id`. The corp gateway at maasapi.* uses this field
   for sticky upstream-key load balancing — requests sharing the
   same user_id land on the same upstream API key, so the second
   request's prefix can read the first request's cache write. Without
   sticky routing, the gateway round-robins anonymous requests across
   ~10+ upstream accounts each with its own cache namespace, and
   cache_read never hits in test scenarios.

   This was misdiagnosed initially as "gateway doesn't support cache
   reads". Prod logs show that requests carrying device_id in their
   end_user blob do achieve cache_read>0 (e.g. spend_logs entries
   e0ab6f96 and eec6a70d both hitting cache_read=128k-135k tokens
   under the same b0cb8e4d... device_id). Without device_id, the
   same gateway shows cache_read=0 on subsequent calls.

2. All Anthropic e2e cases (01, 02, 03, 04, 08, 09) now pass
   `--user-id` so cache-related assertions are deterministic
   regardless of upstream-account routing. Case 03 (cache READ) flips
   from SKIP-on-miss to required GREEN now that cache_read can be
   reliably triggered.

3. `e2e/tools/run-all-cases` ships as a first-class tool (was a
   /tmp scratch script). One PASS/FAIL/SKIP line per case, summary
   at the bottom, exits 0 iff every case PASSes (SKIPs allowed).
   Per-case logic split into `case_NN()` functions; adding a new case
   means dropping a fixture + adding one function + one invocation.

   Fixes a metric-snapshotting bug in the old runner: the previous
   `snap()` returned the LAST matching `/metrics` series, which silently
   compared two different series across before/after snapshots once
   later cases minted new virtual keys / teams. New `snap_sum()` sums
   across all series matching the label selector, giving stable totals
   even as the series set grows. Case 01 now passes cleanly through
   the runner.

   `--skip-paid` flag skips real-provider cases (05, 06, 13) for a
   ~$0 smoke pass against the harness itself.

4. `e2e/cases/data/11_error_information_message_populated.sh` polls
   spend_logs up to 20s instead of a flat 2s sleep — async logger lag
   was producing flake.

Test plan
  - `e2e/tools/run-all-cases` → 13/13 PASS
  - Case 03 cache_read deterministically hits 1827 tokens with --user-id
  - Case 01 no longer false-FAIL when runner is invoked after case 09
    has minted virtual keys/teams
  - Black 24.10.0 formatting clean on e2e/tools/call

* test(e2e): add case 16 — reset_budget_windows must not raise Prisma error

Verifies the cherry-pick of BerriAI#26346 at the full-stack
level: seed a key with `budget_limits` set, wait two ticks of the
background `ResetBudgetJob.reset_budget_windows`, then grep the
proxy container's logs for either the raw
`prisma.errors.MissingRequiredValueError` exception or the
"Failed to reset budget windows" wrapper line. Empty grep = PASS.

The case is intentionally key-path only — the team path is symmetric
and covered by unit tests
(`test_reset_budget_windows_resets_expired_team_window`,
`test_reset_budget_windows_query_error_does_not_break_team_path`).
Seeding a team via `/team/new` with `budget_limits` hits an unrelated
Prisma serialization bug in the team endpoint, which would mask the
result of this regression check.

Companion changes that the case relies on:

- `e2e/_config/docker-compose.yml`: pin
  `PROXY_BUDGET_RESCHEDULER_MIN_TIME=10` /
  `PROXY_BUDGET_RESCHEDULER_MAX_TIME=15`. Upstream default is ~600s
  (10 min) which makes the case impossible to verify inside a sane
  observation window; for dev/e2e there's no production reason to
  wait that long. The fixture skips (exit 77, not fail) if a stale
  container still has the upstream default — so an out-of-date
  environment doesn't masquerade as a regression.

- `e2e/tools/proxy`: split rebuild into two commands. `build` is the
  cached path (30-90s, default for source-only edits); `rebuild`
  keeps the original `--no-cache` semantics (3-5 min, for Dockerfile
  / dep changes). The cached path was always implicitly possible via
  `docker compose build`, but only `--no-cache` was surfaced through
  the tool, forcing a full rebuild on every fix-and-verify cycle.

- `CLAUDE.md`: document that root-level `e2e/` is the project's
  Claude-driven end-to-end harness, and that full-stack regression
  verification (DB schema, background jobs, real HTTP flow) belongs
  under `e2e/cases/` rather than under `tests/`.
songkuan-zheng added a commit that referenced this pull request Jun 4, 2026
… 6e) (#60)

* fix(passthrough): record real TTFT and start_time for streaming requests

Pass-through streaming requests (/v1/messages, /vertex_ai/*, /gemini/*,
/cohere/*, /assemblyai/*, /openai/*, /cursor/*) all share
PassThroughStreamingHandler.chunk_processor, which had two timing bugs
that interacted to collapse spend_logs.completionStartTime onto
spend_logs.endTime (off by ~1ms of clock resolution, not literally
identical) — making the streaming phase (endTime - completionStartTime)
round to roughly zero and TTFT effectively soak up the entire request
duration for every pass-through streaming row.

Root cause

1. `start_time` arg too late. The caller's start_time originates in
   BaseAnthropicMessagesStreamingIterator.__init__, which runs AFTER
   the upstream HTTP response has already been received. SpendLogs.
   startTime therefore reflects "moment we started reading the
   stream", not "moment the client request entered the proxy" — the
   real TTFT window is silently subtracted from Duration.

2. First-chunk arrival never recorded. The chunk loop yielded bytes
   to the client and collected them for logging, but never noted when
   the first byte arrived. With litellm_logging_obj.completion_start_time
   left as None, the fallback at litellm_logging.py:1834-1837 sets it
   to end_time — completionStartTime lands within ~1ms of endTime and
   streaming_phase rounds to 0.

Both bugs hide each other. Fixing only #2 gives TTFT close to 0 with
Duration deflated by ~TTFT. Fixing only #1 leaves completionStartTime
still pinned to endTime. Both must be fixed for the math to be correct.

Fix

In chunk_processor, at the top of the try block:
  - Override start_time with litellm_logging_obj.start_time when the
    latter is an earlier datetime — that's the true request-entry
    timestamp set in common_request_processing.base_process_llm_request.
  - On the first chunk yielded by response.aiter_bytes(), call
    litellm_logging_obj._update_completion_start_time(datetime.now())
    to populate the field that downstream payload builders look for.

Verified end-to-end (Anthropic claude-sonnet-4-6, 200-word stream):

  Before fix:  Duration=6528ms TTFT=6527ms streaming_phase=1ms
  After fix:   Duration=8496ms TTFT=2373ms streaming_phase=6123ms
  Control:     /v1/chat/completions (same model+prompt)
               Duration=8143ms TTFT=2071ms streaming_phase=6072ms

Test plan
  - New e2e case 13 (`13_passthrough_streaming_ttft.md` + data/13_*.sh)
    sends a real ~200-word streamed completion through /v1/messages,
    polls spend_logs for up to 30s, asserts:
      * streaming_phase_ms > 1000   (catches bug #2 regression)
      * ttft_ms > 300              (catches bug #1 regression)
      * ttft_ms < duration_ms / 2  (catches either regression)
    Plus a soft parity check vs /v1/chat/completions for the same
    upstream model.
  - Cost ~$0.005 per case run.
  - GREEN against the fix; was RED before (streaming_phase=1ms,
    ttft=6527ms) — assertions 1 and 3 fired.

Blast radius: all pass-through endpoints — they all flow through this
single chunk_processor. User noticed the bug via Anthropic; the same
fix improves TTFT observability for Vertex AI, Gemini, Cohere, and
the other pass-through providers in the same release.

* hotfix(bedrock-anthropic): clean stray conflict markers from Wave 6d merge

PR #59 (Wave 6d) accidentally committed unresolved cherry-pick conflict
markers in the Bedrock anthropic_claude3_transformation Invoke filter.
The file is broken in ship/v1.87.0 — the Python module fails to import.

Fix: apply the intended Wave 6d resolution — keep upstream's single
`filtered_betas = sorted(...)` shape from PR BerriAI#26148, layer our
`overrides=_overrides` argument on top so the
`anthropic_beta_overrides` per-deployment config still threads through.

This commit only touches the unresolved hunk; no behavioral change vs
the intent of Wave 6d.
songkuan-zheng added a commit that referenced this pull request Jun 10, 2026
…axonomy (#78)

* fix(streaming): reset Anthropic message_start cursor (output_tokens=1) when no message_delta arrives

The Anthropic streaming protocol emits `message_start.usage.output_tokens=1`
as a placeholder cursor; the real cumulative output count only arrives in
the final `message_delta` event. When a stream is cancelled before
`message_delta` lands (common for thinking models on long-tail prompts),
ChunkProcessor._calculate_usage_per_chunk's last-wins accumulator left
completion_tokens stuck at 1. Because 1 is truthy, the
`completion_tokens or token_counter(text=...)` fallback in
calculate_usage() never fired, and requests were billed for 1 output
token even when several thousand tokens of text had actually streamed.

Fix: track whether any chunk's completion_tokens exceeded 1
(saw_non_cursor_completion). If the only update we saw was the cursor,
reset completion_tokens to 0 so the text-based fallback estimates from
the real completion content.

Legitimate 1-token completions (model returns "Yes." etc.) are unaffected
in practice — token_counter on a 1-token completion_output also yields
~1, so billing stays approximately correct.

Tier: C (universal bug fix — affects every LiteLLM user calling Anthropic
with streaming + cancel/timeout). Upstream PR candidate once landed here.

Tests added:
- TestAnthropicCursorBug (6 cases) — pins the post-fix behavior
- TestNonAnthropicStreamingIntact (2 cases) — guards against regression on
  providers without the cursor pattern

All 9 existing streaming_chunk_builder_utils tests still pass.

* feat(spend_logs): add success_partial status + cancellation metadata fields

Adds the schema scaffolding for tracking client-cancelled requests as a
billable "success_partial" status (instead of dropping them into the
failure bucket, which both pollutes the proxy failure rate metric and
zeroes out billing for compute the upstream provider already charged us
for).

New StandardLoggingPayloadStatus value:
  - "success_partial": client disconnected mid-flight (or upstream cut
    early) but upstream consumed billable compute. Bills prompt + the
    output that was generated up to cancellation. Does NOT count toward
    the proxy failure rate.

New Literal types in litellm/types/utils.py (used by PR #3 cancel-billing
path to populate metadata fields):
  - CancelPhase: before_upstream | during_upstream | streaming_partial |
                 during_parsing
  - CancelUsageSource: upstream_truth | tokenizer_estimate |
                       upstream_completed_after_cancel | shield_timeout |
                       no_completion

New StandardLoggingMetadata + SpendLogsMetadata fields (all Optional,
populated only when the cancel-billing path fires):
  - cancellation_indicator: client_disconnect | upstream_disconnect
  - cancel_phase: lifecycle phase when cancel was detected
  - bytes_delivered_to_client: total bytes ACK'd to client before disconnect
  - upstream_completed: whether upstream call ran to completion
  - usage_source: provenance of the recorded usage numbers

DB compatibility:
  - LiteLLM_SpendLogs.status is `String?` (not enum) — accepts new value
    without migration.
  - LiteLLM_SpendLogs.metadata is `Json?` — new fields land inside the JSON
    blob, no migration required.
  - Existing `metadata.error_information.error_code` (used by nginx/
    Prometheus dashboards) is unchanged; cancellation fields are additive.

Tier: B (internal change to a TypedDict + helper init; downstream
behavior change ships in PR #3). The success_partial status name and
field semantics are coordinated with the upcoming cancel-billing module.

Tests:
  - 10 new cases covering literal acceptance, None-init, round-trip,
    error_code coexistence, and enum value coverage.
  - All 73 existing spend_tracking tests pass.

* feat(cancel): catch CancelledError + re-route streaming/non-stream cancels to success_partial path

Closes the 499 black hole. asyncio.CancelledError is a BaseException
subclass since Python 3.8 — every `except Exception` in LiteLLM was
letting it slip through silently:
  * SpendLogs got no row for the cancelled request
  * No callback fired (Langfuse trace stuck "Running", Prometheus
    failed_requests counter not incremented)
  * The upstream provider continued generating after the client TCP
    close and billed us for compute that never reached our ledger

This PR introduces a small finalize layer that catches the cancel
signal, marks the Logging object with cancel_phase / cancellation
metadata (matching PR #4's schema), dispatches the partial response
through the normal async_success_handler path with the new
status="success_partial" classification, and re-raises CancelledError
so asyncio's cancellation contract is preserved.

Files
-----
* litellm/litellm_core_utils/cancel_finalize.py (new)
    Public helpers:
      - mark_logging_obj_cancelled(logging_obj, phase, indicator,
        bytes_delivered): idempotent dict mutation; pathological inputs
        (None logging_obj, missing model_call_details) are no-ops.
      - finalize_streaming_cancel(stream_wrapper, logging_obj, ...):
        runs stream_chunk_builder on accumulated chunks → dispatches
        through async_success_handler with success_partial intent.
        Shielded internally so the SpendLogs write completes even if
        the runtime injects more cancel signals during teardown.
        Falls back to post_call_failure_hook if there are zero chunks.
      - finalize_non_stream_cancel(upstream_task, logging_obj, ...,
        shield_timeout_s=60.0): shields the upstream call past the
        cancel and waits for real usage (strategy A). On timeout,
        cancels upstream and records usage_source="shield_timeout".

* litellm/proxy/proxy_server.py — async_data_generator
    Added `except asyncio.CancelledError:` before `except Exception:`
    that calls finalize_streaming_cancel and re-raises.

* litellm/proxy/common_request_processing.py — async_streaming_data_generator
    Same catch for the /v1/messages and /v1beta/...generateContent
    paths (Anthropic + Google).

* litellm/litellm_core_utils/streaming_handler.py — __anext__
    Defense-in-depth: if cancel reaches the stream wrapper directly
    (not via the proxy generator), mark the Logging object and
    re-raise so whoever does finalize the request gets the marker.

* CLAUDE.md — Test discipline
    Added "No theater tests" rule: mock the boundary (httpx transport,
    DB cursor), not the unit under test. Use a small spy fake to
    capture handler args; assert on the captured data, not on
    `.called`. Example: instead of mocking `stream_chunk_builder`,
    pass real ModelResponseStream chunks in. The rule is enforced in
    this PR's test file.

Tests (18 cases, all real — no mocking of the unit under test):
  * mark_logging_obj_cancelled: dict-mutation contract incl.
    idempotency, None handling, no-pollution-with-Nones
  * is_logging_obj_cancelled: marker detection
  * _get_accumulated_chunks: stream wrapper variants
  * finalize_streaming_cancel:
      - Real Anthropic-shaped chunks (message_start cursor=1 + content
        deltas, NO message_delta) run through real stream_chunk_builder
        → captured response.usage.completion_tokens > 1, exercising
        the cross-PR contract with PR #1's cursor reset
      - No-chunks edge: real post_call_failure_hook spy verifies
        fallback fires with a CancelledError, not a swallowed signal
      - Downstream success_handler raising must NOT propagate
  * finalize_non_stream_cancel:
      - Shield-wait happy path: real asyncio task, real response
        flows through to captured success kwargs
      - Shield timeout: upstream task actually cancelled
      - Upstream exception during shield window
      - No upstream task (before_upstream cancel)

Cost computation for the partial response — the actual implementation
of "modified strategy 5'" billing — lands in PR #3 (cancel_billing.py)
which reads the markers this PR sets. Until PR #3 lands, the
reassembled partial response runs through the existing cost calculator
unchanged, which is already a significant improvement over the prior
"drop everything" behavior because PR #1's cursor=1 fix produces
sensible partial usage for Anthropic.

Tier: C+D (Tier C for the BaseException catch fix — purely a bug;
Tier D for the success_partial taxonomy — opinionated mechanism with
billing implications we want carried in the fork until upstream agrees
on semantics).

* feat(cancel): compute partial cost + bridge cancel metadata to SpendLogs

Completes the cancel-billing path started in PR #2. cancel_finalize was
catching cancels and dispatching through async_success_handler with a
reassembled partial response, but two gaps remained:

1. The cancellation markers written to logging_obj.model_call_details
   never reached SpendLogs — _get_spend_logs_metadata pulls fields from
   request_data.litellm_params.metadata, not from model_call_details.
   Result: SpendLogs rows existed (good, no more black hole) but had
   no cancel_phase / usage_source / etc., so dashboards couldn't tell
   normal vs partial.

2. The zero-chunk fallback path (finalize → failure hook) still wrote
   response_cost=0.0 — the upstream provider received the prompt and
   started processing, but our ledger said "free". For thinking models
   on long prompts this is a real $ leak ($0.003-$0.01 per 499 in
   production).

Files
-----
* litellm/litellm_core_utils/cancel_billing.py (new):
    - compute_prompt_only_cost(messages, model, custom_llm_provider):
      best-effort prompt-only $ via the real LiteLLM cost map. Returns
      0.0 on any pricing failure — caller is in the failure hook and
      cannot afford an exception. Wraps token_counter + cost_per_token
      so the cost-map lookup path is exercised by real provider IDs.
    - enrich_request_metadata_with_cancel_markers(request_data,
      logging_obj): copies the five cancel markers from
      model_call_details into litellm_params.metadata. Idempotent;
      no-op when there's no marker (normal requests pay zero
      overhead). Overrides metadata.status to "success_partial".

* litellm/litellm_core_utils/cancel_finalize.py:
    - finalize_streaming_cancel now calls
      enrich_request_metadata_with_cancel_markers right after
      mark_logging_obj_cancelled, so the partial response dispatched
      through async_success_handler carries the markers into the
      cost-tracking pipeline.
    - finalize_non_stream_cancel does the same bridge, AND a second
      bridge after the shield-wait outcome is known (upstream_completed
      / usage_source can change after the asyncio.wait_for resolves).

* litellm/proxy/hooks/proxy_track_cost_callback.py:
    - async_post_call_failure_hook: when the failure is actually a
      CancelledError (i.e. the cancel-finalize fallback path because
      no chunks were received), compute the prompt-only cost via
      cancel_billing.compute_prompt_only_cost instead of the hardcoded
      0.0. Other failures (provider 5xx, real timeouts) still write
      0.0 — there's no compute on our side to bill in those cases.
    - Same hook also calls enrich_request_metadata_with_cancel_markers
      so failure-path SpendLogs rows tagged success_partial get the
      right metadata even when called by external code paths that
      didn't go through cancel_finalize first.

Tests (12 new cases, no mocking of cost_per_token / token_counter):
  * TestComputePromptOnlyCost:
      - Real Anthropic & OpenAI pricing paths exercised → cost > 0
        with sanity upper bound (catches both "always 0" and "wrong
        rate" regressions)
      - Unknown model degrades to 0.0 without raising
      - Long-prompt cost scales with token count (sanity that we're
        actually counting, not constants)
  * TestEnrichRequestMetadataWithCancelMarkers:
      - No-op when no marker (must not pollute normal requests)
      - All five markers + status override propagate
      - Creates missing litellm_params.metadata dict
      - Partial markers don't insert keys for absent ones
      - Overwrites stale status (success / failure → success_partial)

All 43 Phase 1 tests pass together (cursor + schema + cancel_finalize
+ cancel_billing). 82 streaming + spend_tracking regression tests pass.

Tier: D (mechanism with billing policy — bills > 0 for cancelled
requests, which is a behavior change downstream consumers can observe).

* test(e2e): add case 26 cancel-billing + plumb success_partial through SpendLogs

Adds an e2e case that verifies the full Phase 1 cancel-billing chain
end-to-end against the live proxy + Postgres, and fixes three real
plumbing bugs the case exposed:

1. FallbackStreamWrapper.__anext__ did not accumulate chunks
   FallbackStreamWrapper inherits from CustomStreamWrapper but overrides
   __anext__ to delegate to its own async generator — bypassing the
   parent's `self.chunks.append(chunk)`. Result: any streaming request
   that went through the Router fallback wrapper had empty
   wrapper.chunks, so finalize_streaming_cancel's reassembly path saw
   "no chunks" and fell through to the failure-hook fallback. Fix:
   append in __anext__ to mirror the parent.

2. enrich_request_metadata_with_cancel_markers only touched one of two
   metadata dicts
   The proxy failure path reads metadata from
   request_data["litellm_params"]["metadata"]; the success path reads
   from logging_obj.litellm_params["metadata"]. These are usually the
   same object but can diverge after construction-time copies. Fix:
   update both. Detected via case 26 showing status="success"
   completion_tokens=200 with all cancel markers blank — the success
   path was getting the partial response but not the markers.

3. _get_status_for_spend_log only knew "success" / "failure"
   metadata.status was set to "success_partial" by the cancel-finalize
   bridge, but _get_status_for_spend_log collapsed any non-"failure"
   value back to "success" before writing the LiteLLM_SpendLogs.status
   column. Dashboards filtering `WHERE status='success_partial'` would
   have matched zero rows. Fix: widen the Literal to three values and
   pass success_partial through unchanged.

4. _ProxyDBLogger.async_post_call_failure_hook clobbered cancel markers
   and hardcoded status="failure"
   The failure hook rebuilt litellm_params.metadata from scratch,
   preserving only "tags" — cancel markers set by cancel_finalize were
   lost on the zero-chunk fallback path. Also hardcoded metadata.status
   = "failure" even for CancelledError. Fix: detect CancelledError and
   classify as success_partial; preserve the five cancel marker fields
   alongside the existing tags preservation.

5. time.time() vs datetime.now() in finalize success_handler calls
   async_success_handler expects datetime instances for start_time /
   end_time (it does arithmetic on them); passing time.time() floats
   raised TypeError in the cost calculator and silently dropped the
   row. Fix: datetime.now().

Test plan (live proxy + Postgres):
* Case 26 C1: streaming cancel mid-flight       → status=success_partial,
                                                  completion_tokens=214,
                                                  cancel_phase=streaming_partial
* Case 26 C2: long stream cancelled near end    → status=success_partial,
                                                  completion_tokens=760
* Case 26 C3: zero-chunk cancel                 → status=success_partial,
                                                  completion_tokens=14
* Full mock-only e2e suite: 17/17 pass, 0 fail, 8 skip (expected
  real-provider skips). No regression in cases 04, 06, 07, 10, 11,
  12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 23, 25.

Unit tests:
* 43 Phase 1 tests still pass (cursor + schema + cancel_finalize +
  cancel_billing).
* 82 streaming + spend_tracking regression tests still pass.

Tier: C (FallbackStreamWrapper accumulation, _get_status_for_spend_log
widening — universal bug fixes) + D (success_partial classification,
metadata bridging policy).

* test(e2e): cases 27-32 cancel-billing expansion + Phase 2 gaps documented

Extends the cancel-billing e2e coverage with five new mock-only cases.
Two PASS outright, three honestly SKIP with detailed deferral reasons
that point at the remaining Phase 2 work.

PASSING (new):
  29 failure-not-polluted    — Forces upstream 503; asserts SpendLogs
                               row classifies as status=failure with
                               error metadata populated. Critical
                               regression guard: without this, the
                               CancelledError check in
                               proxy_track_cost_callback could silently
                               grow to catch other exceptions and
                               start polluting success_partial.
  32 concurrent-cancels      — Fires 10 concurrent stream cancels,
                               asserts all 10 SpendLogs rows write
                               with success_partial + the marker,
                               and a control health probe returns
                               in < 3s (no asyncio task leak under
                               contention).

DEFERRED (skip 77, with detailed runbook comments):
  27 non-stream-cancel       — finalize_non_stream_cancel exists and
                               is unit-tested, but no caller wires it
                               up. LiteLLM removed the
                               check_request_disconnection polling
                               task (it's dead code in proxy_server.py),
                               so non-stream cancel detection requires
                               a Phase 2 polling revive. Until then
                               non-stream cancels look like normal
                               successes from the proxy's view (which
                               is at least billing-correct, just
                               unclassified).
  28 v1-messages-cancel      — The CancelledError catch in
                               common_request_processing.py IS firing
                               (verified: SpendLogs row gets correct
                               partial completion_tokens reflecting
                               bytes received before cancel), but the
                               success_partial markers don't propagate
                               because response.logging_obj is None
                               on the anthropic_messages path.
                               Phase 2: pull logging_obj from
                               request_data["litellm_logging_obj"]
                               instead of from response.
  30 shield-timeout          — Depends on case 27 wiring. Implemented
                               the LITELLM_CANCEL_SHIELD_TIMEOUT_S env
                               var (default 60s, e2e sets 2s) so the
                               case can run deterministically once
                               polling is wired.

Other changes:
  * e2e/_config/mock_provider.py: added time.sleep(ttft_ms) to the
    Anthropic non-stream path so case 27/30 can deterministically
    drive client-cancel timing once the proxy plumbing lands.
  * e2e/_config/docker-compose.yml: sets
    LITELLM_CANCEL_SHIELD_TIMEOUT_S=2 on the proxy container.
  * litellm/litellm_core_utils/cancel_finalize.py: reads
    LITELLM_CANCEL_SHIELD_TIMEOUT_S env var at module load for the
    DEFAULT_CANCEL_SHIELD_TIMEOUT_S knob.

Suite outcome: 30 cases · 19 PASS · 0 FAIL · 11 SKIP
(11 skips = 6 pre-existing Tier=real + 1 pre-existing thinking-signature
+ 1 pre-existing anthropic-beta-overrides + 3 new Phase-2-deferred).

Tier: B (e2e infrastructure + documentation of Phase 2 gaps).

* feat(cancel): Phase 2 — non-stream cancel detection + /v1/messages markers

Closes the remaining Phase 2 gaps documented in commit 987df5f:

1. /v1/messages cancel markers didn't reach SpendLogs (case 28)
   The CancelledError catch in async_streaming_data_generator was firing
   correctly but `getattr(response, "logging_obj", None)` returned None
   on the Anthropic native path (response is an async iterator without
   that attribute), so mark_logging_obj_cancelled became a no-op.
   Additionally, `_get_litellm_metadata_from_kwargs` prefers the newer
   `litellm_metadata` dict over `metadata` — and the cancel markers were
   only written to the latter, so they were filtered out by the SpendLogs
   pipeline even when present on the Logging instance.

2. Non-stream cancel was a silent black hole (case 27, 30)
   No proxy code path detected client disconnect during non-stream LLM
   calls. The dead-code `check_request_disconnection` helper had been in
   proxy_server.py for ages but with zero callers.

Fixes
-----

litellm/proxy/proxy_server.py + litellm/proxy/common_request_processing.py
  Both streaming catches: when response.logging_obj is None, fall back to
  request_data["litellm_logging_obj"]. Required for /v1/messages,
  /v1beta/...streamGenerateContent and any other endpoint where the
  response object isn't a CustomStreamWrapper.

litellm/litellm_core_utils/cancel_billing.py
  enrich_request_metadata_with_cancel_markers now writes to BOTH the
  `metadata` and `litellm_metadata` keys on each litellm_params dict
  (request_data + logging_obj.litellm_params). litellm_metadata is the
  newer-endpoint variant; without writing to it, /v1/messages
  cancellations couldn't surface success_partial markers because the
  cost callback's metadata extractor returns litellm_metadata when both
  are present.

litellm/proxy/common_request_processing.py — base_process_llm_request
  Added a _disconnect_watcher task that polls request.is_disconnected()
  every second while the LLM call is in flight. On detection:

  * Records the disconnect time and KEEPS the LLM call running. We
    don't want to cancel it: the upstream provider has almost certainly
    started processing and will charge us regardless, so the right
    thing is to record real upstream usage.
  * If the LLM call completes within LITELLM_CANCEL_SHIELD_TIMEOUT_S
    (default 60s, env-tunable) of the disconnect, the row gets tagged
    with usage_source=upstream_completed_after_cancel and
    upstream_completed=True.
  * If the budget elapses with the LLM still running, give up waiting,
    synthesize an empty response so the handler can return, and tag
    usage_source=shield_timeout / upstream_completed=False. The real
    response (when it eventually comes back from upstream) will fire
    a SECOND SpendLogs row via the normal success path — that's
    acceptable because it preserves the upstream-billing alignment.

  Implemented as a `wait_for(shield(...), poll_interval)` loop rather
  than a single wait_for + shield, so the watcher task gets a chance
  to set disconnect_flag between polls without being interrupted by
  the timeout itself. Cleaned up via a finally that cancels the
  watcher task regardless of how the LLM call resolved.

e2e/cases/27, 28, 30
  Un-skipped; all three now PASS against the live proxy.

Test plan (mock-only e2e suite, full run):
  * 30 cases · 22 PASS · 0 FAIL · 8 SKIP
  * 8 SKIP = 6 Tier=real + 1 thinking-signature (pre-existing) + 1
    anthropic-beta-overrides (Tier=real, requires Bedrock). All cancel-
    related cases now PASS — no remaining Phase 2 deferrals.
  * 125 unit tests still pass (cancel_billing 12 + cancel_finalize 18 +
    cursor 8 + spend_logs metadata 5 + streaming regression 9 + spend
    tracking regression 73).

Tier: C (universal bug fix — anyone relying on non-stream cancel
detection or the /v1/messages cancel taxonomy was getting wrong rows)
+ D (the disconnect-watcher's shield-and-wait policy is opinionated
billing behavior we want to carry until upstream agrees on semantics).

* feat(cancel): Phase 3 — orthogonal delivery_status/billing_status taxonomy

Replace the single `status='success_partial'` marker with two orthogonal
dimensions derived at SpendLog write time from the five cancel markers
plus the raw status:

  delivery_status ∈ {full, partial, none}   # what reached the client
  billing_status  ∈ {full, partial, none}   # what we billed

The top-level SpendLogs status filter stays binary (success | failure);
cancellations carry `status='success'` with metadata markers and the UI
surfaces them through the new derived columns. Callers that need to
target the cancel slice in SQL can filter on
`metadata::jsonb->>'delivery_status'` / `billing_status` directly.

Scope:
- core derivation in `spend_tracking_utils._derive_delivery_billing_status`
  is the single source of truth — no other writer sets these fields.
- e2e cases 26–30 and 32 updated to assert both dimensions; case 33
  (real-anthropic-cancel) added end-to-end.
- UI: view_logs columns and filter_options surface the new fields.
- Tests added for the new derivation and for the binary status filter
  (test_status_filter_condition).

* docs(cancel): post-Phase-3 cleanup — rename case 26, refresh case 33 doc, harden case 30 preflight, log upstream PR candidates

- e2e/cases/26 renamed: success_partial → partial (refactor dropped the
  three-valued status sentinel). Updated dispatcher reference in
  run-all-cases.
- e2e/cases/26 + 33 .md docs rewritten to describe the new binary
  status + orthogonal delivery_status / billing_status taxonomy,
  including the failure-mode lookup table for case 33 (real Anthropic).
- e2e/cases/data/30_shield_timeout.sh: add the mock-container
  healthcheck preflight that the other mock-only cases already have,
  so real-mode runs SKIP (rc=77) instead of FAIL.
- UPSTREAM_PR_QUEUE.md: log two Tier-C candidates (cursor=1 reset +
  CancelledError catch) and the Phase 3 taxonomy as a Tier-D
  ISSUE-FIRST entry. Refresh Last reviewed.

Tier: B (internal infra / docs).
Tried upstream first: N/A — purely internal cleanup.
songkuan-zheng pushed a commit that referenced this pull request Jun 15, 2026
* restore an explicit no-match policy

* fix(jwt): fix AUTO_REGISTER sentinel bypass, race condition, and inline import comment

- AUTO_REGISTER now evicts stale __NO_MAPPING__ sentinel instead of silently
  returning None when cached under a prior fallback_team_mapping config
- Race condition in _auto_register_jwt_mapping: catch P2002 unique-constraint
  violation on concurrent creates, fetch the winning mapping, proceed cleanly
- Added comment on inline generate_key_helper_fn import explaining the circular
  dependency (key_management_endpoints imports user_api_key_auth at line 51)
- 3 new tests: stale sentinel eviction, race condition winner fallback, and the
  existing auto_register happy path

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(jwt): cache __NO_MAPPING__ sentinel before raising 403 in REJECT mode

REJECT mode was raising HTTPException immediately on a DB miss without writing
the __NO_MAPPING__ sentinel, causing every subsequent rejected request to
re-query the DB. Write the sentinel first so repeated rejections are served
from cache within virtual_key_mapping_cache_ttl.

Adds test asserting DB is not hit on the second reject after a cache-warm miss.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(jwt): enforce no-match policy when prisma_client is None

The early `if prisma_client is None: return None` guard ran before the
no-match policy check, silently bypassing REJECT and AUTO_REGISTER — every
JWT client fell through to team auth regardless of configuration.

Fix: treat prisma_client=None as a definitive DB miss and fall through to the
same policy block as a real miss. REJECT now raises 403, AUTO_REGISTER raises
500 with a clear message (can't create keys without a DB), FALLBACK_TEAM_MAPPING
returns None unchanged.

Adds three tests: REJECT/403 with no DB, FALLBACK returns None with no DB,
AUTO_REGISTER/500 with no DB.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(jwt): consistent AUTO_REGISTER on cached sentinel; clean up race orphans

Addresses Greptile review on PR BerriAI#25570 cherry-pick.

1. Inconsistent AUTO_REGISTER when __NO_MAPPING__ sentinel is cached:
   The cached-sentinel branch silently returned None when prisma_client was
   None, while the fresh path raised HTTP 500 under the same config. Same
   request, different access-control outcome depending on cache state. Both
   paths now raise the same 500.

2. Orphaned virtual keys from race-condition losers:
   On unique-constraint conflict, generate_key_helper_fn had already persisted
   an unrestricted virtual key in LiteLLM_VerificationToken with the cleartext
   in request memory. Under sustained concurrency these accumulated
   indefinitely. The loser now deletes its orphan before falling back to the
   winner's mapping; failure to delete is logged but does not fail the request.

Also corrects a latent FK bug surfaced while fixing #2: the mapping row was
storing the plaintext key in LiteLLM_JWTKeyMapping.token, but that column FKs
to the hashed LiteLLM_VerificationToken.token — now hashed at the call site.

Tests:
- updated test_auto_register_creates_key_and_mapping to assert the hashed
  token is stored, not the plaintext
- updated test_auto_register_race_condition_unique_conflict to assert the
  orphan is deleted with the correct hashed token
- added test_auto_register_raises_500_when_sentinel_cached_and_no_db
- added test_auto_register_race_conflict_tolerates_delete_failure

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(jwt): close REJECT bypass when JWT omits the configured claim field

A JWT presented without the configured `virtual_key_claim_field` previously
returned None at the `claim_value is None` guard before the
`unregistered_jwt_client_behavior` check ran. A caller who knows the configured
claim-field name could bypass REJECT by simply omitting that field and falling
through to team-based JWT auth.

Apply the no-match policy on a missing claim:
  - REJECT          → 403
  - AUTO_REGISTER   → 403 (no stable identity to map; refuse rather than
                     create a sentinel-keyed record)
  - FALLBACK_TEAM_MAPPING → return None (unchanged, backward-compatible)

Adds three tests covering each branch of the missing-claim path.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(jwt): AUTO_REGISTER inherits team_id so keys are bounded by team limits

Auto-registered virtual keys were created with no team, model, route, rate, or
budget constraints — broader access than the standard team-based JWT auth path
the same client would have taken. Under AUTO_REGISTER, resolve the team_id
from the JWT (via the operator-configured team_id_jwt_field / team_id_default)
and stamp it on the new key. Downstream auth then applies the team's
budget/models/tpm/rpm/allowed_routes via the existing virtual-key flow.

Policy when team_id_jwt_field is configured:
  - JWT carries team claim → stamp resolved team_id
  - JWT lacks claim + team_id_default set → stamp default
  - JWT lacks claim + no default → 403 (refuse to create an unbounded key)

When neither team_id_jwt_field nor team_id_default is configured, the
operator has explicitly opted out of team-based limits — the auto-created
key has no team_id (matches what team-auth would do in the same config).

Adds 4 tests covering each branch.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(jwt): make AUTO_REGISTER functional in prod; raise on missing winner

Two correctness fixes flagged by Greptile on the AUTO_REGISTER path:

1. generate_key_helper_fn was called without table_name="key". Without that,
   the helper falls into the user-upsert branch (table_name in (None, "user"))
   and tries to insert into LiteLLM_UserTable with user_id=None, which hits
   the NOT NULL @id constraint. AUTO_REGISTER would never have succeeded in
   production. Now passes table_name="key" explicitly, matching the
   /key/generate caller.

2. When the race loser refetches the winner's mapping and gets None (winner
   row concurrently deleted), the previous code returned None — and the
   caller in _resolve_jwt_to_virtual_key then fell through to less-
   restrictive team-based JWT auth, silently bypassing the configured
   AUTO_REGISTER policy. Now raises HTTP 503 so the caller retries against
   a stable state rather than getting unintended fallback access.

Adds one test for the 503 winner-vanishes path.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(jwt): defer AUTO_REGISTER until JWT policy is enforced by auth_builder

Closes the JWT policy bypass on the AUTO_REGISTER path flagged by veria-ai.

Before: when unregistered_jwt_client_behavior=auto_register and the JWT's
claim was unmapped, _resolve_jwt_to_virtual_key validated the JWT signature
and then immediately created a virtual key + mapping. JWTAuthManager.auth_builder
never ran for the first request (the new key short-circuited the team-auth
path), and every subsequent request hit the cached mapping — so custom_validate,
RBAC, scope_mappings, and user_allowed_email_domain were never enforced for
auto-registered clients.

After: _resolve_jwt_to_virtual_key returns a _PendingAutoRegister signal
instead of creating the key. The caller in _user_api_key_auth_builder runs
JWTAuthManager.auth_builder, then — only on a validated, policy-passing
result — calls _auto_register_jwt_mapping with the team_id / user_id from
that result. The created key inherits team + user limits from the validated
identity, and future cache hits load that already-policy-checked key.

Also drops the interim _resolve_inherited_team_id helper that pulled team_id
from raw JWT claims — same bypass risk; team_id now comes exclusively from
auth_builder.

Tests:
  - Rewrote two existing tests to assert _resolve_jwt_to_virtual_key returns
    _PendingAutoRegister (no key created yet) for both the fresh-DB-miss
    and stale-sentinel branches
  - Added a contract test that _auto_register_jwt_mapping stamps the
    validated team_id/user_id onto generate_key_helper_fn
  - Removed four stale team-binding tests that exercised the prior
    raw-claim helper

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* Update user_api_key_auth.py

* fix(jwt): cache proxy-admin AUTO_REGISTER path to avoid repeated DB lookups

Cache-miss regression introduced by the deferred-auto-register refactor:
when a JWT under AUTO_REGISTER resolved to a proxy admin, the is_proxy_admin
early-return in _user_api_key_auth_builder ran *before* the pending
auto-register cache-write block. Result: no cache entry, so every
subsequent proxy-admin request re-queried get_jwt_key_mapping_object
indefinitely.

Fix: write a __JWT_PROXY_ADMIN__ sentinel to user_api_key_cache before the
early return when a pending auto-register existed. _resolve_jwt_to_virtual_key
treats that sentinel as "skip mapping, fall through to auth_builder", so
future requests from the same JWT identity hit the cache instead of the DB.
auth_builder still runs full JWT policy on every request — only the
mapping DB lookup is short-circuited.

Adds one test asserting the sentinel cache-hit returns None without
hitting prisma_client.db.litellm_jwtkeymapping.find_first.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(proxy): stamp org context on JWT auto-registered keys

AUTO_REGISTER keys were created with team_id and user_id only, so org budget checks were skipped after switching to the key-scoped path.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
songkuan-zheng added a commit that referenced this pull request Jun 17, 2026
veria-ai's review of a501b0e raised two Medium-severity findings:

1. /v2/model/info default branch did not apply the BYOK team-scope
   filter when caller's key/team/user model lists were all empty.
   Result: a virtual key with empty restrictions could see other
   teams' BYOK deployment metadata (api_base, model_info.id, etc.).

2. /v1/model/info?litellm_model_id=<id> (by-id branch) applied only
   the user.models filter after the option-B refactor, skipping
   key.models / team_models / BYOK scope. A key restricted by
   key.models with no personal model restriction could pull metadata
   for any deployment.

Finding #2 is a regression introduced by switching the by-id branch
from a 403 gate to filter-chain semantics; the option-B refactor
applied only the personal-models filter and missed the listing-path's
key/team intersection plus the BYOK scope filter that the v1 default
branch already used.

Fix mirrors the v1 default branch's filter chain on the by-id path
and adds the BYOK filter to v2:
- v1 by-id: insert _get_v1_model_info_allowed_model_names +
  _filter_v1_model_info_deployments (key/team), keep the existing
  user.models filter, append the BYOK _get_caller_byok_team_scope +
  _byok_row_outside_caller_teams filter.
- v2 default: append the same BYOK filter after user.models.

All filters use existing helpers; no new code paths introduced.
Filter (drop deployment → empty data) preserved instead of 403, per
the route_checks.py:189 spec.

Local verification: 87/87 across
- tests/test_litellm/proxy/discovery_endpoints/
- tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
- tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
songkuan-zheng added a commit that referenced this pull request Jun 17, 2026
veria-ai's review of a501b0e raised two Medium-severity findings:

1. /v2/model/info default branch did not apply the BYOK team-scope
   filter when caller's key/team/user model lists were all empty.
   Result: a virtual key with empty restrictions could see other
   teams' BYOK deployment metadata (api_base, model_info.id, etc.).

2. /v1/model/info?litellm_model_id=<id> (by-id branch) applied only
   the user.models filter after the option-B refactor, skipping
   key.models / team_models / BYOK scope. A key restricted by
   key.models with no personal model restriction could pull metadata
   for any deployment.

Finding #2 is a regression introduced by switching the by-id branch
from a 403 gate to filter-chain semantics; the option-B refactor
applied only the personal-models filter and missed the listing-path's
key/team intersection plus the BYOK scope filter that the v1 default
branch already used.

Fix mirrors the v1 default branch's filter chain on the by-id path
and adds the BYOK filter to v2:
- v1 by-id: insert _get_v1_model_info_allowed_model_names +
  _filter_v1_model_info_deployments (key/team), keep the existing
  user.models filter, append the BYOK _get_caller_byok_team_scope +
  _byok_row_outside_caller_teams filter.
- v2 default: append the same BYOK filter after user.models.

All filters use existing helpers; no new code paths introduced.
Filter (drop deployment → empty data) preserved instead of 403, per
the route_checks.py:189 spec.

Local verification: 87/87 across
- tests/test_litellm/proxy/discovery_endpoints/
- tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
- tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
songkuan-zheng added a commit that referenced this pull request Jun 19, 2026
veria-ai's review of a501b0e raised two Medium-severity findings:

1. /v2/model/info default branch did not apply the BYOK team-scope
   filter when caller's key/team/user model lists were all empty.
   Result: a virtual key with empty restrictions could see other
   teams' BYOK deployment metadata (api_base, model_info.id, etc.).

2. /v1/model/info?litellm_model_id=<id> (by-id branch) applied only
   the user.models filter after the option-B refactor, skipping
   key.models / team_models / BYOK scope. A key restricted by
   key.models with no personal model restriction could pull metadata
   for any deployment.

Finding #2 is a regression introduced by switching the by-id branch
from a 403 gate to filter-chain semantics; the option-B refactor
applied only the personal-models filter and missed the listing-path's
key/team intersection plus the BYOK scope filter that the v1 default
branch already used.

Fix mirrors the v1 default branch's filter chain on the by-id path
and adds the BYOK filter to v2:
- v1 by-id: insert _get_v1_model_info_allowed_model_names +
  _filter_v1_model_info_deployments (key/team), keep the existing
  user.models filter, append the BYOK _get_caller_byok_team_scope +
  _byok_row_outside_caller_teams filter.
- v2 default: append the same BYOK filter after user.models.

All filters use existing helpers; no new code paths introduced.
Filter (drop deployment → empty data) preserved instead of 403, per
the route_checks.py:189 spec.

Local verification: 87/87 across
- tests/test_litellm/proxy/discovery_endpoints/
- tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
- tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
songkuan-zheng added a commit that referenced this pull request Jun 23, 2026
veria-ai's review of a501b0e raised two Medium-severity findings:

1. /v2/model/info default branch did not apply the BYOK team-scope
   filter when caller's key/team/user model lists were all empty.
   Result: a virtual key with empty restrictions could see other
   teams' BYOK deployment metadata (api_base, model_info.id, etc.).

2. /v1/model/info?litellm_model_id=<id> (by-id branch) applied only
   the user.models filter after the option-B refactor, skipping
   key.models / team_models / BYOK scope. A key restricted by
   key.models with no personal model restriction could pull metadata
   for any deployment.

Finding #2 is a regression introduced by switching the by-id branch
from a 403 gate to filter-chain semantics; the option-B refactor
applied only the personal-models filter and missed the listing-path's
key/team intersection plus the BYOK scope filter that the v1 default
branch already used.

Fix mirrors the v1 default branch's filter chain on the by-id path
and adds the BYOK filter to v2:
- v1 by-id: insert _get_v1_model_info_allowed_model_names +
  _filter_v1_model_info_deployments (key/team), keep the existing
  user.models filter, append the BYOK _get_caller_byok_team_scope +
  _byok_row_outside_caller_teams filter.
- v2 default: append the same BYOK filter after user.models.

All filters use existing helpers; no new code paths introduced.
Filter (drop deployment → empty data) preserved instead of 403, per
the route_checks.py:189 spec.

Local verification: 87/87 across
- tests/test_litellm/proxy/discovery_endpoints/
- tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
- tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
songkuan-zheng added a commit that referenced this pull request Jun 26, 2026
veria-ai's review of a501b0e raised two Medium-severity findings:

1. /v2/model/info default branch did not apply the BYOK team-scope
   filter when caller's key/team/user model lists were all empty.
   Result: a virtual key with empty restrictions could see other
   teams' BYOK deployment metadata (api_base, model_info.id, etc.).

2. /v1/model/info?litellm_model_id=<id> (by-id branch) applied only
   the user.models filter after the option-B refactor, skipping
   key.models / team_models / BYOK scope. A key restricted by
   key.models with no personal model restriction could pull metadata
   for any deployment.

Finding #2 is a regression introduced by switching the by-id branch
from a 403 gate to filter-chain semantics; the option-B refactor
applied only the personal-models filter and missed the listing-path's
key/team intersection plus the BYOK scope filter that the v1 default
branch already used.

Fix mirrors the v1 default branch's filter chain on the by-id path
and adds the BYOK filter to v2:
- v1 by-id: insert _get_v1_model_info_allowed_model_names +
  _filter_v1_model_info_deployments (key/team), keep the existing
  user.models filter, append the BYOK _get_caller_byok_team_scope +
  _byok_row_outside_caller_teams filter.
- v2 default: append the same BYOK filter after user.models.

All filters use existing helpers; no new code paths introduced.
Filter (drop deployment → empty data) preserved instead of 403, per
the route_checks.py:189 spec.

Local verification: 87/87 across
- tests/test_litellm/proxy/discovery_endpoints/
- tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
- tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
songkuan-zheng added a commit that referenced this pull request Jun 26, 2026
veria-ai's review of a501b0e raised two Medium-severity findings:

1. /v2/model/info default branch did not apply the BYOK team-scope
   filter when caller's key/team/user model lists were all empty.
   Result: a virtual key with empty restrictions could see other
   teams' BYOK deployment metadata (api_base, model_info.id, etc.).

2. /v1/model/info?litellm_model_id=<id> (by-id branch) applied only
   the user.models filter after the option-B refactor, skipping
   key.models / team_models / BYOK scope. A key restricted by
   key.models with no personal model restriction could pull metadata
   for any deployment.

Finding #2 is a regression introduced by switching the by-id branch
from a 403 gate to filter-chain semantics; the option-B refactor
applied only the personal-models filter and missed the listing-path's
key/team intersection plus the BYOK scope filter that the v1 default
branch already used.

Fix mirrors the v1 default branch's filter chain on the by-id path
and adds the BYOK filter to v2:
- v1 by-id: insert _get_v1_model_info_allowed_model_names +
  _filter_v1_model_info_deployments (key/team), keep the existing
  user.models filter, append the BYOK _get_caller_byok_team_scope +
  _byok_row_outside_caller_teams filter.
- v2 default: append the same BYOK filter after user.models.

All filters use existing helpers; no new code paths introduced.
Filter (drop deployment → empty data) preserved instead of 403, per
the route_checks.py:189 spec.

Local verification: 87/87 across
- tests/test_litellm/proxy/discovery_endpoints/
- tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
- tests/test_litellm/proxy/proxy_server/test_team_model_name_translation.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant