agent_budget.rego

Caps how much a session can consume — three independent budgets you can mix and match: tokens, API calls, cost (cents). When no budget is set, the policy short-circuits to allow.

Package: kitelogik.agent_budget · Source: kitelogik/policies/agent_budget.rego

Budget fields on SessionContext

Total	Used	Triggers deny when
budget_total_tokens	budget_used_tokens	used >= total
budget_total_api_calls	budget_used_api_calls	used >= total
budget_total_cost_cents	budget_used_cost_cents	used >= total

Each pair is independent. A budget is "set" when the total_* field is non-null. A budget is "exhausted" when both total_* and used_* are non-null and used_* >= total_*.

Rules

agent.budget events

Outcome	Trigger
Allow	No budget set (all total_* are null)
Allow	All set budgets are within limits
Deny	Any of: token / API-call / cost budget exhausted

tool_call events (opportunistic enforcement)

Outcome	Trigger
Deny	Token / API-call / cost budget exhausted, regardless of which tool is being called

The tool_call deny rules run on every governed tool invocation — this means you don't have to fire explicit agent.budget events to get budget enforcement. As long as your runtime maintains the budget_used_* fields on SessionContext, every tool call is gated.

Selected source

# Allow when no budget is set
allow if {
    input.event_type == "agent.budget"
    not _has_any_budget
}

# Allow when all set budgets are within limits
allow if {
    input.event_type == "agent.budget"
    _has_any_budget
    not _token_budget_exhausted
    not _api_call_budget_exhausted
    not _cost_budget_exhausted
}

# Opportunistic enforcement on every tool call
deny if {
    input.event_type == "tool_call"
    _token_budget_exhausted
}

_token_budget_exhausted if {
    input.context.budget_total_tokens != null
    input.context.budget_used_tokens != null
    input.context.budget_used_tokens >= input.context.budget_total_tokens
}

How main.rego uses this module

Two routing paths:

# Explicit agent.budget events
allow if { not deny; input.event_type == "agent.budget"; agent_budget.allow }
deny if { input.event_type == "agent.budget"; agent_budget.deny }

# Budget enforcement on tool_call events (deny if budget exhausted)
deny if { agent_budget.deny }

The second deny if { agent_budget.deny } is what makes per-tool-call enforcement work without needing an explicit agent.budget event between every call.

Setting budgets on SessionContext

from kitelogik import SessionContext

context = SessionContext(
    session_id="sess_001",
    user_role="support_agent",
    session_scopes=["read_customer"],
    budget_total_tokens=100_000,
    budget_used_tokens=0,
    budget_total_api_calls=200,
    budget_used_api_calls=0,
    # cost cents omitted — no cost budget for this session
)

After every model interaction, your runtime updates the used_* counters before the next gate evaluation. Producing a derived context with model_copy() is the cleanest way:

context = context.model_copy(update={
    "budget_used_tokens": context.budget_used_tokens + tokens_this_call,
    "budget_used_api_calls": (context.budget_used_api_calls or 0) + 1,
})

Extending in your own project

Add tighter caps without touching the OSS module:

package kitelogik.agent_budget

import future.keywords.if

# Hard cap: no more than 50K tokens per session, regardless of total
deny if {
    input.context.budget_used_tokens != null
    input.context.budget_used_tokens > 50000
}

# Cap cost at $5 for low-trust roles
deny if {
    input.context.user_role == "guest"
    input.context.budget_used_cost_cents != null
    input.context.budget_used_cost_cents > 500
}

Same package, denies merge across files — the stricter rule wins.