Skip to content

feat(spend): warn when spend-attribution metadata diverges from the resolved key - #31005

Open
yassin-berriai wants to merge 1 commit into
litellm_internal_stagingfrom
litellm_principal_identity_divergence
Open

feat(spend): warn when spend-attribution metadata diverges from the resolved key#31005
yassin-berriai wants to merge 1 commit into
litellm_internal_stagingfrom
litellm_principal_identity_divergence

Conversation

@yassin-berriai

@yassin-berriai yassin-berriai commented Jun 22, 2026

Copy link
Copy Markdown
Contributor

Relevant issues

Linear ticket

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

  • I have added meaningful tests
  • My PR passes all unit tests on make test-unit
  • My PR's scope is as isolated as possible; it only solves 1 specific problem
  • I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

CI (LiteLLM team)

  • Branch creation CI run
    Link:

  • CI run for the last commit
    Link:

  • Merge / cherry-pick CI run
    Links:

Screenshots / Proof of Fix

Verified against a live proxy on localhost:4000 backed by Postgres, hitting the real Anthropic API, running this branch with PYTHONPATH pinned to the worktree (the divergence module was confirmed loaded from the worktree, not the main checkout).

The guard sits in the failure-path fallback _enrich_failure_metadata_with_key_info, so it only evaluates when a failed request reaches spend logging. On a team-scoped key carrying a user, team, and org, one successful completion and one deliberately failing call (a model the key is not allowed to use, a real 403) both write spend rows with the correct identity:

$ psql -c 'select substring(request_id,1,18) request_id, "user" user_id, team_id, organization_id, status
           from "LiteLLM_SpendLogs" order by "startTime" desc limit 2;'

 19ab44a9-5e7d-48ab | a745550f-...-dd7e39faea53 | c1014da1-...-98ef76fb62b8 | 78153153-...-c7d04a35b0b4 | failure
 chatcmpl-36b51980- | a745550f-...-dd7e39faea53 | c1014da1-...-98ef76fb62b8 | 78153153-...-c7d04a35b0b4 | success

The failure row carries the resolved user, team, and org, so the enrichment fallback found identity already present and the guard correctly stayed silent: no Spend attribution metadata diverges from the resolved key identity warnings and no guard errors in the proxy log (the only traceback is the expected key_model_access_denied ProxyException from the 403). This shows the guard is non-disruptive on healthy traffic and that failure-path attribution is intact. The firing behavior, where the consumed metadata identity is missing or different from the freshly fetched key (the empty-metadata misattribution case the design cites), is pinned by the unit tests in tests/test_litellm/proxy/auth/test_resolvers_divergence.py: an empty consumed identity against a populated key warns naming exactly user_id, team_id, and org_id, and a matching identity stays silent

Type

🆕 New Feature

Changes

Phase 1 of the internal Caller Identity design ("Resolve Once, Consume Everywhere"); the first consumer of the Phase 0 resolved identity work (PR #30887)

The design's failure table calls out spend being mis- or un-attributed when the flattened user_api_key_* request metadata is empty or stale, which the spend path silently papers over by falling back to a key lookup. That fallback lives in _enrich_failure_metadata_with_key_info; it fetches the key object and fills metadata["user_api_key_user_id"/"team_id"/"org_id"] whenever they are None, with no signal that the attribution was missing in the first place. This change turns that silent fill into an observable warning so the divergence surfaces before deploy instead of hiding in the spend logs

litellm/proxy/auth/resolvers/divergence.py holds a pure, fully typed comparator. spend_identity_divergence compares a resolved SpendIdentity against the consumed one and returns the fields that are missing-in-consumed or hold a different value (user_id, team_id, org_id) as an immutable tuple of frozen records. resolved_identity_from_key projects the authoritative identity off the fetched key object, SpendIdentity is a frozen model so the untyped metadata values are validated at the boundary, and log_identity_divergence emits a single structured verbose_proxy_logger.warning naming the diverging fields (resolved value vs consumed value) keyed by the hashed token, and is a no-op when nothing diverges

The guard is wired into _enrich_failure_metadata_with_key_info right after the key object is fetched and before the existing fill runs, comparing the pre-fill metadata against the key. The existing fill behavior is unchanged; this only adds the warning, and it is wrapped in its own try/except so a guard failure can never affect the enrichment or cost tracking. The change is additive and zero-behavior-change; the only externally observable effect is a warning log on divergence

Tests in tests/test_litellm/proxy/auth/test_resolvers_divergence.py cover matching identity (silent), empty consumed metadata against a populated key naming exactly the missing fields (the misattribution case), and a consumed user_id that differs from the key. The enrichment fallback is not cleanly injectable (get_key_object and the proxy-server globals it reads are module-level rather than injected), so the warning is asserted through the thin log_identity_divergence wrapper rather than by mocking that dependency

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@codecov

codecov Bot commented Jun 22, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.54839% with 2 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
litellm/proxy/hooks/proxy_track_cost_callback.py 60.00% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

@yassin-berriai yassin-berriai changed the title feat(auth): warn when downstream identity diverges from the resolved Principal feat(spend): warn when spend-attribution identity diverges from the resolved key Jun 22, 2026
@yassin-berriai
yassin-berriai force-pushed the litellm_principal_identity_divergence branch from 7a3305e to 10e98ac Compare June 22, 2026 20:12
@yassin-berriai yassin-berriai changed the title feat(spend): warn when spend-attribution identity diverges from the resolved key feat(spend): warn when spend-attribution metadata diverges from the resolved key Jun 22, 2026
@yassin-berriai
yassin-berriai force-pushed the litellm_principal_identity_divergence branch 2 times, most recently from fedbe0e to c9cef3c Compare June 22, 2026 20:17
@yassin-berriai
yassin-berriai force-pushed the litellm_principal_identity_divergence branch from c9cef3c to 9933c50 Compare June 22, 2026 20:20
@yassin-berriai
yassin-berriai marked this pull request as ready for review June 22, 2026 20:24
@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai

@greptile-apps

greptile-apps Bot commented Jun 22, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds Phase 1 of a "Resolve Once, Consume Everywhere" identity design: a divergence guard that emits a structured warning whenever the user_api_key_* metadata fields carried into spend logging differ from what the freshly-fetched key object actually holds. The change is purely additive — the existing fill logic in _enrich_failure_metadata_with_key_info is untouched, and the guard is wrapped in its own try/except so any failure is non-fatal.

  • litellm/proxy/auth/resolvers/divergence.py introduces a self-contained, fully-typed comparator (SpendIdentity, FieldDivergence, spend_identity_divergence, log_identity_divergence) with no proxy globals or DB calls of its own.
  • proxy_track_cost_callback.py wires the guard into the failure-path enrichment function immediately after get_key_object returns and before the fill; test coverage is provided through the thin log_identity_divergence wrapper since the enrichment path's module-level dependencies are not injectable.

Confidence Score: 5/5

Safe to merge — the change is additive and zero-behavior-changing; the only externally visible effect is a new warning log on divergence, and the guard is wrapped so it can never disrupt spend tracking.

All three files contain self-contained, well-isolated additions. The divergence comparator is a pure function with no DB access or global state. The wiring in the enrichment hook is inside its own try/except and does not alter any existing fill logic. Unit tests are mock-only and cover the three critical paths.

No files require special attention.

Important Files Changed

Filename Overview
litellm/proxy/auth/resolvers/divergence.py New pure module introducing SpendIdentity, FieldDivergence, spend_identity_divergence, and log_identity_divergence; correctly frozen/typed, no side-effects beyond the warning log.
litellm/proxy/hooks/proxy_track_cost_callback.py Wires the divergence guard into _enrich_failure_metadata_with_key_info before the existing fill; guarded by its own try/except so any failure is non-fatal; existing fill logic is unchanged.
tests/test_litellm/proxy/auth/test_resolvers_divergence.py New unit tests (no network calls) covering silent match, all-None consumed vs populated resolved, and single-field user_id divergence; all mock-only as required by repo rules.

Reviews (2): Last reviewed commit: "feat(spend): warn when spend-attribution..." | Re-trigger Greptile

Comment on lines +63 to +68
verbose_proxy_logger.warning(
"Spend attribution metadata diverges from the resolved key identity "
"[credential_ref=%s]: %s",
resolved_key.token,
fields,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 WARNING volume for valid-key 401 failures

For auth errors where the bearer token resolves to a real key in the DB but the request failed before UserAPIKeyAuth fields were populated (e.g., expired key, revoked key that still exists, early rate-limit checks), all three metadata fields will be None while the resolved key_obj carries actual user_id/team_id/org_id. This fires a WARNING for every such request. In deployments with many invalid auth attempts against valid-but-revoked keys the log can become very noisy. Consider adding a counter metric or sampling, or downgrading to DEBUG for the "all-fields-None-consumed vs populated-resolved" case (i.e., total metadata absence) vs the "some fields present but wrong" case (genuine staleness).

@yassin-berriai

Copy link
Copy Markdown
Contributor Author

Thanks for the review. On the gate and the log-volume points, both are deliberate and narrower than they look.

The alias-populated path is intentionally not covered because there is nothing independent to compare against there. The key object (the authoritative reference) is only fetched inside the alias-None branch. In the alias-populated post-auth-failure path the consumed user_api_key_* metadata was itself built from the same resolved UserAPIKeyAuth in async_post_call_failure_hook, so comparing the two would be tautological. The only place an independent identity exists without a new lookup is the empty-metadata branch, where a fresh key fetch can genuinely disagree with what the request carried, which is exactly where the guard sits. Extending it to the alias-populated path would require a second key fetch purely for the guard, the redundant identity lookup this initiative is removing. Fetch-free coverage of every failure path arrives in Phase 2, when identity is resolved once at the seam and carried on the request.

On volume: the warning is a no-op unless the freshly fetched key actually disagrees with the metadata, and it only compares user_id, team_id, and org_id. The common fallback case is a post-auth failure where only team_alias is missing from the SQL view while user/team/org are present; there the guard finds no divergence and stays silent. It fires only when user/team/org identity itself was missing or different, which is the real misattribution this is meant to surface, not per-request noise

@yassin-berriai

Copy link
Copy Markdown
Contributor Author

@greptileai

@greptile-apps

greptile-apps Bot commented Jun 22, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR introduces Phase 1 of the "Resolve Once, Consume Everywhere" design: a divergence guard that emits a structured WARNING log whenever the flattened user_api_key_* spend metadata disagrees with the freshly-fetched key identity, turning the previously silent _enrich_failure_metadata_with_key_info fill into an observable signal.

  • litellm/proxy/auth/resolvers/divergence.py adds a pure, frozen comparator (SpendIdentity, FieldDivergence, spend_identity_divergence) and a thin logging wrapper (log_identity_divergence) that names each diverging field with its resolved and consumed values.
  • The guard is inserted into _enrich_failure_metadata_with_key_info (failure-path spend logging only) right after the key is fetched, wrapped in its own try/except so it is strictly non-fatal, and the existing metadata-fill behavior is unmodified.
  • Three targeted unit tests cover the matching (silent), all-absent consumed metadata, and differing single-field cases without any network calls.

Confidence Score: 4/5

The change is purely additive: new module, new import, new try/except-wrapped log call. The existing fill path is untouched, and a guard failure cannot affect cost tracking.

The warning emitted by log_identity_divergence fires whenever consumed fields are None and the resolved key has them populated — which is the normal state for every 401 auth error, not just genuine misattributions. High-auth-failure deployments will see constant three-field warnings on each failed unauthenticated request, making the signal harder to act on in production. Everything else (comparison logic, wiring, non-fatality, tests) looks correct.

litellm/proxy/auth/resolvers/divergence.py — specifically the log level decision in log_identity_divergence for the all-None consumed case

Important Files Changed

Filename Overview
litellm/proxy/auth/resolvers/divergence.py New pure comparator module: SpendIdentity model, FieldDivergence dataclass, and log_identity_divergence helper. Logic is correct; warning fires for both absent and wrong consumed fields including all-None cases from 401 errors.
litellm/proxy/hooks/proxy_track_cost_callback.py Divergence guard wired into _enrich_failure_metadata_with_key_info after key lookup; wrapped in its own try/except so it is non-fatal. Existing fill behavior is unchanged.
tests/test_litellm/proxy/auth/test_resolvers_divergence.py Three unit tests covering: matching identity (silent), empty consumed metadata (warns all three fields), differing user_id (warns correctly). No network calls; consistent with existing test patterns in this directory.

Reviews (3): Last reviewed commit: "feat(spend): warn when spend-attribution..." | Re-trigger Greptile

Comment on lines +54 to +68
divergences = spend_identity_divergence(
resolved_identity_from_key(resolved_key), consumed
)
if not divergences:
return
fields = ", ".join(
f"{d.field} (resolved={d.resolved_value!r}, consumed={d.consumed_value!r})"
for d in divergences
)
verbose_proxy_logger.warning(
"Spend attribution metadata diverges from the resolved key identity "
"[credential_ref=%s]: %s",
resolved_key.token,
fields,
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Warning fires on every 401 auth error, not just genuine misattributions

In scenario 1 from the docstring (auth errors where UserAPIKeyAuth is created with only api_key set), the consumed metadata will always have all three identity fields as None while the looked-up key has them populated. This means every 401 auth failure will emit a 3-field divergence warning — even when nothing is actually wrong with spend attribution.

There is currently no way to distinguish "metadata was absent because request never authenticated" (expected) from "metadata holds a stale/wrong identity" (the actual misattribution signal this guard is designed to surface). High-auth-failure deployments will see constant noise, making the signal hard to act on. Consider logging at DEBUG for the all-None consumed case and reserving WARNING for when consumed_value is a non-None value that differs from resolved_value.

@yassin-berriai
yassin-berriai enabled auto-merge (squash) June 26, 2026 19:59
@yassin-berriai
yassin-berriai disabled auto-merge June 26, 2026 19:59
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.

2 participants