diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index b4a4fd571d03..8618a6ca382c 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,51 @@ 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 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 + + 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..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 @@ -1042,6 +1042,178 @@ 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_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. 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()), + 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, + 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" + ) + + 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 +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,