symfonic.core.observability.budget

budget

TokenBudgetTracker — per-tenant cost ceilings + usage aggregation.

Gate 2 of v5.5.0; Sprint B (v5.6.1) adds multi-replica hydration via an optional :class:BudgetStore. In-memory-first on the hot path with an asyncio.Lock around record() for correctness under asyncio.gather. Persistence is optional; without a store the tracker behaves exactly like pre-5.5 agents. Data types live in :mod:.budget_models and hydration helpers in :mod:.budget_hydrate so each file stays below the 300-LOC cap.

BudgetCheckResult dataclass

BudgetCheckResult(
    allowed: bool,
    reason: str | None = None,
    daily_used_usd: float = 0.0,
    daily_limit_usd: float | None = None,
    monthly_used_usd: float = 0.0,
    monthly_limit_usd: float | None = None,
)

Result of a pre-call budget check.

BudgetStore

Bases: Protocol

Optional persistence layer for daily aggregates.

Implementations should be idempotent on conflict — the tracker calls upsert_daily after every LLM call, so the underlying storage must handle UPSERT semantics (ON CONFLICT ... DO UPDATE).

The optional read_today / read_month methods hydrate the tracker's in-memory buckets on demand so that multiple app replicas converge on the same view of a tenant's cumulative usage. A store that only implements upsert_daily still works (writes only) — multi-replica correctness requires all three.

read_month async

read_month(
    tenant_id: str, year_month: str
) -> BudgetUsage | None

Return the aggregated month-to-date usage for tenant_id.

year_month is the YYYY-MM key produced by :func:month_key. Optional — see :meth:read_today.

Source code in src/symfonic/core/observability/budget_models.py

async def read_month(  # pragma: no cover
    self, tenant_id: str, year_month: str,
) -> BudgetUsage | None:
    """Return the aggregated month-to-date usage for ``tenant_id``.

    ``year_month`` is the ``YYYY-MM`` key produced by
    :func:`month_key`.  Optional — see :meth:`read_today`.
    """
    ...

read_today async

read_today(tenant_id: str) -> BudgetUsage | None

Return today's usage row for tenant_id (None if absent).

Optional — implementations without this method get write-through semantics only (single-process correctness).

Source code in src/symfonic/core/observability/budget_models.py

async def read_today(  # pragma: no cover
    self, tenant_id: str,
) -> BudgetUsage | None:
    """Return today's usage row for ``tenant_id`` (None if absent).

    Optional — implementations without this method get write-through
    semantics only (single-process correctness).
    """
    ...

BudgetUsage dataclass

BudgetUsage(
    tenant_id: str,
    window: str,
    window_key: str,
    input_tokens: int = 0,
    output_tokens: int = 0,
    cached_tokens: int = 0,
    cache_write_tokens: int = 0,
    cost_usd: float = 0.0,
    request_count: int = 0,
    updated_at: datetime = (lambda: datetime.now(UTC))(),
)

Cumulative usage for a single (tenant, window) pair.

TokenBudgetTracker

TokenBudgetTracker(
    *,
    default_daily_budget_usd: float | None = None,
    default_monthly_budget_usd: float | None = None,
    store: BudgetStore | None = None,
)

Tracks cumulative token usage + cost per tenant.

Two in-memory buckets per tenant (day, month) roll over lazily on date/month boundaries. Limits configured globally via default_*_budget_usd or per-tenant via meth:set_tenant_budget; None means unlimited. Admin callers bypass the breaker.

Source code in src/symfonic/core/observability/budget.py

def __init__(
    self,
    *,
    default_daily_budget_usd: float | None = None,
    default_monthly_budget_usd: float | None = None,
    store: BudgetStore | None = None,
) -> None:
    self._default_daily = default_daily_budget_usd
    self._default_monthly = default_monthly_budget_usd
    self._store = store

    # {tenant_id: BudgetUsage}  — always keyed on the CURRENT window.
    # Rollover is lazy: if the window key no longer matches, a fresh
    # BudgetUsage is created on first access.
    self._daily: dict[str, BudgetUsage] = {}
    self._monthly: dict[str, BudgetUsage] = {}

    # Per-tenant overrides.  None value = explicit "unlimited".
    self._tenant_daily: dict[str, float | None] = {}
    self._tenant_monthly: dict[str, float | None] = {}

    # Lazy per-loop lock — binding at __init__ breaks mixed-loop tests.
    self._lock: asyncio.Lock | None = None
    self._lock_loop: Any = None

check_budget async

check_budget(
    tenant_id: str, *, is_admin: bool = False
) -> BudgetCheckResult

Return whether tenant_id may incur another LLM call.

When a BudgetStore with read_today / read_month is registered, the in-memory buckets are first hydrated from the store so every replica shares one view of the running total. Store errors are logged at WARNING and fall through to in-memory state — the ledger is never allowed to break auth. Admin callers bypass the breaker.

Source code in src/symfonic/core/observability/budget.py

async def check_budget(
    self,
    tenant_id: str,
    *,
    is_admin: bool = False,
) -> BudgetCheckResult:
    """Return whether ``tenant_id`` may incur another LLM call.

    When a ``BudgetStore`` with ``read_today`` / ``read_month`` is
    registered, the in-memory buckets are first hydrated from the
    store so every replica shares one view of the running total.
    Store errors are logged at WARNING and fall through to in-memory
    state — the ledger is never allowed to break auth.  Admin
    callers bypass the breaker.
    """
    from symfonic.core.observability.budget_hydrate import (
        hydrate_from_store,
    )
    await hydrate_from_store(self, tenant_id)

    day = await self.usage_today(tenant_id)
    month = await self.usage_this_month(tenant_id)
    daily_limit, monthly_limit = self.get_tenant_limits(tenant_id)

    if is_admin:
        return BudgetCheckResult(
            allowed=True,
            daily_used_usd=day.cost_usd,
            daily_limit_usd=daily_limit,
            monthly_used_usd=month.cost_usd,
            monthly_limit_usd=monthly_limit,
        )

    if daily_limit is not None and day.cost_usd >= daily_limit:
        return BudgetCheckResult(
            allowed=False,
            reason=(
                f"daily limit ${daily_limit:.2f} reached "
                f"(used: ${day.cost_usd:.4f})"
            ),
            daily_used_usd=day.cost_usd,
            daily_limit_usd=daily_limit,
            monthly_used_usd=month.cost_usd,
            monthly_limit_usd=monthly_limit,
        )
    if monthly_limit is not None and month.cost_usd >= monthly_limit:
        return BudgetCheckResult(
            allowed=False,
            reason=(
                f"monthly limit ${monthly_limit:.2f} reached "
                f"(used: ${month.cost_usd:.4f})"
            ),
            daily_used_usd=day.cost_usd,
            daily_limit_usd=daily_limit,
            monthly_used_usd=month.cost_usd,
            monthly_limit_usd=monthly_limit,
        )

    return BudgetCheckResult(
        allowed=True,
        daily_used_usd=day.cost_usd,
        daily_limit_usd=daily_limit,
        monthly_used_usd=month.cost_usd,
        monthly_limit_usd=monthly_limit,
    )

clear_tenant_budget

clear_tenant_budget(tenant_id: str) -> None

Remove per-tenant overrides so the defaults apply again.

Source code in src/symfonic/core/observability/budget.py

def clear_tenant_budget(self, tenant_id: str) -> None:
    """Remove per-tenant overrides so the defaults apply again."""
    self._tenant_daily.pop(tenant_id, None)
    self._tenant_monthly.pop(tenant_id, None)

get_tenant_limits

get_tenant_limits(
    tenant_id: str,
) -> tuple[float | None, float | None]

Return (daily_limit, monthly_limit) applied to tenant_id.

Source code in src/symfonic/core/observability/budget.py

def get_tenant_limits(
    self, tenant_id: str,
) -> tuple[float | None, float | None]:
    """Return ``(daily_limit, monthly_limit)`` applied to ``tenant_id``."""
    daily = self._tenant_daily.get(tenant_id, self._default_daily) \
        if tenant_id in self._tenant_daily else self._default_daily
    monthly = self._tenant_monthly.get(tenant_id, self._default_monthly) \
        if tenant_id in self._tenant_monthly else self._default_monthly
    return daily, monthly

record async

record(
    tenant_id: str, model: str, usage: dict[str, int]
) -> float

Record token usage for a single LLM call; returns its cost_usd.

Source code in src/symfonic/core/observability/budget.py

async def record(
    self,
    tenant_id: str,
    model: str,
    usage: dict[str, int],
) -> float:
    """Record token usage for a single LLM call; returns its cost_usd."""
    input_tokens = int(usage.get("input_tokens", 0) or 0)
    output_tokens = int(usage.get("output_tokens", 0) or 0)
    cached = int(usage.get("cached_tokens", 0) or 0)
    cache_write = int(
        usage.get("cache_write_tokens")
        or usage.get("cache_creation_tokens")
        or 0,
    )
    cost = calculate_cost(
        model, input_tokens, output_tokens, cached, cache_write,
    )

    async with self._get_lock():
        day = self._bucket(tenant_id, "day")
        month = self._bucket(tenant_id, "month")
        for b in (day, month):
            b.input_tokens += input_tokens
            b.output_tokens += output_tokens
            b.cached_tokens += cached
            b.cache_write_tokens += cache_write
            b.cost_usd = round(b.cost_usd + cost, 6)
            b.request_count += 1
            b.updated_at = datetime.now(UTC)

        daily_snapshot = copy_usage(day)

    # Persistence happens OUTSIDE the lock — the store may do I/O.
    # Log but don't propagate; telemetry failures never break LLM flow.
    if self._store is not None:
        try:
            await self._store.upsert_daily(daily_snapshot)
        except Exception:  # pragma: no cover - integration tested
            logger.warning(
                "BudgetStore.upsert_daily failed for tenant=%s "
                "(in-memory counters still accurate)",
                tenant_id,
                exc_info=True,
            )
    return cost

reset async

reset(tenant_id: str) -> None

Reset in-memory day + month counters for a tenant.

When a persistent store is wired, the shared Postgres row is NOT cleared — use the admin billing endpoint to zero the ledger.

Source code in src/symfonic/core/observability/budget.py

async def reset(self, tenant_id: str) -> None:
    """Reset in-memory day + month counters for a tenant.

    When a persistent store is wired, the shared Postgres row is NOT
    cleared — use the admin billing endpoint to zero the ledger.
    """
    async with self._get_lock():
        self._daily.pop(tenant_id, None)
        self._monthly.pop(tenant_id, None)

set_tenant_budget

set_tenant_budget(
    tenant_id: str,
    *,
    daily_usd: float | None = None,
    monthly_usd: float | None = None,
) -> None

Override budget ceilings for a tenant; None = unlimited.

Source code in src/symfonic/core/observability/budget.py

def set_tenant_budget(
    self,
    tenant_id: str,
    *,
    daily_usd: float | None = None,
    monthly_usd: float | None = None,
) -> None:
    """Override budget ceilings for a tenant; ``None`` = unlimited."""
    self._tenant_daily[tenant_id] = daily_usd
    self._tenant_monthly[tenant_id] = monthly_usd

get_budget_tracker

get_budget_tracker() -> TokenBudgetTracker | None

Return the currently registered tracker, if any.

Source code in src/symfonic/core/observability/budget.py

def get_budget_tracker() -> TokenBudgetTracker | None:
    """Return the currently registered tracker, if any."""
    return _budget_tracker

set_budget_tracker

set_budget_tracker(
    tracker: TokenBudgetTracker | None,
) -> None

Register (or clear) the process-wide budget tracker.

Consulted by check_budget_before_llm and SymfonicAgent.run(); when unset both skip the check.

Source code in src/symfonic/core/observability/budget.py

def set_budget_tracker(tracker: TokenBudgetTracker | None) -> None:
    """Register (or clear) the process-wide budget tracker.

    Consulted by ``check_budget_before_llm`` and ``SymfonicAgent.run()``;
    when unset both skip the check.
    """
    global _budget_tracker
    _budget_tracker = tracker