feat(agents): define persisted budget windows and accounting keys

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

@@ -0,0 +1,15 @@
+ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "budget_id" TEXT;
+
+-- CreateIndex
+CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentsTable_budget_id_key" ON "LiteLLM_AgentsTable"("budget_id");
+
+-- AddForeignKey
+DO $$
+BEGIN
+    IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_AgentsTable_budget_id_fkey') THEN
+        ALTER TABLE "LiteLLM_AgentsTable" ADD CONSTRAINT "LiteLLM_AgentsTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
+    END IF;
+END $$;
+
+
+ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "spend_window" TIMESTAMP(3);

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -36,6 +36,7 @@ model LiteLLM_BudgetTable {
   model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget
   team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team 
   organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
+  agents LiteLLM_AgentsTable[]
 }
 
 // Models on proxy
@@ -83,6 +84,9 @@ model LiteLLM_AgentsTable {
   execution_mode String @default("autonomous")
   identity LiteLLM_AgentIdentity?
   retired_identities LiteLLM_RetiredAgentIdentity[]
+  budget_id String? @unique
+  spend_window        DateTime?
+  litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
   tpm_limit  Int?
   rpm_limit  Int?
   session_tpm_limit Int?

litellm/proxy/schema.prisma

@@ -36,6 +36,7 @@ model LiteLLM_BudgetTable {
   model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget
   team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team 
   organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
+  agents LiteLLM_AgentsTable[]
 }
 
 // Models on proxy
@@ -83,6 +84,9 @@ model LiteLLM_AgentsTable {
   execution_mode String @default("autonomous")
   identity LiteLLM_AgentIdentity?
   retired_identities LiteLLM_RetiredAgentIdentity[]
+  budget_id String? @unique
+  spend_window        DateTime?
+  litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
   tpm_limit  Int?
   rpm_limit  Int?
   session_tpm_limit Int?

litellm/types/agents.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping, Sequence
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias
 from urllib.parse import urlsplit
 
@@ -15,6 +15,7 @@ from litellm.types.proxy.agent_identity import (
 
 if TYPE_CHECKING:
     from a2a.types import SendMessageResponse
+    from prisma.types import LiteLLM_AgentsTableWhereInput
 
 
 # AgentProvider
@@ -311,6 +312,27 @@ class AgentKeySummary(BaseModel):
     key_name: str | None = None
 
 
+def agent_budget_counter_key(agent_id: str, reset_at: datetime | None) -> str:
+    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)
+    window: Final = aware.astimezone(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ")
+    return f"spend:agent_window:{window}:{agent_id}"
+
+
+def agent_spend_filter(counter_key: str) -> "LiteLLM_AgentsTableWhereInput":
+    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)
+        windowed: Final[LiteLLM_AgentsTableWhereInput] = {"agent_id": agent_id, "spend_window": window}
+        return windowed
+    cumulative: Final[LiteLLM_AgentsTableWhereInput] = {
+        "agent_id": counter_key.removeprefix("spend:agent:"),
+        "spend_window": None,
+    }
+    return cumulative
+
+
 class AgentResponse(BaseModel):
     identity: AgentIdentityBinding | None = None
     identity_managed: bool = False

litellm/types/proxy/agent_identity.py

@@ -46,6 +46,22 @@ class AgentIdentityBinding(BaseModel):
     last_authenticated_at: datetime | None = None
 
 
+class AgentBudgetConfig(BaseModel):
+    model_config = ConfigDict(frozen=True, extra="forbid")
+
+    max_budget: float = Field(ge=0, allow_inf_nan=False)
+    budget_duration: str | None = None
+
+
+class AgentBudgetState(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    budget_id: str
+    max_budget: float | None = None
+    budget_duration: str | None = None
+    budget_reset_at: datetime | None = None
+
+
 class AgentSubject(BaseModel):
     model_config = ConfigDict(frozen=True)
 

schema.prisma

@@ -36,6 +36,7 @@ model LiteLLM_BudgetTable {
   model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget
   team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team 
   organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
+  agents LiteLLM_AgentsTable[]
 }
 
 // Models on proxy
@@ -83,6 +84,9 @@ model LiteLLM_AgentsTable {
   execution_mode String @default("autonomous")
   identity LiteLLM_AgentIdentity?
   retired_identities LiteLLM_RetiredAgentIdentity[]
+  budget_id String? @unique
+  spend_window        DateTime?
+  litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
   tpm_limit  Int?
   rpm_limit  Int?
   session_tpm_limit Int?

tests/unit/types/proxy/test_agent_identity.py

@@ -0,0 +1,20 @@
+import pytest
+from pydantic import ValidationError
+
+from litellm.types.proxy.agent_identity import AgentBudgetConfig
+
+
+@pytest.mark.parametrize("amount", [-1, float("inf"), float("-inf"), float("nan")])
+def test_agent_budget_rejects_negative_or_nonfinite_caps(amount: float) -> None:
+    with pytest.raises(ValidationError):
+        AgentBudgetConfig(max_budget=amount)
+
+
+@pytest.mark.parametrize("amount", [0, 0.01, 100])
+def test_agent_budget_preserves_a_finite_nonnegative_cap(amount: float) -> None:
+    assert AgentBudgetConfig(max_budget=amount).max_budget == amount
+
+
+def test_agent_budget_rejects_unknown_policy_fields() -> None:
+    with pytest.raises(ValidationError):
+        AgentBudgetConfig.model_validate({"max_budget": 1, "unknown_control": True})

tests/unit/types/test_agents.py

@@ -0,0 +1,28 @@
+from datetime import datetime, timedelta, timezone
+from typing import Final
+
+import pytest
+
+from litellm.types.agents import agent_budget_counter_key, agent_spend_filter
+
+
+@pytest.mark.parametrize("offset", [None, timezone.utc, timezone(timedelta(hours=5, minutes=30))])
+def test_agent_window_key_round_trip_preserves_the_admitted_instant(offset) -> None:
+    instant: Final = datetime(2030, 1, 1, 12, 30, 1, 123000, tzinfo=offset)
+    expected: Final = instant.replace(tzinfo=timezone.utc) if offset is None else instant.astimezone(timezone.utc)
+    key: Final = agent_budget_counter_key("agent:with:colons", instant)
+    assert agent_spend_filter(key) == {"agent_id": "agent:with:colons", "spend_window": expected}
+    assert key == agent_budget_counter_key("agent:with:colons", expected)
+
+
+def test_unbudgeted_key_is_filtered_to_an_unbudgeted_row() -> None:
+    key: Final = agent_budget_counter_key("agent-one", None)
+    assert agent_spend_filter(key) == {"agent_id": "agent-one", "spend_window": None}
+
+
+def test_different_budget_windows_never_share_a_settlement_filter() -> None:
+    first: Final = datetime(2030, 1, 1, tzinfo=timezone.utc)
+    second: Final = first + timedelta(days=1)
+    assert agent_spend_filter(agent_budget_counter_key("agent-one", first)) != agent_spend_filter(
+        agent_budget_counter_key("agent-one", second)
+    )