feat(proxy): attribute usage to every team a user belongs to (opt-in) - #37809
Draft
Hardell wants to merge 3 commits into
Draft
feat(proxy): attribute usage to every team a user belongs to (opt-in)#37809Hardell wants to merge 3 commits into
Hardell wants to merge 3 commits into
Conversation
Today a request is attributed to the single team stamped on the virtual key (or resolved from one JWT claim) and to that team's organization. A user who belongs to many teams contributes spend to whichever team the key happens to name and nothing to the rest, so "what did this team consume?" only has an answer for keys that name it. Add two opt-in general_settings, both defaulting to off: - track_spend_across_all_user_teams: spend increments, daily rollups, and budget gates apply to every team the caller belongs to, and to every organization reached through those teams. - enforce_rate_limits_across_all_user_teams: the same expansion for the RPM/TPM limiter, so a request must fit inside every membership's limit. They are separate settings because they carry different costs. Spend attribution is additive bookkeeping. Rate-limit expansion makes the caller's effective limit the minimum across their memberships, which is a live behavior change for anyone in a busy team. Memberships come from LiteLLM_UserTable.teams, already maintained by SCIM and JWT sync. Resolution happens once in the auth path and is carried on server-only UserAPIKeyAuth fields, stripped from validated input like mcp_source_team_rpm_limits so a caller cannot choose which buckets they are charged against. No migration. LiteLLM_DailyTeamSpend is already unique per (team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint), and the spend queue aggregates by entity_type:entity_id, so N teams means N rows and N counters. Also make the rate limiter's per-hash-tag Lua calls concurrent. They were sequential, so on Redis Cluster one descriptor per team would have meant one round trip per team on the hot path. Non-cluster Redis was and remains a single call. Known limitation: organization rate limits still apply to the stamped org only. Multi-org spend attribution works; multi-org rate limiting would need per-org limits precomputed in auth and is left out deliberately to keep this change reviewable.
|
mzolota seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Contributor
added 2 commits
August 21, 2026 12:29
- Import attribution_targets in proxy_server. increment_spend_counters referenced it without an import, raising NameError on every spend increment. Caught by proxy-server and proxy-infra shards; my local runs never exercised that function. - Run ruff format over auth_checks.py. - Regenerate the two ConfigGeneralSettings entries in schema.d.ts. - Widen a test stub in test_spend_tracking_utils.py whose fixed signature predates the two new update_database parameters. Also close a real test gap the coverage report exposed: the two new budget gates and the request-metadata stamping had no tests at all. Adds nine, covering an over-budget non-stamped team, an over-budget non-stamped org, the stamped team being skipped so it cannot raise twice, an unloadable team contributing no ceiling, and both settings-off paths. Adds two counter tests pinning that spend fans out to every attributed team and that it does not when the setting is off.
The strict gate ratchets per-rule totals against ruff-strict-budget.json, and this change had pushed two of them up by one each: - ANN401/TID251: membership_attribution.py used typing.Any for the otel span and the team-object tuples. Both now carry concrete types (Span, LiteLLM_TeamTableCachedObj) imported under TYPE_CHECKING, so the module no longer references Any at all. - ANN202: the two new budget-check functions had no return annotation. Also moves a return out of a try block into an else (TRY300). Verified per-file against the base commit across every rule the strict config selects, not just the two the gate reported: zero new violations in all nine changed files. scripts/ruff_strict_gate.py itself needs fcntl and cannot run on Windows, so this was checked by counting each rule in each changed file at the base blob and at HEAD and diffing.
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
2 participants
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.
TLDR
Problem this solves:
How it solves it:
general_settings, both default offUser Flow
Before: a platform admin who provisions teams from an IdP cannot tell what any team consumed, because a developer's usage only ever lands on one team.
POST https://litellm-domain/v1/chat/completionswith their key and gets a normal200with a completion.https://litellm-domain/ui/?page=teams. One team shows the spend. The other four show$0.After: the same request charges every team the developer belongs to, so each team's spend and budget become real.
track_spend_across_all_user_teams: trueingeneral_settingsand restarts the proxy.POST https://litellm-domain/v1/chat/completionswith the same key and gets the same200with a completion.https://litellm-domain/ui/?page=teams. All five teams now show that request's spend.Relevant issues
Linear ticket
Pre-Submission checklist
Test files run locally, all passing:
tests/test_litellm/proxy/auth/test_membership_attribution.py(21 new tests)tests/test_litellm/proxy/db/test_db_spend_update_writer.py(69)tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py(47)tests/test_litellm/proxy/auth/test_multi_budget_windows.pyruff check litellm/proxy/andruff format --checkcleanMost of the new tests are OFF-case regression guards, since both settings default to off and the untouched default path is the thing most worth protecting.
Screenshots / Proof Of Fix
Outstanding — this is why the PR is a draft.
I have not yet run the end-to-end proof this template requires: a live proxy, real provider calls, and the before/after team spend captured at the merge base and at the PR tip. I did not want to open a PR claiming e2e verification I have not performed.
What is verified so far is unit-level only, listed in the checklist above.
I will add the full Before/After section against a live proxy before marking this ready for review. If maintainers would rather see a specific shape of proof — particular endpoints, or a particular team/org fixture — say so and I will capture that instead.
Type
🆕 New Feature
Caveats (if any)
Notes on two of those
The adjacent Redis fix.
_execute_lua_script_by_hash_tagawaited one round trip per hash-tag group in sequence. On non-cluster Redis there is a single group, so it never showed. On Redis Cluster the groups are per-slot, and one descriptor per team would have meant one round trip per team on the hot path. The calls now run concurrently, results still consumed in group order. It is correct independently of this feature — happy to split it into its own PR if you would prefer that for scope.Relationship to the existing MCP stance.
_admitted_subject_team_rpm_limitsdeliberately avoids charging several team buckets for one call, so a cross-team user cannot drain several teams' shared buckets. That reasoning holds when the stamped team is the real owner of the call and the others are incidental. Under membership attribution there is no primary team — the caller's activity genuinely belongs to all their teams, which is the premise of the feature — so charging all of them is the consistent choice. The default stays off, so today's behavior remains the default everywhere.Implementation notes
Memberships come from
LiteLLM_UserTable.teams, already maintained by SCIM and JWT sync, so there is no new source of truth. Resolution happens once per request in the auth path and is carried on server-onlyUserAPIKeyAuthfields withexclude=True, popped in thecheck_api_keyvalidator exactly likemcp_source_team_rpm_limits— a caller who could set them would pick their own budget and rate-limit buckets. There is a test for that.Every consumer reads through one helper that falls back to the stamped id, so with both settings off the resolved list is exactly
[team_id].Resolution fails open: a team that cannot be loaded is skipped, not raised. Attribution is bookkeeping layered on an authorization decision already made, and it must not turn an authorized request into a 500.
Attributed team rate-limit descriptors reuse
key="team", so a team shares one bucket whether it is stamped on the key or merely a membership.should_rate_limitalready rejects when any descriptor is over limit, so the limiter engine itself is unchanged.No migration:
LiteLLM_DailyTeamSpendis already unique on(team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint), and the spend queue aggregates byentity_type:entity_id.LiteLLM_SpendLogskeeps naming the stamped team — fanning out per-request log rows would multiply the highest-volume table, so the daily rollups carry the multi-team truth instead.Final Attestation