diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql index 6f1db719855..033cd4101c2 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql @@ -13,3 +13,5 @@ END $$; ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "spend_window" TIMESTAMP(3); + +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "lifetime_budget_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 6823deadbd3..fc7325ddac3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -86,6 +86,7 @@ model LiteLLM_AgentsTable { retired_identities LiteLLM_RetiredAgentIdentity[] budget_id String? @unique spend_window DateTime? + lifetime_budget_spend Float @default(0.0) litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) tpm_limit Int? rpm_limit Int? diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 1c760d8dc6b..89bdd6647ce 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3246,6 +3246,11 @@ } ] }, + "lifetime_budget_spend": { + "default": 0.0, + "title": "Lifetime Budget Spend", + "type": "number" + }, "litellm_budget_table": { "anyOf": [ { diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 7dbe409226a..5104d904f26 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -91,6 +91,15 @@ class AgentRecord(Protocol): @property def spend(self) -> float: ... + @property + def lifetime_budget_spend(self) -> float: ... + + @property + def budget_id(self) -> str | None: ... + + @property + def litellm_budget_table(self) -> "prisma_models.LiteLLM_BudgetTable | None": ... + def model_dump(self) -> AgentRecordDump: ... def __iter__(self) -> Iterator[tuple[str, object]]: ... diff --git a/litellm/proxy/agent_endpoints/auth/managed_authorization.py b/litellm/proxy/agent_endpoints/auth/managed_authorization.py index ef437b029fe..e254430def3 100644 --- a/litellm/proxy/agent_endpoints/auth/managed_authorization.py +++ b/litellm/proxy/agent_endpoints/auth/managed_authorization.py @@ -214,7 +214,7 @@ async def check_agent_budget(auth: UserAPIKeyAuth) -> None: budget: Final = agent.litellm_budget_table.max_budget spend: Final = await get_current_spend( counter_key=agent.budget_counter_key, - fallback_spend=agent.spend or 0.0, + fallback_spend=agent.budget_spend, max_budget=budget, fallback_authoritative=True, ) @@ -254,7 +254,12 @@ async def prepare_agent_invocation( if target is None and registered_managed: raise_identity_failure(AgentIdentityFailure(message="Invoked agent no longer exists")) effective: Final = target if target is not None else registered - if not effective.identity_managed and effective.litellm_budget_table is None and auth.managed_agent_policy is None: + if ( + not effective.identity_managed + and effective.litellm_budget_table is None + and auth.managed_agent_policy is None + and auth.billing_agent_policy is None + ): return if not await AgentRequestHandler.is_agent_allowed(effective.agent_id, auth): raise_identity_failure(AgentIdentityFailure(message="The caller is not permitted to invoke this agent")) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 5ea8266d5ed..a76ec57b630 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -13,7 +13,7 @@ import os import uuid from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import Annotated, Final, TypedDict +from typing import TYPE_CHECKING, Annotated, Final, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query, Request from pydantic import ValidationError @@ -77,6 +77,7 @@ from litellm.types.agents import ( ) from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.proxy.agent_identity import ( + AgentBudgetState, AgentIdentityBinding, AgentIdentityFailure, EntraIdentityConfig, @@ -129,6 +130,12 @@ def _build_merged_agent_card( ) +if TYPE_CHECKING: + from prisma.types import LiteLLM_AgentsTableInclude + +_AGENT_BUDGET_INCLUDE: Final["LiteLLM_AgentsTableInclude"] = {"litellm_budget_table": True} + + router: Final = APIRouter() @@ -374,8 +381,10 @@ async def get_agents( if agent_ids: db_agents: Final = await agents_table(prisma_client).find_many( where={"agent_id": {"in": agent_ids}}, + include=_AGENT_BUDGET_INCLUDE, ) - spend_map: Final = {a.agent_id: a.spend for a in db_agents} + spend_map: Final = MappingProxyType({a.agent_id: a.spend for a in db_agents}) + budget_map: Final = MappingProxyType({a.agent_id: a for a in db_agents}) for agent in returned_agents: matched_spends = tuple( spend_map[alias_id] @@ -384,6 +393,14 @@ async def get_agents( ) if matched_spends: agent.spend = sum(matched_spends) + if (budget_row := budget_map.get(agent.agent_id)) is not None: + agent.lifetime_budget_spend = budget_row.lifetime_budget_spend + agent.budget_id = budget_row.budget_id + agent.litellm_budget_table = ( + AgentBudgetState.model_validate(budget_row.litellm_budget_table.model_dump()) + if budget_row.litellm_budget_table is not None + else None + ) await _attach_keys_to_agents(returned_agents, prisma_client) # add is_public field to each agent - we do it this way, to allow setting config agents as public @@ -676,7 +693,7 @@ async def get_agent_by_id( if agent is None: agent_row: Final = await agents_table(prisma_client).find_unique( where={"agent_id": agent_id}, - include={"object_permission": True, "identity": True}, + include={"object_permission": True, "identity": True, "litellm_budget_table": True}, ) if agent_row is not None: agent_dict: Final = agent_row.model_dump() @@ -688,9 +705,18 @@ async def get_agent_by_id( agent = AgentResponse(**agent_dict) else: # Agent found in memory — refresh spend from DB - db_row: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id}) + db_row: Final = await agents_table(prisma_client).find_unique( + where={"agent_id": agent_id}, include=_AGENT_BUDGET_INCLUDE + ) if db_row is not None: agent.spend = db_row.spend + agent.lifetime_budget_spend = db_row.lifetime_budget_spend + agent.budget_id = db_row.budget_id + agent.litellm_budget_table = ( + AgentBudgetState.model_validate(db_row.litellm_budget_table.model_dump()) + if db_row.litellm_budget_table is not None + else None + ) if agent is None: raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found") diff --git a/litellm/proxy/agent_endpoints/managed_identity.py b/litellm/proxy/agent_endpoints/managed_identity.py index 87041e92459..e2289e3bf47 100644 --- a/litellm/proxy/agent_endpoints/managed_identity.py +++ b/litellm/proxy/agent_endpoints/managed_identity.py @@ -82,6 +82,7 @@ class ManagedWriteFields(TypedDict, total=False): litellm_budget_table: ReadOnly[BudgetRelationWrite] spend_window: ReadOnly[datetime | None] spend: ReadOnly[float] + lifetime_budget_spend: ReadOnly[float] def raise_identity_failure(failure: AgentIdentityFailure, status_code: int = 403) -> NoReturn: @@ -202,6 +203,11 @@ def _budget_write(raw: object, existing: AgentResponse | None, updated_by: str) duration_error: Final = budget_duration_error(budget.budget_duration) if duration_error is not None: raise ValueError(duration_error) + creating_lifetime: Final = budget.budget_duration is None and ( + existing is None + or existing.litellm_budget_table is None + or existing.litellm_budget_table.budget_duration is not None + ) fields: Final[BudgetFields] = { "max_budget": budget.max_budget, "budget_duration": budget.budget_duration, @@ -218,6 +224,7 @@ def _budget_write(raw: object, existing: AgentResponse | None, updated_by: str) } result: Final[ManagedWriteFields] = { "spend_window": fields["budget_reset_at"], + **({"lifetime_budget_spend": 0.0} if creating_lifetime else {}), **( {"spend": 0.0} if fields["budget_reset_at"] is not None @@ -229,7 +236,9 @@ def _budget_write(raw: object, existing: AgentResponse | None, updated_by: str) else {} ), "litellm_budget_table": ( - {"update": fields} if existing and existing.budget_id else {"create": {**fields, "created_by": updated_by}} + {"update": fields} + if existing and existing.budget_id and not creating_lifetime + else {"create": {**fields, "created_by": updated_by}} ), } return result diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c88ef2544d6..55839620831 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -86,6 +86,8 @@ from litellm.types.agents import agent_spend_filter from litellm.types.utils import CallTypes if TYPE_CHECKING: + from prisma.types import LiteLLM_AgentsTableUpdateManyMutationInput, LiteLLM_AgentsTableWhereInput + from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -163,6 +165,26 @@ def _entity_spend_table(batcher: _SpendBatch, table_accessor: _EntitySpendTable) return _ENTITY_SPEND_TABLES[table_accessor](batcher) +def _queue_lifetime_agent_spend(table: BatchTable, counter_key: str, response_cost: float) -> None: + lifetime_filter: Final = agent_spend_filter(counter_key) + lifetime_data: Final[LiteLLM_AgentsTableUpdateManyMutationInput] = { + "lifetime_budget_spend": {"increment": response_cost} + } + history_filter: Final[LiteLLM_AgentsTableWhereInput] = { + "agent_id": lifetime_filter.get("agent_id"), + "spend_window": None, + } + history_data: Final[LiteLLM_AgentsTableUpdateManyMutationInput] = {"spend": {"increment": response_cost}} + table.update_many( + where=lifetime_filter, + data=lifetime_data, + ) + table.update_many( + where=history_filter, + data=history_data, + ) + + class _SpendBatchManager(Protocol): async def __aenter__(self) -> _SpendBatch: ... @@ -1361,7 +1383,7 @@ class DBSpendUpdateWriter: try: if agent_id is None or prisma_client is None: return - if counter_key is not None and agent_spend_filter(counter_key)["agent_id"] != agent_id: + if counter_key is not None and agent_spend_filter(counter_key).get("agent_id") != agent_id: raise ValueError("Agent spend counter does not match the billed agent") await self.spend_update_queue.add_update( @@ -2348,6 +2370,13 @@ class DBSpendUpdateWriter: entity_id, response_cost, ) + if table_accessor == "litellm_agentstable" and entity_id.startswith( + "spend:agent_lifetime:" + ): + _queue_lifetime_agent_spend( + _entity_spend_table(batcher, table_accessor), entity_id, response_cost + ) + continue _entity_spend_table(batcher, table_accessor).update_many( where=( agent_spend_filter(entity_id) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index fc2fbd4dec8..3222868dad8 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -42,7 +42,11 @@ from litellm.repositories.verification_token_repository import ( from litellm.types.agents import agent_budget_counter_key if TYPE_CHECKING: - from prisma.types import LiteLLM_EndUserTableWhereUniqueInput + from prisma.types import ( + LiteLLM_AgentsTableInclude, + LiteLLM_AgentsTableWhereUniqueInput, + LiteLLM_EndUserTableWhereUniqueInput, + ) from litellm.caching.dual_cache import DualCache from litellm.proxy.utils import PrismaClient @@ -172,21 +176,37 @@ class SpendCounterReseed: return await OrganizationRepository(prisma_client).table.find_unique( where={"organization_id": counter_key[len("spend:org:") :]} ) + if counter_key.startswith("spend:agent_lifetime:"): + _, _, budget_id, agent_id = counter_key.split(":", 3) + lifetime_where: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id} + lifetime: Final = await AgentsRepository(prisma_client, use_writer=True).table.find_unique( + where=lifetime_where + ) + if lifetime is None: + return None + return lifetime.model_copy( + update=MappingProxyType( + {"spend": lifetime.lifetime_budget_spend if lifetime.budget_id == budget_id else 0.0} + ) + ) if counter_key.startswith("spend:agent_window:"): parts: Final = counter_key.split(":", 3) if len(parts) != 4: return None + window_where: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": parts[3]} + window_include: Final[LiteLLM_AgentsTableInclude] = {"litellm_budget_table": True} row: Final = await AgentsRepository(prisma_client, use_writer=True).table.find_unique( - where={"agent_id": parts[3]}, include={"litellm_budget_table": True} + where=window_where, include=window_include ) if row is None: return None current_key: Final = agent_budget_counter_key(row.agent_id, row.spend_window) - return row if current_key == counter_key else row.model_copy(update={"spend": 0.0}) + return row if current_key == counter_key else row.model_copy(update=MappingProxyType({"spend": 0.0})) if counter_key.startswith("spend:agent:"): - return await AgentsRepository(prisma_client).table.find_unique( - where={"agent_id": counter_key[len("spend:agent:") :]} - ) + agent_where: Final[LiteLLM_AgentsTableWhereUniqueInput] = { + "agent_id": counter_key[len("spend:agent:") :] + } + return await AgentsRepository(prisma_client).table.find_unique(where=agent_where) if counter_key.startswith("spend:project:"): return await ProjectRepository(prisma_client).table.find_unique( where={"project_id": counter_key[len("spend:project:") :]} diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c914e04e779..9cf1e0d6eed 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -742,6 +742,8 @@ async def _update_database_and_spend_counters( tags=request_tags, model_access_groups=model_access_groups, project_id=project_id, + billing_agent_id=billing_agent_id, + billing_agent_counter_key=billing_agent_counter_key, ) with spend_counter_batch_scope(spend_counter_cache.redis_cache, counter_keys=counter_keys): return await _update_database_and_spend_counters_in_batch( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 6823deadbd3..fc7325ddac3 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -86,6 +86,7 @@ model LiteLLM_AgentsTable { retired_identities LiteLLM_RetiredAgentIdentity[] budget_id String? @unique spend_window DateTime? + lifetime_budget_spend Float @default(0.0) litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) tpm_limit Int? rpm_limit Int? diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 93cba37a90c..62a6bdbe0fd 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -511,7 +511,7 @@ async def _get_budget_counters( counter_key=agent.budget_counter_key, source_cache_key=None, max_budget=agent.litellm_budget_table.max_budget, - fallback_spend=agent.spend or 0.0, + fallback_spend=agent.budget_spend, entity_type="Agent", entity_id=agent.agent_id, ) diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index b3d0e8646e8..5ae30a5c027 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -207,8 +207,11 @@ def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> fr billing_agent: Final = token.billing_agent_policy charged_agent_id: Final = billing_agent.agent_id if billing_agent is not None else token.agent_id agent_keys: Final = ( - frozenset((billing_agent.budget_counter_key if billing_agent is not None else f"spend:agent:{charged_agent_id}",)) - if charged_agent_id is not None else frozenset() + frozenset( + (billing_agent.budget_counter_key if billing_agent is not None else f"spend:agent:{charged_agent_id}",) + ) + if charged_agent_id is not None + else frozenset() ) return agent_keys | frozenset( _iter_entity_counter_keys( @@ -251,13 +254,10 @@ def post_call_counter_keys( for group in model_access_groups or () if group and isinstance(group, str) ) - agent_key: Final = billing_agent_counter_key or (f"spend:agent:{billing_agent_id}" if billing_agent_id is not None else None) - return ( - entity_keys - | tag_keys - | group_keys - | (frozenset((agent_key,)) if agent_key else frozenset()) + agent_key: Final = billing_agent_counter_key or ( + f"spend:agent:{billing_agent_id}" if billing_agent_id is not None else None ) + return entity_keys | tag_keys | group_keys | (frozenset((agent_key,)) if agent_key else frozenset()) def bind_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> None: diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 964ef94d227..d59b6fa6ce9 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -316,7 +316,9 @@ class AgentKeySummary(BaseModel): key_name: str | None = None -def agent_budget_counter_key(agent_id: str, reset_at: datetime | None) -> str: +def agent_budget_counter_key(agent_id: str, reset_at: datetime | None, budget_id: str | None = None) -> str: + if reset_at is None and budget_id is not None: + return f"spend:agent_lifetime:{budget_id}:{agent_id}" if reset_at is None: return f"spend:agent:{agent_id}" aware: Final = reset_at if reset_at.tzinfo is not None else reset_at.replace(tzinfo=timezone.utc) @@ -325,6 +327,14 @@ def agent_budget_counter_key(agent_id: str, reset_at: datetime | None) -> str: def agent_spend_filter(counter_key: str) -> "LiteLLM_AgentsTableWhereInput": + if counter_key.startswith("spend:agent_lifetime:"): + _, _, budget_id, agent_id = counter_key.split(":", 3) + lifetime: Final[LiteLLM_AgentsTableWhereInput] = { + "agent_id": agent_id, + "budget_id": budget_id, + "spend_window": None, + } + return lifetime if counter_key.startswith("spend:agent_window:"): _, _, raw_window, agent_id = counter_key.split(":", 3) window: Final = datetime.strptime(raw_window, "%Y%m%dT%H%M%S.%fZ").replace(tzinfo=timezone.utc) @@ -339,6 +349,7 @@ def agent_spend_filter(counter_key: str) -> "LiteLLM_AgentsTableWhereInput": class AgentResponse(BaseModel): budget_id: str | None = None + lifetime_budget_spend: float = 0.0 litellm_budget_table: AgentBudgetState | None = None identity: AgentIdentityBinding | None = None identity_managed: bool = False @@ -369,9 +380,17 @@ class AgentResponse(BaseModel): @property def budget_counter_key(self) -> str: return agent_budget_counter_key( - self.agent_id, self.litellm_budget_table.budget_reset_at if self.litellm_budget_table else None + self.agent_id, + self.litellm_budget_table.budget_reset_at if self.litellm_budget_table else None, + self.litellm_budget_table.budget_id if self.litellm_budget_table else None, ) + @property + def budget_spend(self) -> float: + if self.litellm_budget_table is not None and self.litellm_budget_table.budget_duration is None: + return self.lifetime_budget_spend + return self.spend or 0.0 + class ListAgentsResponse(BaseModel): agents: list[AgentResponse] diff --git a/schema.prisma b/schema.prisma index 6823deadbd3..fc7325ddac3 100644 --- a/schema.prisma +++ b/schema.prisma @@ -86,6 +86,7 @@ model LiteLLM_AgentsTable { retired_identities LiteLLM_RetiredAgentIdentity[] budget_id String? @unique spend_window DateTime? + lifetime_budget_spend Float @default(0.0) litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) tpm_limit Int? rpm_limit Int? diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py index a2c9fca3bd0..c9b4d301018 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py @@ -108,7 +108,7 @@ async def test_agent_budget_accumulates_across_credentials_and_denies_the_next_a from litellm.proxy.agent_endpoints.auth.managed_authorization import check_agent_budget counters: Final = DualCache() - counters.set_cache("spend:agent:agent", 0.0) + counters.set_cache("spend:agent_lifetime:budget:agent", 0.0) counters.set_cache("spend:key:first", 0.0) counters.set_cache("spend:key:second", 0.0) monkeypatch.setattr(proxy_server, "spend_counter_cache", counters) @@ -117,15 +117,15 @@ async def test_agent_budget_accumulates_across_credentials_and_denies_the_next_a auth.billing_agent_policy = policy await check_agent_budget(auth) await proxy_server.increment_spend_counters( - token="first", team_id=None, user_id=None, response_cost=0.3, billing_agent_id="agent" + token="first", team_id=None, user_id=None, response_cost=0.3, billing_agent_id="agent", billing_agent_counter_key=policy.budget_counter_key ) await check_agent_budget(auth) await proxy_server.increment_spend_counters( - token="second", team_id=None, user_id=None, response_cost=0.3, billing_agent_id="agent" + token="second", team_id=None, user_id=None, response_cost=0.3, billing_agent_id="agent", billing_agent_counter_key=policy.budget_counter_key ) with pytest.raises(litellm.BudgetExceededError): await check_agent_budget(auth) - assert await counters.async_get_cache("spend:agent:agent") == pytest.approx(0.6) + assert await counters.async_get_cache("spend:agent_lifetime:budget:agent") == pytest.approx(0.6) assert await counters.async_get_cache("spend:key:first") == pytest.approx(0.3) assert await counters.async_get_cache("spend:key:second") == pytest.approx(0.3) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index 9ed9859717d..33b9a5d4b9d 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -1231,6 +1231,7 @@ def _stored_agent_row(values: Mapping[str, object] | SimpleNamespace) -> LiteLLM "created_by": "admin", "updated_by": "admin", "spend": 0, + "lifetime_budget_spend": 0, "identity_managed": False, "enabled": True, "execution_mode": "autonomous", @@ -1431,6 +1432,7 @@ async def test_agent_listing_preserves_stored_identity_bindings(bound: bool) -> enabled=True, execution_mode="autonomous", spend=0.0, + lifetime_budget_spend=0.0, agent_access_groups=[], access_group_ids=[], extra_headers=[], diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 2cf81892db7..0bf9afa832d 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -1648,3 +1648,31 @@ def test_invalid_identity_and_untrusted_tenant_cannot_be_registered( with pytest.raises(HTTPException, match=message) as failure: agent_endpoints._validate_managed_identity_request(request) assert failure.value.status_code == 400 + + +@pytest.mark.parametrize("cached,path", [(False, "/v1/agents/agent-123"), (True, "/v1/agents/agent-123"), (True, "/v1/agents")]) +def test_agent_budget_readback_refreshes_consumption_and_limit(monkeypatch: pytest.MonkeyPatch, cached: bool, path: str) -> None: + from prisma.models import LiteLLM_BudgetTable + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints import agent_registry + from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry + + row: Final = LiteLLM_AgentsTable.model_construct( + agent_id="agent-123", agent_name="Agent", agent_card_params={}, spend=12.5, + lifetime_budget_spend=0.75, budget_id="budget", litellm_params=None, + litellm_budget_table=LiteLLM_BudgetTable.model_construct(budget_id="budget", max_budget=2.0), + ) + registry: Final = AgentRegistry() + if cached: + registry.register_agent(_sample_agent_response()) + table: Final = SimpleNamespace(find_unique=AsyncMock(return_value=row), find_many=AsyncMock(return_value=[row])) + database: Final = SimpleNamespace(litellm_agentstable=table, litellm_verificationtoken=SimpleNamespace(find_many=AsyncMock(return_value=[]))) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=database, writer_db=database)) + monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry) + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + response: Final = client.get(path) + assert response.status_code == 200, response.text + payload: Final = response.json()[0] if path == "/v1/agents" else response.json() + assert payload["spend"] == 12.5 + assert payload["lifetime_budget_spend"] == 0.75 + assert payload["litellm_budget_table"]["max_budget"] == 2.0 diff --git a/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py b/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py index 424330e078e..55b1a0d366e 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py @@ -311,3 +311,46 @@ def test_budget_write_stamps_the_same_window_on_the_agent_row() -> None: assert not isinstance(result, AgentIdentityFailure) assert result["spend_window"] == result["litellm_budget_table"]["create"]["budget_reset_at"] assert result["spend"] == 0.0 + + +def test_new_lifetime_budget_starts_unused_without_erasing_historical_spend() -> None: + existing: Final = managed_agent().model_copy(update={"spend": 12.5}) + result: Final = managed_write_fields({"budget": {"max_budget": 1.0}}, existing, "admin") + assert not isinstance(result, AgentIdentityFailure) + assert "spend" not in result + assert result["lifetime_budget_spend"] == 0.0 + assert result["litellm_budget_table"]["create"]["max_budget"] == 1.0 + assert existing.spend == 12.5 + + +def test_editing_lifetime_budget_preserves_consumption() -> None: + from litellm.types.proxy.agent_identity import AgentBudgetState + + existing: Final = managed_agent().model_copy(update={ + "spend": 12.5, "lifetime_budget_spend": 0.75, "budget_id": "budget", + "litellm_budget_table": AgentBudgetState(budget_id="budget", max_budget=1.0), + }) + result: Final = managed_write_fields({"budget": {"max_budget": 2.0}}, existing, "admin") + assert not isinstance(result, AgentIdentityFailure) + assert "spend" not in result + assert "lifetime_budget_spend" not in result + assert result["litellm_budget_table"]["update"]["max_budget"] == 2.0 + + +@pytest.mark.parametrize("previous_duration", (None, "1d")) +def test_recreated_or_converted_lifetime_budget_gets_a_fresh_allowance(previous_duration: str | None) -> None: + from litellm.types.proxy.agent_identity import AgentBudgetState + + existing: Final = managed_agent().model_copy(update={ + "spend": 12.5, "lifetime_budget_spend": 0.75, + "budget_id": "previous" if previous_duration else None, + "litellm_budget_table": AgentBudgetState( + budget_id="previous", max_budget=1.0, budget_duration=previous_duration, + ) if previous_duration else None, + }) + result: Final = managed_write_fields({"budget": {"max_budget": 2.0}}, existing, "admin") + assert not isinstance(result, AgentIdentityFailure) + assert result["lifetime_budget_spend"] == 0.0 + assert "spend" not in result + assert "create" in result["litellm_budget_table"] + assert "update" not in result["litellm_budget_table"] 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 9d0bc1c3475..a60d64bd562 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 @@ -9762,6 +9762,8 @@ async def test_human_agent_discovery_does_not_reserve_target_budget_but_send_and user_id="human", object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="grant", agents=["agent"]), ) + database.get_data = AsyncMock(return_value=auth) + proxy_server.proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock(return_value=None) with patch( "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", new_callable=AsyncMock, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index ca6a9411005..ae338c518b3 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -4885,3 +4885,32 @@ async def test_agent_settlement_charges_only_the_matching_current_window(capture "agent_id": "window-agent", "spend_window": captured_window, } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("captured_budget", ("current", "retired")) +async def test_lifetime_settlement_preserves_history_without_charging_another_budget(captured_budget: str) -> None: + from types import SimpleNamespace + + row: Final = SimpleNamespace( + agent_id="agent", budget_id="current", spend_window=None, spend=12.5, lifetime_budget_spend=0.25 + ) + + def apply_update(*, where: dict[str, object], data: dict[str, dict[str, float]]) -> None: + if all(getattr(row, key) == value for key, value in where.items()): + for field, operation in data.items(): + setattr(row, field, getattr(row, field) + operation["increment"]) + + batcher: Final = MagicMock() + batcher.litellm_agentstable.update_many.side_effect = apply_update + transaction: Final = AsyncMock() + transaction.batch_ = MagicMock(return_value=AsyncMock(__aenter__=AsyncMock(return_value=batcher))) + client: Final = MagicMock() + client.db.tx.return_value = AsyncMock(__aenter__=AsyncMock(return_value=transaction)) + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Agent", transactions={f"spend:agent_lifetime:{captured_budget}:agent": 0.25}, + table_accessor="litellm_agentstable", where_field="agent_id", n_retry_times=0, + prisma_client=client, proxy_logging_obj=MagicMock(), + ) + assert row.spend == 12.75 + assert row.lifetime_budget_spend == (0.5 if captured_budget == "current" else 0.25) diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 391cede1005..8e742c06b51 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -523,3 +523,17 @@ async def test_agent_window_reseed_handles_missing_rows_and_malformed_keys(missi key: Final = "spend:agent_window:20260102T000000.000000Z:missing" if missing else "spend:agent_window:malformed" assert await SpendCounterReseed.from_db(client, key) is None assert lookup.await_count == int(missing) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget_id,expected", [("current", 0.25), ("retired", 0.0), ("deleted", None)]) +async def test_lifetime_counter_reseed_excludes_historical_and_other_budget_spend(budget_id: str, expected: float | None) -> None: + from prisma.models import LiteLLM_AgentsTable + + row: Final = LiteLLM_AgentsTable.model_construct( + agent_id="agent", budget_id="current", spend=12.5, lifetime_budget_spend=0.25, + ) + lookup: Final = AsyncMock(return_value=row if expected is not None else None) + client: Final = SimpleNamespace(writer_db=SimpleNamespace(litellm_agentstable=SimpleNamespace(find_unique=lookup))) + assert await SpendCounterReseed.from_db(client, f"spend:agent_lifetime:{budget_id}:agent") == expected + lookup.assert_awaited_once_with(where={"agent_id": "agent"}) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 7096bc7c632..b5ca6f5c42d 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -4759,6 +4759,7 @@ def _agent_db_row(agent_id: str, agent_name: str): agent_access_groups=[], access_group_ids=[], spend=0.0, + lifetime_budget_spend=0.0, identity_managed=False, enabled=True, execution_mode="autonomous", diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 9d64fdb5aab..02686a13a85 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -317,13 +317,13 @@ async def test_agent_invocation_reserves_exact_fee_and_reconciles_without_child_ from litellm.types.agents import AgentResponse cache: Final = DualCache() - counter_key: Final = f"spend:agent_window:20260102T000000.000000Z:{charged_agent}" if window else f"spend:agent:{charged_agent}" + counter_key: Final = f"spend:agent_window:20260102T000000.000000Z:{charged_agent}" if window else f"spend:agent_lifetime:agent-budget:{charged_agent}" cache.set_cache(counter_key, 0.1) monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) auth: Final = UserAPIKeyAuth(agent_id="caller-agent" if charged_agent == "caller-agent" else None) auth.billing_agent_policy = AgentResponse( - agent_id=charged_agent, agent_name="Charged agent", agent_card_params={}, spend=0.1, - litellm_budget_table={"budget_id": "agent-budget", "max_budget": 0.5, "budget_reset_at": window}, + agent_id=charged_agent, agent_name="Charged agent", agent_card_params={}, spend=0.1, lifetime_budget_spend=0.1, + litellm_budget_table={"budget_id": "agent-budget", "max_budget": 0.5, "budget_reset_at": window, "budget_duration": "1d" if window else None}, ) auth.invoked_agent_id = "target-agent" auth.agent_invocation_cost = 0.2 @@ -359,11 +359,11 @@ async def test_agent_invocation_over_budget_is_rejected_and_reservation_is_refun from litellm.types.agents import AgentResponse cache: Final = DualCache() - cache.set_cache("spend:agent:agent", 0.4) + cache.set_cache("spend:agent_lifetime:agent-budget:agent", 0.4) monkeypatch.setattr(proxy_server, "spend_counter_cache", cache) auth: Final = UserAPIKeyAuth(agent_id="agent") auth.billing_agent_policy = AgentResponse( - agent_id="agent", agent_name="Charged agent", agent_card_params={}, spend=0.4, + agent_id="agent", agent_name="Charged agent", agent_card_params={}, spend=0.4, lifetime_budget_spend=0.4, litellm_budget_table={"budget_id": "agent-budget", "max_budget": 0.5}, ) auth.agent_invocation_cost = 0.2 @@ -374,7 +374,7 @@ async def test_agent_invocation_over_budget_is_rejected_and_reservation_is_refun user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), fail_closed_budget_enforcement=True, ) - assert await cache.async_get_cache("spend:agent:agent") == pytest.approx(0.4) + assert await cache.async_get_cache("spend:agent_lifetime:agent-budget:agent") == pytest.approx(0.4) @pytest.mark.asyncio @@ -413,3 +413,68 @@ async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_call assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"]) assert reservation["finalized"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outcome", ("failed", "cancelled", "completed")) +async def test_budgeted_caller_reserves_unmanaged_agent_fees_before_concurrent_admission( + spend_counter_cache: DualCache, monkeypatch: pytest.MonkeyPatch, outcome: str +) -> None: + import asyncio + + from litellm.proxy.agent_endpoints import agent_registry + from litellm.proxy.agent_endpoints.auth.managed_authorization import prepare_agent_invocation + from litellm.proxy.spend_tracking.budget_reservation import ( + reconcile_budget_reservation, + release_budget_reservation, + release_budget_reservation_on_cancel, + ) + from litellm.types.agents import AgentResponse + + caller: Final = AgentResponse( + agent_id="caller", agent_name="Caller", agent_card_params={}, spend=0.0, + litellm_budget_table={"budget_id": "caller-budget", "max_budget": 0.5}, + ) + target: Final = AgentResponse( + agent_id="target", agent_name="Target", agent_card_params={}, litellm_params={"cost_per_query": 0.25}, + ) + registry: Final = agent_registry.AgentRegistry() + registry.register_agent(target) + monkeypatch.setattr(agent_registry, "global_agent_registry", registry) + + async def admit() -> dict[str, object] | None: + auth: Final = UserAPIKeyAuth(agent_id="caller", user_role="proxy_admin") + auth.billing_agent_policy = caller + await prepare_agent_invocation(auth, "target", None) + return await reserve_budget_for_request( + request_body={"method": "message/send"}, route="/a2a/target", llm_router=None, + valid_token=auth, team_object=None, user_object=None, prisma_client=None, + user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + fail_closed_budget_enforcement=True, + ) + + results: Final = await asyncio.gather(*(admit() for _ in range(8)), return_exceptions=True) + accepted: Final = tuple(result for result in results if isinstance(result, dict)) + rejected: Final = tuple(result for result in results if isinstance(result, litellm.BudgetExceededError)) + assert all(result is None or isinstance(result, (dict, litellm.BudgetExceededError)) for result in results), results + assert len(accepted) == 2, results + assert len(rejected) == 6 + assert await spend_counter_cache.async_get_cache(caller.budget_counter_key) == pytest.approx(0.5) + first: Final = accepted[0] + if outcome == "completed": + await proxy_server.increment_spend_counters( + token=None, team_id=None, user_id=None, response_cost=0.25, + billing_agent_id=caller.agent_id, billing_agent_counter_key=caller.budget_counter_key, + budget_reservation=first, + ) + await reconcile_budget_reservation(first, actual_cost=0.25) + assert await spend_counter_cache.async_get_cache(caller.budget_counter_key) == pytest.approx(0.5) + with pytest.raises(litellm.BudgetExceededError): + await admit() + else: + release: Final = release_budget_reservation_on_cancel if outcome == "cancelled" else release_budget_reservation + await release(first) + await release(first) + assert await spend_counter_cache.async_get_cache(caller.budget_counter_key) == pytest.approx(0.25) + assert await admit() is not None + assert await spend_counter_cache.async_get_cache(caller.budget_counter_key) == pytest.approx(0.5) diff --git a/tests/unit/types/test_agents.py b/tests/unit/types/test_agents.py index a583ab9e8a2..9acb7c41b2a 100644 --- a/tests/unit/types/test_agents.py +++ b/tests/unit/types/test_agents.py @@ -26,3 +26,16 @@ def test_different_budget_windows_never_share_a_settlement_filter() -> None: assert agent_spend_filter(agent_budget_counter_key("agent-one", first)) != agent_spend_filter( agent_budget_counter_key("agent-one", second) ) + + +def test_lifetime_budget_consumption_is_separate_from_agent_history() -> None: + from litellm.types.agents import AgentResponse + + agent: Final = AgentResponse( + agent_id="agent", agent_name="Agent", agent_card_params={}, spend=12.5, + lifetime_budget_spend=0.75, + litellm_budget_table={"budget_id": "budget", "max_budget": 1.0}, + ) + assert agent.budget_spend == 0.75 + assert agent.spend == 12.5 + assert agent.budget_counter_key == "spend:agent_lifetime:budget:agent" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.test.ts index 10ca2114186..f54d6364ddc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest"; import { + agentBudgetSpend, buildIdentityParams, entraTenantFromIssuer, parseIdentityForForm, @@ -84,3 +85,25 @@ describe("agent identity configuration", () => { expect(entraTenantFromIssuer("https://login.microsoftonline.com/common/v2.0")).toBeNull(); }); }); + +describe("agent budget consumption", () => { + it("keeps historical spend separate from lifetime budget consumption", () => { + expect( + agentBudgetSpend({ + spend: 12.5, + lifetime_budget_spend: 0.75, + litellm_budget_table: { budget_id: "budget", max_budget: 1 }, + }), + ).toBe(0.75); + }); + it("preserves recurring window spend and defaults missing lifetime consumption to zero", () => { + expect( + agentBudgetSpend({ + spend: 0.5, + lifetime_budget_spend: 9, + litellm_budget_table: { budget_id: "budget", max_budget: 1, budget_duration: "1d" }, + }), + ).toBe(0.5); + expect(agentBudgetSpend({ spend: 12.5, litellm_budget_table: { budget_id: "budget", max_budget: 1 } })).toBe(0); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.ts index 2db4f97e128..24f25a54f65 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.ts @@ -100,23 +100,26 @@ export const withAgentIdentity = ( const hasCard = !existing || cardEdited || Object.keys(existing.agent_card_params ?? {}).length > 0; const identityFields = buildIdentityParams(values, existing?.identity); const managed = values.identity_provider === "microsoft_entra" || Boolean(readAgentIdentity(existing?.identity)); - const budgetIsSet = - values.agent_max_budget !== undefined && values.agent_max_budget !== "" && values.agent_max_budget !== null; - const budgetWasSet = existing?.litellm_budget_table?.max_budget != null; return { ...settings, ...(hasCard && agent_card_params ? { agent_card_params } : {}), ...identityFields, ...(managed && values.execution_mode !== undefined ? { execution_mode: values.execution_mode } : {}), ...(managed && values.enabled !== undefined ? { enabled: values.enabled } : {}), - ...(budgetIsSet - ? { - budget: { - max_budget: Number(values.agent_max_budget), - budget_duration: values.agent_budget_duration || null, - }, - } - : {}), - ...(!budgetIsSet && budgetWasSet && values.agent_max_budget !== undefined ? { budget: null } : {}), + ...(values.agent_max_budget === undefined + ? {} + : { + budget: + values.agent_max_budget === "" || values.agent_max_budget === null + ? null + : { max_budget: Number(values.agent_max_budget), budget_duration: values.agent_budget_duration || null }, + }), }; }; + +export const agentBudgetSpend = ( + agent: Pick, +): number => + agent.litellm_budget_table && !agent.litellm_budget_table.budget_duration + ? agent.lifetime_budget_spend ?? 0 + : agent.spend ?? 0; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index e08cab776c4..3ba83920589 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -156,6 +156,44 @@ describe("AgentInfoView update payload", () => { .mockResolvedValue({} as never); }); + it.each(["preserve", "edit", "clear"])( + "%s lifetime budget without displaying historical spend as consumption", + async (action) => { + const budgetedAgent = { + ...A2A_AGENT, + spend: 12.5, + lifetime_budget_spend: 0.5, + litellm_budget_table: { budget_id: "budget", max_budget: 1, budget_duration: null }, + }; + vi.mocked(networking.getAgentInfo).mockResolvedValue(budgetedAgent); + const user = setup(); + renderView(); + expect(await screen.findByText("$0.5 / $1")).toBeInTheDocument(); + await openEditor(user); + const limit = screen.getByLabelText("Aggregate Agent Budget ($)"); + expect(limit).toHaveValue(1); + expect(screen.getByLabelText("Budget Reset Period")).toHaveValue(""); + if (action !== "preserve") { + fireEvent.change(limit, { target: { value: action === "edit" ? "2" : "" } }); + } + await save(user); + expect(patchedPayload().budget).toEqual( + action === "clear" ? null : { max_budget: action === "edit" ? 2 : 1, budget_duration: null }, + ); + }, + ); + + it("creates a recurring budget from the budget fields", async () => { + const user = setup(); + renderView(); + expect(await screen.findByText("No aggregate limit")).toBeInTheDocument(); + await openEditor(user); + fireEvent.change(screen.getByLabelText("Aggregate Agent Budget ($)"), { target: { value: "0.75" } }); + fireEvent.change(screen.getByLabelText("Budget Reset Period"), { target: { value: "1d" } }); + await save(user); + expect(patchedPayload().budget).toEqual({ max_budget: 0.75, budget_duration: "1d" }); + }); + it.each([ { card: "complete", editCard: false }, { card: "empty", editCard: false }, @@ -222,7 +260,8 @@ describe("AgentInfoView update payload", () => { await save(user); - expect(patchedPayload()).toEqual({ + const expectedPayload = { + budget: null, agent_name: "my-agent", agent_card_params: { protocolVersion: "1.0", @@ -241,7 +280,8 @@ describe("AgentInfoView update payload", () => { session_rpm_limit: 444, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, access_group_ids: [], - }); + }; + expect(patchedPayload()).toEqual(expectedPayload); }); it("sends the loaded values of every panel the user opens", async () => { @@ -259,7 +299,8 @@ describe("AgentInfoView update payload", () => { await save(user); - expect(patchedPayload()).toEqual({ + const expectedPayload = { + budget: null, agent_name: "my-agent", agent_card_params: { protocolVersion: "1.0", @@ -283,7 +324,8 @@ describe("AgentInfoView update payload", () => { session_rpm_limit: 444, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, access_group_ids: [], - }); + }; + expect(patchedPayload()).toEqual(expectedPayload); }); it("clamps a rate limit typed below its minimum up to that minimum", async () => { @@ -342,7 +384,8 @@ describe("AgentInfoView update payload", () => { await save(user); - expect(patchedPayload()).toEqual({ + const expectedPayload = { + budget: null, agent_name: "lg-agent", agent_card_params: { protocolVersion: "1.0", @@ -362,7 +405,8 @@ describe("AgentInfoView update payload", () => { }, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, access_group_ids: [], - }); + }; + expect(patchedPayload()).toEqual(expectedPayload); }); it("keeps the agent's existing MCP grants in the update payload", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index 24fd7e9274f..491496b8a07 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -1,6 +1,6 @@ import { AgentIdentityFields } from "./AgentIdentityFields"; import { AgentIdentityDetails } from "./AgentIdentityDetails"; -import { withAgentIdentity } from "./agent_identity"; +import { agentBudgetSpend, withAgentIdentity } from "./agent_identity"; import React, { useState, useEffect, useMemo } from "react"; import { cx } from "@/lib/cva.config"; import { FormProvider, useForm, useWatch } from "react-hook-form"; @@ -82,7 +82,7 @@ const AgentBudgetDetails = ({ agent }: { agent: Agent }) => ( <> {agent.litellm_budget_table?.max_budget != null - ? `$${agent.spend ?? 0} / $${agent.litellm_budget_table.max_budget}` + ? `$${agentBudgetSpend(agent)} / $${agent.litellm_budget_table.max_budget}` : "No aggregate limit"} diff --git a/ui/litellm-dashboard/src/components/agents/types.ts b/ui/litellm-dashboard/src/components/agents/types.ts index 703d6170764..2df1c53333c 100644 --- a/ui/litellm-dashboard/src/components/agents/types.ts +++ b/ui/litellm-dashboard/src/components/agents/types.ts @@ -33,6 +33,7 @@ export interface Agent { kill_switch?: AgentKillSwitchConfig | null; keys?: AgentAttachedKey[] | null; spend?: number; + lifetime_budget_spend?: number; tpm_limit?: number | null; rpm_limit?: number | null; session_tpm_limit?: number | null; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a44326e4a7c..f199f132812 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25016,6 +25016,11 @@ export interface components { /** Keys */ keys?: components["schemas"]["AgentKeySummary"][] | null; kill_switch?: components["schemas"]["AgentKillSwitchConfig"] | null; + /** + * Lifetime Budget Spend + * @default 0 + */ + lifetime_budget_spend: number; litellm_budget_table?: components["schemas"]["AgentBudgetState"] | null; /** Litellm Params */ litellm_params?: {