From d4ca6c859befa8fd0932b9bae322e4bb2f3d594e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 19 Jun 2026 01:43:05 +0000 Subject: [PATCH 1/3] fix(proxy): attribute org spend via team when key has no org_id Co-authored-by: Mateo Wang --- .../proxy/hooks/proxy_track_cost_callback.py | 48 ++++++- .../hooks/test_proxy_track_cost_callback.py | 124 ++++++++++++++++++ 2 files changed, 171 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b4a4fd571d03..ce722046cd02 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -205,7 +205,10 @@ async def _PROXY_track_cost_callback( ) user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None)) team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None)) - org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None)) + org_id = await self._resolve_spend_tracking_org_id( + org_id=cast(Optional[str], metadata.get("user_api_key_org_id", None)), + team_id=team_id, + ) key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None)) end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) sl_object: Optional[StandardLoggingPayload] = kwargs.get( @@ -333,6 +336,49 @@ async def _PROXY_track_cost_callback( spend_log_error("Error in tracking cost callback - %s", str(e), exc=e) + @staticmethod + async def _resolve_spend_tracking_org_id( + org_id: Optional[str], team_id: Optional[str] + ) -> Optional[str]: + """Resolve the organization a request's spend rolls up to. + + A key created under a team that belongs to an organization carries no + organization_id of its own, so without this fallback its spend never + reaches the org's spend column or daily org aggregate. Org budgets then + fail to enforce once the in-memory reservation counter expires or across + workers that don't share it, because the persisted floor stays at zero. + This mirrors the auth-time fallback in + auth_checks._organization_max_budget_check so spend attribution and + budget enforcement agree on the same org. + """ + if org_id is not None or team_id is None: + return org_id + + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + if prisma_client is None: + return org_id + + try: + team_object = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception: + verbose_proxy_logger.debug( + "Failed to resolve org_id from team_id=%s for spend tracking", + team_id, + ) + return org_id + + return team_object.organization_id + @staticmethod async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict: """ diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 771e10a54a06..9c0507267571 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1042,6 +1042,130 @@ async def test_failure_hook_keeps_error_information_traceback_by_default(monkeyp assert error_information["traceback"], "expected a non-empty traceback by default" +@pytest.mark.asyncio +async def test_resolve_spend_tracking_org_id_falls_back_to_team_org(): + """A key with no org_id of its own must attribute spend to its team's org, + so org spend/budget tracking sees the usage. Mirrors the auth-time fallback + in auth_checks._organization_max_budget_check.""" + mock_team_obj = MagicMock() + mock_team_obj.organization_id = "org-from-team" + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ) as mock_get_team, + ): + resolved = await _ProxyDBLogger._resolve_spend_tracking_org_id( + org_id=None, team_id="team-123" + ) + + assert resolved == "org-from-team" + mock_get_team.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_resolve_spend_tracking_org_id_prefers_explicit_org_id(): + """When the key already carries an org_id, keep it and skip the team lookup.""" + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + resolved = await _ProxyDBLogger._resolve_spend_tracking_org_id( + org_id="explicit-org", team_id="team-123" + ) + + assert resolved == "explicit-org" + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_resolve_spend_tracking_org_id_no_team_returns_none(): + """No org_id and no team_id means there is nothing to attribute org spend to.""" + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + resolved = await _ProxyDBLogger._resolve_spend_tracking_org_id( + org_id=None, team_id=None + ) + + assert resolved is None + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_track_cost_callback_attributes_org_spend_via_team(): + """End-to-end through the cost callback: a team key with no org_id of its own + must still record spend against the team's organization. Before the fallback, + org_id stayed None and org spend/budget were never tracked.""" + logger = _ProxyDBLogger() + + mock_team_obj = MagicMock() + mock_team_obj.organization_id = "org-from-team" + + kwargs = { + "model": "gpt-4", + "litellm_params": { + "metadata": { + "user_api_key": "hashed_key", + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-123", + # no user_api_key_org_id - the key inherits org via its team + }, + }, + "standard_logging_object": { + "response_cost": 0.5, + "request_tags": None, + }, + "stream": False, + } + + mock_proxy_logging = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + mock_increment_spend_counters = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), + patch( + "litellm.proxy.proxy_server.increment_spend_counters", + mock_increment_spend_counters, + ), + patch( + "litellm.proxy.proxy_server.update_cache", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ), + ): + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + assert ( + mock_proxy_logging.db_spend_update_writer.update_database.call_args.kwargs[ + "org_id" + ] + == "org-from-team" + ) + mock_increment_spend_counters.assert_awaited_once() + assert mock_increment_spend_counters.call_args.kwargs["org_id"] == "org-from-team" + + @pytest.mark.asyncio async def test_failure_hook_drops_error_information_traceback_when_env_set( monkeypatch, From 0bb265720a3a7e14af6fcd26a14bd24d800352b3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 19 Jun 2026 01:57:42 +0000 Subject: [PATCH 2/3] test(proxy): cover org-id resolution fallbacks for codecov patch coverage Add cases for the no-db-client early return and the swallowed team-lookup failure in _resolve_spend_tracking_org_id so every new line in the patch is exercised, satisfying codecov's 100% patch coverage gate. --- .../hooks/test_proxy_track_cost_callback.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 9c0507267571..d4058841b30b 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1098,6 +1098,47 @@ async def test_resolve_spend_tracking_org_id_no_team_returns_none(): mock_get_team.assert_not_called() +@pytest.mark.asyncio +async def test_resolve_spend_tracking_org_id_skips_lookup_without_db(): + """With no DB client we cannot resolve the team's org, so return the original + org_id without attempting a lookup rather than crashing cost tracking.""" + with ( + patch("litellm.proxy.proxy_server.prisma_client", None), + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team, + ): + resolved = await _ProxyDBLogger._resolve_spend_tracking_org_id( + org_id=None, team_id="team-123" + ) + + assert resolved is None + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_resolve_spend_tracking_org_id_swallows_team_lookup_failure(): + """A failing team lookup must not break cost tracking; fall back to the + original org_id instead of letting the exception propagate.""" + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + side_effect=Exception("team not found"), + ) as mock_get_team, + ): + resolved = await _ProxyDBLogger._resolve_spend_tracking_org_id( + org_id=None, team_id="team-123" + ) + + assert resolved is None + mock_get_team.assert_awaited_once() + + @pytest.mark.asyncio async def test_track_cost_callback_attributes_org_spend_via_team(): """End-to-end through the cost callback: a team key with no org_id of its own From 51e7fd4294e86305b6e5ad7f32bedf31659bd6ae Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 19 Jun 2026 02:14:37 +0000 Subject: [PATCH 3/3] fix(proxy): log team-lookup failures in org spend resolution at warning level A team lookup failure during spend tracking silently dropped org attribution. At debug level this gap is invisible in production, so an unenforced org budget has no signal to investigate. Log at warning with the impact, and assert the warning in the regression test so it can't quietly fall back to debug. --- litellm/proxy/hooks/proxy_track_cost_callback.py | 8 +++++--- .../proxy/hooks/test_proxy_track_cost_callback.py | 9 ++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index ce722046cd02..8618a6ca382c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -370,10 +370,12 @@ async def _resolve_spend_tracking_org_id( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - except Exception: - verbose_proxy_logger.debug( - "Failed to resolve org_id from team_id=%s for spend tracking", + except Exception as e: + verbose_proxy_logger.warning( + "Failed to resolve org_id from team_id=%s for spend tracking; " + "org spend for this request will be unattributed: %s", team_id, + e, ) return org_id diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index d4058841b30b..a4f3394c5b3b 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1120,7 +1120,9 @@ async def test_resolve_spend_tracking_org_id_skips_lookup_without_db(): @pytest.mark.asyncio async def test_resolve_spend_tracking_org_id_swallows_team_lookup_failure(): """A failing team lookup must not break cost tracking; fall back to the - original org_id instead of letting the exception propagate.""" + original org_id instead of letting the exception propagate. The failure is + a silent org-attribution gap in production, so it must surface at warning + level (not debug) for operators to notice.""" with ( patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), @@ -1130,6 +1132,9 @@ async def test_resolve_spend_tracking_org_id_swallows_team_lookup_failure(): new_callable=AsyncMock, side_effect=Exception("team not found"), ) as mock_get_team, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.verbose_proxy_logger.warning" + ) as mock_warning, ): resolved = await _ProxyDBLogger._resolve_spend_tracking_org_id( org_id=None, team_id="team-123" @@ -1137,6 +1142,8 @@ async def test_resolve_spend_tracking_org_id_swallows_team_lookup_failure(): assert resolved is None mock_get_team.assert_awaited_once() + mock_warning.assert_called_once() + assert "team-123" in mock_warning.call_args.args @pytest.mark.asyncio