feat: enforce x-litellm-trace-id in header, if required

litellm/proxy/_types.py

@@ -1,59 +1,40 @@
 import enum
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
+from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal,
+                    Optional, Union)
 
 import httpx
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-    Json,
-    field_validator,
-    model_validator,
-)
+from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator,
+                      model_validator)
 from typing_extensions import Required, TypedDict
 
 from litellm._uuid import uuid
 from litellm.types.integrations.slack_alerting import AlertType
-from litellm.types.llms.openai import (
-    AllMessageValues,
-    OpenAIFileObject,
-    ResponsesAPIResponse,
-)
-from litellm.types.mcp import (
-    MCPAuthType,
-    MCPCredentials,
-    MCPTransport,
-    MCPTransportType,
-)
+from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject,
+                                       ResponsesAPIResponse)
+from litellm.types.mcp import (MCPAuthType, MCPCredentials, MCPTransport,
+                               MCPTransportType)
 from litellm.types.mcp_server.mcp_server_manager import MCPInfo
 from litellm.types.router import RouterErrors, UpdateRouterConfig
 from litellm.types.secret_managers.main import KeyManagementSystem
-from litellm.types.utils import (
-    CallTypes,
-    CostBreakdown,
-    EmbeddingResponse,
-    GenericBudgetConfigType,
-    ImageResponse,
-    LiteLLMBatch,
-    LiteLLMFineTuningJob,
-    LiteLLMPydanticObjectBase,
-    ModelResponse,
-    ProviderField,
-    StandardCallbackDynamicParams,
-    StandardLoggingGuardrailInformation,
-    StandardLoggingMCPToolCall,
-    StandardLoggingModelInformation,
-    StandardLoggingPayloadErrorInformation,
-    StandardLoggingPayloadStatus,
-    StandardLoggingVectorStoreRequest,
-    StandardPassThroughResponseObject,
-    TextCompletionResponse,
-)
+from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse,
+                                 GenericBudgetConfigType, ImageResponse,
+                                 LiteLLMBatch, LiteLLMFineTuningJob,
+                                 LiteLLMPydanticObjectBase, ModelResponse,
+                                 ProviderField, StandardCallbackDynamicParams,
+                                 StandardLoggingGuardrailInformation,
+                                 StandardLoggingMCPToolCall,
+                                 StandardLoggingModelInformation,
+                                 StandardLoggingPayloadErrorInformation,
+                                 StandardLoggingPayloadStatus,
+                                 StandardLoggingVectorStoreRequest,
+                                 StandardPassThroughResponseObject,
+                                 TextCompletionResponse)
 from litellm.types.videos.main import VideoObject
 
-from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type
+from .types_utils.utils import (get_instance_fn,
+                                validate_custom_validate_return_type)
 
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
@@ -201,6 +182,7 @@ class Litellm_EntityType(enum.Enum):
     ORGANIZATION = "organization"
     PROJECT = "project"
     TAG = "tag"
+    AGENT = "agent"
 
     # global proxy level entity
     PROXY = "proxy"
@@ -2399,7 +2381,8 @@ class UserAPIKeyAuth(
 
         This is used to track number of requests/spend for health check calls.
         """
-        from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
+        from litellm.constants import \
+            LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
 
         return cls(
             api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
@@ -2431,7 +2414,8 @@ class UserAPIKeyAuth(
 
         This is used to track actions performed by automated system jobs.
         """
-        from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
+        from litellm.constants import \
+            LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
 
         return cls(
             api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
@@ -2837,7 +2821,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase):
 
     @model_validator(mode="after")
     def mask_api_keys(self):
-        from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
+        from litellm.litellm_core_utils.sensitive_data_masker import \
+            SensitiveDataMasker
 
         masker = SensitiveDataMasker(sensitive_patterns={"key"})
 
@@ -4205,6 +4190,7 @@ class DBSpendUpdateTransactions(TypedDict):
     team_member_list_transactions: Optional[Dict[str, float]]
     org_list_transactions: Optional[Dict[str, float]]
     tag_list_transactions: Optional[Dict[str, float]]
+    agent_list_transactions: Optional[Dict[str, float]]
 
 
 class SpendUpdateQueueItem(TypedDict, total=False):

litellm/proxy/auth/auth_checks.py

@@ -11,7 +11,8 @@ Run checks for:
 import asyncio
 import re
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
+from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union,
+                    cast)
 
 from fastapi import HTTPException, Request, status
 from pydantic import BaseModel
@@ -20,42 +21,29 @@ import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
 from litellm.caching.dual_cache import LimitedSizeOrderedDict
-from litellm.constants import (
-    CLI_JWT_EXPIRATION_HOURS,
-    CLI_JWT_TOKEN_NAME,
-    DEFAULT_ACCESS_GROUP_CACHE_TTL,
-    DEFAULT_IN_MEMORY_TTL,
-    DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
-    DEFAULT_MAX_RECURSE_DEPTH,
-    EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
-)
+from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME,
+                               DEFAULT_ACCESS_GROUP_CACHE_TTL,
+                               DEFAULT_IN_MEMORY_TTL,
+                               DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
+                               DEFAULT_MAX_RECURSE_DEPTH,
+                               EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE)
 from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
-from litellm.proxy._types import (
-    RBAC_ROLES,
-    CallInfo,
-    LiteLLM_AccessGroupTable,
-    LiteLLM_BudgetTable,
-    LiteLLM_EndUserTable,
-    Litellm_EntityType,
-    LiteLLM_JWTAuth,
-    LiteLLM_ObjectPermissionTable,
-    LiteLLM_OrganizationMembershipTable,
-    LiteLLM_OrganizationTable,
-    LiteLLM_ProjectTableCachedObj,
-    LiteLLM_TagTable,
-    LiteLLM_TeamMembership,
-    LiteLLM_TeamTable,
-    LiteLLM_TeamTableCachedObj,
-    LiteLLM_UserTable,
-    LiteLLMRoutes,
-    LitellmUserRoles,
-    NewTeamRequest,
-    ProxyErrorTypes,
-    ProxyException,
-    RoleBasedPermissions,
-    SpecialModelNames,
-    UserAPIKeyAuth,
-)
+from litellm.proxy._types import (RBAC_ROLES, CallInfo,
+                                  LiteLLM_AccessGroupTable,
+                                  LiteLLM_BudgetTable, LiteLLM_EndUserTable,
+                                  Litellm_EntityType, LiteLLM_JWTAuth,
+                                  LiteLLM_ObjectPermissionTable,
+                                  LiteLLM_OrganizationMembershipTable,
+                                  LiteLLM_OrganizationTable,
+                                  LiteLLM_ProjectTableCachedObj,
+                                  LiteLLM_TagTable, LiteLLM_TeamMembership,
+                                  LiteLLM_TeamTable,
+                                  LiteLLM_TeamTableCachedObj,
+                                  LiteLLM_UserTable, LiteLLMRoutes,
+                                  LitellmUserRoles, NewTeamRequest,
+                                  ProxyErrorTypes, ProxyException,
+                                  RoleBasedPermissions, SpecialModelNames,
+                                  UserAPIKeyAuth)
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
 from litellm.proxy.guardrails.tool_name_extraction import (
@@ -326,6 +314,29 @@ async def common_checks(  # noqa: PLR0915
                 code=status.HTTP_401_UNAUTHORIZED,
             )
 
+    # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
+    if valid_token is not None and valid_token.agent_id:
+        from litellm.proxy.agent_endpoints.agent_registry import \
+            global_agent_registry
+        from litellm.proxy.litellm_pre_call_utils import \
+            get_chain_id_from_headers
+
+        agent = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id)
+        if agent is not None:
+            require_trace_id = (agent.litellm_params or {}).get(
+                "require_trace_id_on_calls_by_agent"
+            )
+            if require_trace_id:
+                headers_dict = dict(request.headers)
+                trace_id = get_chain_id_from_headers(headers_dict)
+                if not trace_id:
+                    raise ProxyException(
+                        message="Requests made with this agent's key must include the x-litellm-trace-id header.",
+                        type=ProxyErrorTypes.bad_request_error,
+                        param=None,
+                        code=status.HTTP_400_BAD_REQUEST,
+                    )
+
     ## 2.1 If user can call model (if personal key)
     if _model and team_object is None and user_object is not None:
         await can_user_call_model(
@@ -480,7 +491,8 @@ async def common_checks(  # noqa: PLR0915
     _request_metadata: dict = request_body.get("metadata", {}) or {}
     if _request_metadata.get("guardrails"):
         # check if team allowed to modify guardrails
-        from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails
+        from litellm.proxy.guardrails.guardrail_helpers import \
+            can_modify_guardrails
 
         can_modify: bool = can_modify_guardrails(team_object)
         if can_modify is False:
@@ -1930,9 +1942,8 @@ class ExperimentalUIJWTToken:
     def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str:
         from datetime import timedelta
 
-        from litellm.proxy.common_utils.encrypt_decrypt_utils import (
-            encrypt_value_helper,
-        )
+        from litellm.proxy.common_utils.encrypt_decrypt_utils import \
+            encrypt_value_helper
 
         if user_info.user_role is None:
             raise Exception("User role is required for experimental UI login")
@@ -1978,9 +1989,8 @@ class ExperimentalUIJWTToken:
         """
         from datetime import timedelta
 
-        from litellm.proxy.common_utils.encrypt_decrypt_utils import (
-            encrypt_value_helper,
-        )
+        from litellm.proxy.common_utils.encrypt_decrypt_utils import \
+            encrypt_value_helper
 
         if user_info.user_role is None:
             raise Exception("User role is required for CLI JWT login")
@@ -2019,9 +2029,8 @@ class ExperimentalUIJWTToken:
         import json
 
         from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
-        from litellm.proxy.common_utils.encrypt_decrypt_utils import (
-            decrypt_value_helper,
-        )
+        from litellm.proxy.common_utils.encrypt_decrypt_utils import \
+            decrypt_value_helper
 
         decrypted_token = decrypt_value_helper(
             hashed_token, key="ui_hash_key", exception_type="debug"
@@ -2142,13 +2151,11 @@ async def get_key_object(
         )
 
     # else, check db
-    _valid_token: Optional[BaseModel] = (
-        await _fetch_key_object_from_db_with_reconnect(
-            hashed_token=hashed_token,
-            prisma_client=prisma_client,
-            parent_otel_span=parent_otel_span,
-            proxy_logging_obj=proxy_logging_obj,
-        )
+    _valid_token: Optional[BaseModel] = await _fetch_key_object_from_db_with_reconnect(
+        hashed_token=hashed_token,
+        prisma_client=prisma_client,
+        parent_otel_span=parent_otel_span,
+        proxy_logging_obj=proxy_logging_obj,
     )
 
     if _valid_token is None:
@@ -2294,9 +2301,9 @@ async def get_org_object(
         # Cache the result
         await user_api_key_cache.async_set_cache(
             key=cache_key,
-            value=response.model_dump()
-            if hasattr(response, "model_dump")
-            else response,
+            value=(
+                response.model_dump() if hasattr(response, "model_dump") else response
+            ),
             ttl=DEFAULT_IN_MEMORY_TTL,
         )
 
@@ -2339,8 +2346,10 @@ async def _get_resources_from_access_groups(
     # Lazy import to avoid circular imports
     if prisma_client is None or user_api_key_cache is None:
         from litellm.proxy.proxy_server import prisma_client as _prisma_client
-        from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj
-        from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache
+        from litellm.proxy.proxy_server import \
+            proxy_logging_obj as _proxy_logging_obj
+        from litellm.proxy.proxy_server import \
+            user_api_key_cache as _user_api_key_cache
 
         prisma_client = prisma_client or _prisma_client
         user_api_key_cache = user_api_key_cache or _user_api_key_cache
@@ -3296,7 +3305,8 @@ async def _tag_max_budget_check(
         BudgetExceededError if any tag is over its max budget.
         Triggers a budget alert if any tag is over its max budget.
     """
-    from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
+    from litellm.proxy.common_utils.http_parsing_utils import \
+        get_tags_from_request_body
 
     if prisma_client is None:
         return

litellm/proxy/db/db_spend_update_writer.py

@@ -374,6 +374,18 @@ class DBSpendUpdateWriter:
                 traceback.format_exc(),
             )
 
+        try:
+            await self._update_agent_db(
+                response_cost=response_cost,
+                agent_id=payload_copy.get("agent_id"),
+                prisma_client=prisma_client,
+            )
+        except Exception:
+            verbose_proxy_logger.debug(
+                "_batch_database_updates: _update_agent_db failed: %s",
+                traceback.format_exc(),
+            )
+
         try:
             await self.add_spend_log_transaction_to_daily_user_transaction(
                 payload=payload_copy,
@@ -604,6 +616,37 @@ class DBSpendUpdateWriter:
             )
             raise e
 
+    async def _update_agent_db(
+        self,
+        response_cost: Optional[float],
+        agent_id: Optional[str],
+        prisma_client: Optional[PrismaClient],
+    ):
+        try:
+            if agent_id is None or prisma_client is None:
+                verbose_proxy_logger.debug(
+                    "track_cost_callback: agent_id is None or prisma_client is None. Not tracking spend for agent"
+                )
+                return
+
+            await self.spend_update_queue.add_update(
+                update=SpendUpdateQueueItem(
+                    entity_type=Litellm_EntityType.AGENT,
+                    entity_id=agent_id,
+                    response_cost=response_cost,
+                )
+            )
+        except Exception as e:
+            verbose_proxy_logger.error(
+                "Spend tracking - failed to enqueue agent spend update. "
+                "agent_id=%s, response_cost=%s - %s\n%s",
+                agent_id,
+                response_cost,
+                str(e),
+                traceback.format_exc(),
+            )
+            raise e
+
     async def _update_tag_db(
         self,
         response_cost: Optional[float],
@@ -765,7 +808,7 @@ class DBSpendUpdateWriter:
                 if db_spend_update_transactions is not None:
                     verbose_proxy_logger.info(
                         "Spend tracking - committing spend updates from Redis to DB: "
-                        "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d",
+                        "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d",
                         len(
                             db_spend_update_transactions.get("key_list_transactions")
                             or {}
@@ -798,6 +841,12 @@ class DBSpendUpdateWriter:
                             db_spend_update_transactions.get("tag_list_transactions")
                             or {}
                         ),
+                        len(
+                            db_spend_update_transactions.get(
+                                "agent_list_transactions"
+                            )
+                            or {}
+                        ),
                     )
                     await self._commit_spend_updates_to_db(
                         prisma_client=prisma_client,
@@ -1279,6 +1328,18 @@ class DBSpendUpdateWriter:
             proxy_logging_obj=proxy_logging_obj,
         )
 
+        ### UPDATE AGENT TABLE ###
+        agent_list_transactions = db_spend_update_transactions["agent_list_transactions"]
+        await DBSpendUpdateWriter._update_entity_spend_in_db(
+            entity_name="Agent",
+            transactions=agent_list_transactions,
+            table_accessor="litellm_agentstable",
+            where_field="agent_id",
+            n_retry_times=n_retry_times,
+            prisma_client=prisma_client,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
     @staticmethod
     async def _update_entity_spend_in_db(
         entity_name: str,

litellm/proxy/db/db_transaction_queue/redis_update_buffer.py

@@ -10,33 +10,31 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import RedisCache
-from litellm.constants import (
-    MAX_REDIS_BUFFER_DEQUEUE_COUNT,
-    REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
-    REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
-    REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
-    REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
-    REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
-    REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
-    REDIS_UPDATE_BUFFER_KEY,
-)
+from litellm.constants import (MAX_REDIS_BUFFER_DEQUEUE_COUNT,
+                               REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
+                               REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
+                               REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
+                               REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
+                               REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
+                               REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
+                               REDIS_UPDATE_BUFFER_KEY)
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
-from litellm.proxy._types import (
-    DailyTagSpendTransaction,
-    DailyTeamSpendTransaction,
-    DailyUserSpendTransaction,
-    DailyOrganizationSpendTransaction,
-    DailyEndUserSpendTransaction,
-    DBSpendUpdateTransactions,
-    DailyAgentSpendTransaction,
-)
-from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj
-from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
-    DailySpendUpdateQueue,
-)
-from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
+from litellm.proxy._types import (DailyAgentSpendTransaction,
+                                  DailyEndUserSpendTransaction,
+                                  DailyOrganizationSpendTransaction,
+                                  DailyTagSpendTransaction,
+                                  DailyTeamSpendTransaction,
+                                  DailyUserSpendTransaction,
+                                  DBSpendUpdateTransactions)
+from litellm.proxy.db.db_transaction_queue.base_update_queue import \
+    service_logger_obj
+from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \
+    DailySpendUpdateQueue
+from litellm.proxy.db.db_transaction_queue.spend_update_queue import \
+    SpendUpdateQueue
 from litellm.secret_managers.main import str_to_bool
-from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation
+from litellm.types.caching import (RedisPipelineLpopOperation,
+                                   RedisPipelineRpushOperation)
 from litellm.types.services import ServiceTypes
 
 if TYPE_CHECKING:
@@ -579,6 +577,7 @@ class RedisUpdateBuffer:
             team_member_list_transactions={},
             org_list_transactions={},
             tag_list_transactions={},
+            agent_list_transactions={},
         )
 
         # Define the transaction fields to process
@@ -590,6 +589,7 @@ class RedisUpdateBuffer:
             "team_member_list_transactions",
             "org_list_transactions",
             "tag_list_transactions",
+            "agent_list_transactions",
         ]
 
         # Loop through each transaction and combine the values

litellm/proxy/db/db_transaction_queue/spend_update_queue.py

@@ -3,15 +3,10 @@ from typing import Dict, List, Optional
 
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE
-from litellm.proxy._types import (
-    DBSpendUpdateTransactions,
-    Litellm_EntityType,
-    SpendUpdateQueueItem,
-)
+from litellm.proxy._types import (DBSpendUpdateTransactions,
+                                  Litellm_EntityType, SpendUpdateQueueItem)
 from litellm.proxy.db.db_transaction_queue.base_update_queue import (
-    BaseUpdateQueue,
-    service_logger_obj,
-)
+    BaseUpdateQueue, service_logger_obj)
 from litellm.types.services import ServiceTypes
 
 
@@ -145,6 +140,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
             team_member_list_transactions={},
             org_list_transactions={},
             tag_list_transactions={},
+            agent_list_transactions={},
         )
 
         # Map entity types to their corresponding transaction dictionary keys
@@ -156,6 +152,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
             Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
             Litellm_EntityType.ORGANIZATION: "org_list_transactions",
             Litellm_EntityType.TAG: "tag_list_transactions",
+            Litellm_EntityType.AGENT: "agent_list_transactions",
         }
 
         for update in updates:
@@ -207,6 +204,10 @@ class SpendUpdateQueue(BaseUpdateQueue):
                 transactions_dict = db_spend_update_transactions[
                     "tag_list_transactions"
                 ]
+            elif dict_key == "agent_list_transactions":
+                transactions_dict = db_spend_update_transactions[
+                    "agent_list_transactions"
+                ]
             else:
                 continue
 

litellm/proxy/litellm_pre_call_utils.py

@@ -10,16 +10,12 @@ import litellm
 from litellm._logging import verbose_logger, verbose_proxy_logger
 from litellm._service_logger import ServiceLogging
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
-from litellm.proxy._types import (
-    AddTeamCallback,
-    CommonProxyErrors,
-    LitellmDataForBackendLLMCall,
-    LitellmUserRoles,
-    SpecialHeaders,
-    TeamCallbackMetadata,
-    UserAPIKeyAuth,
-)
-from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
+from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors,
+                                  LitellmDataForBackendLLMCall,
+                                  LitellmUserRoles, SpecialHeaders,
+                                  TeamCallbackMetadata, UserAPIKeyAuth)
+from litellm.proxy.common_utils.http_parsing_utils import \
+    _safe_get_request_headers
 
 # Cache special headers as a frozenset for O(1) lookup performance
 _SPECIAL_HEADERS_CACHE = frozenset(
@@ -28,12 +24,9 @@ _SPECIAL_HEADERS_CACHE = frozenset(
 from litellm.router import Router
 from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
 from litellm.types.services import ServiceTypes
-from litellm.types.utils import (
-    LlmProviders,
-    ProviderSpecificHeader,
-    StandardLoggingUserAPIKeyMetadata,
-    SupportedCacheControls,
-)
+from litellm.types.utils import (LlmProviders, ProviderSpecificHeader,
+                                 StandardLoggingUserAPIKeyMetadata,
+                                 SupportedCacheControls)
 
 service_logger_obj = ServiceLogging()  # used for tracking latency on OTEL
 
@@ -667,6 +660,12 @@ class LiteLLMProxyRequestSetup:
             "user_api_key"
         ] = user_api_key_dict.api_key  # this is just the hashed token
 
+        # Key-owned agent_id for spend attribution; keep existing (e.g. from header) if key has none
+        data[_metadata_variable_name]["agent_id"] = (
+            getattr(user_api_key_dict, "agent_id", None)
+            or data[_metadata_variable_name].get("agent_id")
+        )
+
         data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr(
             user_api_key_dict, "end_user_max_budget", None
         )
@@ -689,8 +688,7 @@ class LiteLLMProxyRequestSetup:
             return data
         from litellm.proxy._types import (
             LiteLLM_ManagementEndpoint_MetadataFields,
-            LiteLLM_ManagementEndpoint_MetadataFields_Premium,
-        )
+            LiteLLM_ManagementEndpoint_MetadataFields_Premium)
 
         # ignore any special fields
         added_metadata = {}
@@ -859,7 +857,8 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
     """
 
     from litellm.proxy.proxy_server import llm_router, premium_user
-    from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields
+    from litellm.types.proxy.litellm_pre_call_utils import (RedactedDict,
+                                                            SecretFields)
 
     _raw_headers: Dict[str, str] = RedactedDict(_safe_get_request_headers(request))
 
@@ -1543,7 +1542,8 @@ async def move_guardrails_to_metadata(
 
     # Only check policy engine if no local config (avoid import + registry lookup)
     if not (has_key_config or has_team_config or has_request_config):
-        from litellm.proxy.policy_engine.policy_registry import get_policy_registry
+        from litellm.proxy.policy_engine.policy_registry import \
+            get_policy_registry
 
         if not get_policy_registry().is_initialized():
             # Nothing configured anywhere - clean up request body fields and return
@@ -1607,14 +1607,16 @@ async def move_guardrails_to_metadata(
 
 def _is_policy_version_id(s: str) -> bool:
     """Return True if string is a policy version ID (starts with policy_<uuid> prefix)."""
-    from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX
+    from litellm.proxy.policy_engine.policy_registry import \
+        POLICY_VERSION_ID_PREFIX
 
     return isinstance(s, str) and s.startswith(POLICY_VERSION_ID_PREFIX)
 
 
 def _extract_policy_id(s: str) -> Optional[str]:
     """Extract raw UUID from policy_<uuid> string, or None if not a valid version ID."""
-    from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX
+    from litellm.proxy.policy_engine.policy_registry import \
+        POLICY_VERSION_ID_PREFIX
 
     if not _is_policy_version_id(s):
         return None
@@ -1635,10 +1637,9 @@ def _match_and_track_policies(
     """
     from litellm._logging import verbose_proxy_logger
     from litellm.proxy.common_utils.callback_utils import (
-        add_policy_sources_to_metadata,
-        add_policy_to_applied_policies_header,
-    )
-    from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
+        add_policy_sources_to_metadata, add_policy_to_applied_policies_header)
+    from litellm.proxy.policy_engine.attachment_registry import \
+        get_attachment_registry
     from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
 
     # Get matching policies via attachments (with match reasons for attribution)
@@ -1783,7 +1784,8 @@ async def add_guardrails_from_policy_engine(
         user_api_key_dict: The user's API key authentication info
     """
     from litellm._logging import verbose_proxy_logger
-    from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
+    from litellm.proxy.common_utils.http_parsing_utils import \
+        get_tags_from_request_body
     from litellm.proxy.policy_engine.policy_registry import get_policy_registry
     from litellm.types.proxy.policy_engine import PolicyMatchContext
 

litellm/proxy/schema.prisma

@@ -66,6 +66,7 @@ model LiteLLM_AgentsTable {
   agent_access_groups String[]       @default([])
   object_permission_id String?
   object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
+  spend      Float    @default(0.0)
   created_at    DateTime               @default(now()) @map("created_at")
   created_by String
   updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")