scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(agents): reserve caller fees and isolate lifetime budget spend

Browse source
This commit is contained in:
Joshua Valluru 2026-09-30 16:24:27 -07:00
parent 1529494bed
commit 8f6d78d652
31 changed files with 457 additions and 55 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql
View file

@ -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;

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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?

5
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -3246,6 +3246,11 @@
}
]
},
"lifetime_budget_spend": {
"default": 0.0,
"title": "Lifetime Budget Spend",
"type": "number"
},
"litellm_budget_table": {
"anyOf": [
{

9
litellm/proxy/agent_endpoints/agent_registry.py
View file

@ -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]]: ...

9
litellm/proxy/agent_endpoints/auth/managed_authorization.py
View file

@ -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"))

34
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -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")

11
litellm/proxy/agent_endpoints/managed_identity.py
View file

@ -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

31
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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)

32
litellm/proxy/db/spend_counter_reseed.py
View file

@ -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:") :]}

2
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -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(

1
litellm/proxy/schema.prisma
View file

@ -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?

2
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -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,
)

16
litellm/proxy/spend_tracking/spend_counter_batch.py
View file

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

23
litellm/types/agents.py
View file

@ -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]

1
schema.prisma
View file

@ -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?

8
tests/test_litellm/proxy/agent_endpoints/auth/test_managed_authorization.py
View file

@ -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)

2
tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py
View file

@ -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=[],

28
tests/test_litellm/proxy/agent_endpoints/test_endpoints.py
View file

@ -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

43
tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py
View file

@ -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"]

2
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -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,

29
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -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)

14
tests/test_litellm/proxy/db/test_spend_counter_reseed.py
View file

@ -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"})

1
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -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",

77
tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
View file

@ -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)

13
tests/unit/types/test_agents.py
View file

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

23
ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.test.ts
View file

@ -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);
});
});

27
ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_identity.ts
View file

@ -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<import("@/components/agents/types").Agent, "spend" | "lifetime_budget_spend" | "litellm_budget_table">,
): number =>
agent.litellm_budget_table && !agent.litellm_budget_table.budget_duration
? agent.lifetime_budget_spend ?? 0
: agent.spend ?? 0;

56
ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx
View file

@ -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 () => {

4
ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx
View file

@ -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 }) => (
<>
<DetailItem label="Agent Budget">
{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"}
</DetailItem>
<DetailItem label="Budget Reset">

1
ui/litellm-dashboard/src/components/agents/types.ts
View file

@ -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;

5
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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?: {