From fa1339a1997d30dff6b503858fa0ff3b1aecbdaa Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 24 Jul 2026 11:12:38 -0700 Subject: [PATCH 1/2] fix(proxy): attribute spend to org for team-linked keys minted without org_id Keys attached to an org-linked team but minted without an organization_id produced spend that was never credited to the org: the spend writer reads user_api_key_dict.org_id with no team fallback, while the org budget check resolves the org from the team. The check therefore ran against a counter fed by almost none of the org's traffic and never tripped. Backfill org_id from the freshly fetched team object in _run_centralized_common_checks, per request only, so the spend writer and the budget check read the same org. A key with an explicitly pinned org_id always wins, and the cached key row is never mutated, so moving a team to a different org takes effect on the next auth once the team cache refreshes. --- litellm/proxy/auth/user_api_key_auth.py | 3 + .../proxy/auth/test_user_api_key_auth.py | 103 ++++++++++++++++++ 2 files changed, 106 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 709df5e64dfc..e368a1a22755 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2320,6 +2320,9 @@ async def _run_centralized_common_checks( None if isinstance(global_spend_result, BaseException) else global_spend_result ) + if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: + user_api_key_auth_obj.org_id = team_object.organization_id + # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and # master_key tokens get admin from the token; the DB row for the diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 0359d974d194..b3b540b6af8f 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3796,6 +3796,109 @@ async def test_centralized_common_checks_user_http_exception_isolates_to_user_on setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_org_id,team_org_id,expected_org_id", + [ + (None, "org-from-team", "org-from-team"), + ("org-pinned-on-key", "org-from-team", "org-pinned-on-key"), + (None, None, None), + ], +) +async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, team_org_id, expected_org_id): + """LIT-4688 regression: a key minted without an organization_id but attached + to an org-linked team must leave auth with org_id set from the team, so the + spend writer (which reads user_api_key_dict.org_id, no team fallback) + credits the org and the org budget cap can actually trip. A key with an + explicitly pinned org_id must win over the team's org.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u", team_id="t1", org_id=key_org_id) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + fetched_team = LiteLLM_TeamTableCachedObj(team_id="t1", organization_id=team_org_id) + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + org_id_seen_by_common_checks = [] + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=fetched_team, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + side_effect=lambda **kw: org_id_seen_by_common_checks.append(kw["valid_token"].org_id), + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.org_id == expected_org_id + assert org_id_seen_by_common_checks == [expected_org_id] + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_centralized_common_checks_org_backfill_survives_team_fetch_failure(): + """When the team DB fetch fails, the token-derived fallback team carries no + organization_id, so the backfill must leave org_id as None rather than + crash or mis-attribute.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u", team_id="t1") + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + side_effect=Exception("DB down"), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ) as mock_checks, + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + mock_checks.assert_awaited_once() + assert token.org_id is None + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_master_key_auth_substitutes_alias_for_api_key(): """ From 22607bf801488d4b6f3a86c77412b0d66ee7e50b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 24 Jul 2026 16:53:35 -0700 Subject: [PATCH 2/2] test(proxy): cover CLI session-token org backfill from team CLI session tokens from /sso/cli/poll are minted with a real team_id but no org_id, and their auth path decrypts the blob without the combined_view team join that fills org for DB keys. Spend from these tokens reached the team but never the org, so org budgets never tripped. The regression test mints a real CLI token, runs it through the centralized checks, and asserts the credential leaves auth with the team's org. --- .../proxy/auth/test_user_api_key_auth.py | 58 +++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index b3b540b6af8f..c0c91ae103f8 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3856,6 +3856,64 @@ async def test_centralized_common_checks_backfills_org_id_from_team(key_org_id, setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_cli_session_token_org_backfilled_from_team(monkeypatch): + """LIT-4688 root cause: CLI session tokens (from /sso/cli/poll) are minted + with a real team_id but no org_id, and their auth path decrypts the blob + without the combined_view team join, so their spend never reached the org. + The centralized-checks backfill must complete the credential from the team + the same way the SQL view does for DB keys.""" + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-lit4688") + + cli_user = LiteLLM_UserTable(user_id="cli-user", user_role="internal_user", teams=["t-cli"], models=[]) + blob = ExperimentalUIJWTToken.get_cli_jwt_auth_token(user_info=cli_user, team_id="t-cli", team_alias="cli-team") + token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(blob) + assert token is not None + assert token.is_session_token is True + assert token.team_id == "t-cli" + assert token.org_id is None + + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + org_linked_team = LiteLLM_TeamTableCachedObj(team_id="t-cli", organization_id="org-infoops") + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.get_team_object", + new_callable=AsyncMock, + return_value=org_linked_team, + ), + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + + assert token.org_id == "org-infoops" + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + @pytest.mark.asyncio async def test_centralized_common_checks_org_backfill_survives_team_fetch_failure(): """When the team DB fetch fails, the token-derived fallback team carries no