fix(proxy): attribute org spend via team when key has no org_id

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:
         """

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,