Budget¶

async def record(self, cost_record: CostRecord) -> None:
    """Append a cost record.

    Args:
        cost_record: Immutable cost record to store.
    """
    async with self._lock:
        self._records.append(cost_record)
        logger.info(
            BUDGET_RECORD_ADDED,
            agent_id=cost_record.agent_id,
            model=cost_record.model,
            cost_usd=cost_record.cost_usd,
        )

async def prune_expired(self, *, now: datetime | None = None) -> int:
    """Remove records older than the 168-hour (7-day) cost window.

    Call periodically from long-running services to bound
    memory growth.

    Args:
        now: Reference time.  Defaults to current UTC time.

    Returns:
        Number of records removed.
    """
    ref = now or datetime.now(UTC)
    cutoff = ref - timedelta(hours=_COST_WINDOW_HOURS)
    async with self._lock:
        pruned = self._prune_before(cutoff)
        if pruned:
            logger.info(
                BUDGET_RECORDS_PRUNED,
                pruned=pruned,
                remaining=len(self._records),
            )
        return pruned

async def get_total_cost(
    self,
    *,
    start: datetime | None = None,
    end: datetime | None = None,
) -> float:
    """Sum of ``cost_usd`` across all records, optionally filtered by time.

    Args:
        start: Inclusive lower bound on ``timestamp``.
        end: Exclusive upper bound on ``timestamp``.

    Returns:
        Rounded total cost in USD (base currency).

    Raises:
        ValueError: If both *start* and *end* are given and
            ``start >= end``.
    """
    _validate_time_range(start, end)
    logger.debug(BUDGET_TOTAL_COST_QUERIED, start=start, end=end)
    snapshot = await self._snapshot()
    filtered = _filter_records(snapshot, start=start, end=end)
    return _aggregate(filtered).cost

async def get_agent_cost(
    self,
    agent_id: str,
    *,
    start: datetime | None = None,
    end: datetime | None = None,
) -> float:
    """Sum of ``cost_usd`` for a single agent, optionally filtered by time.

    Args:
        agent_id: Agent identifier to filter by.
        start: Inclusive lower bound on ``timestamp``.
        end: Exclusive upper bound on ``timestamp``.

    Returns:
        Rounded total cost in USD (base currency) for the agent.

    Raises:
        ValueError: If both *start* and *end* are given and
            ``start >= end``.
    """
    _validate_time_range(start, end)
    logger.debug(
        BUDGET_AGENT_COST_QUERIED,
        agent_id=agent_id,
        start=start,
        end=end,
    )
    snapshot = await self._snapshot()
    filtered = _filter_records(
        snapshot,
        agent_id=agent_id,
        start=start,
        end=end,
    )
    return _aggregate(filtered).cost

async def get_record_count(self) -> int:
    """Total number of recorded cost entries.

    Returns:
        Number of cost records.
    """
    async with self._lock:
        return len(self._records)

async def get_records(
    self,
    *,
    agent_id: str | None = None,
    task_id: str | None = None,
    provider: NotBlankStr | None = None,
    start: datetime | None = None,
    end: datetime | None = None,
) -> tuple[CostRecord, ...]:
    """Return filtered cost records.

    Returns an immutable snapshot of records matching the filters.

    Args:
        agent_id: Filter by agent.
        task_id: Filter by task.
        provider: Filter by provider name.
        start: Inclusive lower bound on ``timestamp``.
        end: Exclusive upper bound on ``timestamp``.

    Returns:
        Immutable tuple of matching cost records.

    Raises:
        ValueError: If both *start* and *end* are given and
            ``start >= end``.
    """
    _validate_time_range(start, end)
    logger.debug(
        BUDGET_RECORDS_QUERIED,
        agent_id=agent_id,
        task_id=task_id,
        provider=provider,
        start=start,
        end=end,
    )
    snapshot = await self._snapshot()
    return _filter_records(
        snapshot,
        agent_id=agent_id,
        task_id=task_id,
        provider=provider,
        start=start,
        end=end,
    )

async def get_provider_usage(
    self,
    provider_name: NotBlankStr,
    *,
    start: datetime | None = None,
    end: datetime | None = None,
) -> ProviderUsageSummary:
    """Return aggregated token and cost totals for a provider.

    Args:
        provider_name: Provider to aggregate usage for.
        start: Inclusive lower bound on ``timestamp``.
        end: Exclusive upper bound on ``timestamp``.

    Returns:
        Total tokens (input + output) and total cost.

    Raises:
        ValueError: If both *start* and *end* are given and
            ``start >= end``.
    """
    _validate_time_range(start, end)
    logger.debug(
        BUDGET_PROVIDER_USAGE_QUERIED,
        provider=provider_name,
        start=start,
        end=end,
    )
    snapshot = await self._snapshot()
    filtered = _filter_records(
        snapshot,
        provider=provider_name,
        start=start,
        end=end,
    )
    if not filtered:
        return ProviderUsageSummary(total_tokens=0, total_cost=0.0)
    agg = _aggregate(filtered)
    return ProviderUsageSummary(
        total_tokens=agg.input_tokens + agg.output_tokens,
        total_cost=agg.cost,
    )

async def build_summary(
    self,
    *,
    start: datetime,
    end: datetime,
) -> SpendingSummary:
    """Build a spending summary for the given period.

    Args:
        start: Inclusive period start.
        end: Exclusive period end.

    Returns:
        Aggregated spending summary with breakdowns and alert level.

    Raises:
        ValueError: If ``start >= end``.
    """
    _validate_time_range(start, end)
    retention_cutoff = datetime.now(UTC) - timedelta(
        hours=_COST_WINDOW_HOURS,
    )
    if start < retention_cutoff:
        logger.warning(
            BUDGET_QUERY_EXCEEDS_RETENTION,
            requested_start=start.isoformat(),
            retention_cutoff=retention_cutoff.isoformat(),
            retention_hours=_COST_WINDOW_HOURS,
        )
    snapshot = await self._snapshot()
    filtered = _filter_records(snapshot, start=start, end=end)
    totals = _aggregate(filtered)

    agent_spendings = _build_agent_spendings(filtered)
    dept_spendings = self._build_dept_spendings(agent_spendings)
    budget_monthly, used_pct, alert = self._build_budget_context(
        totals.cost,
    )

    summary = SpendingSummary(
        period=PeriodSpending(
            start=start,
            end=end,
            total_cost_usd=totals.cost,
            total_input_tokens=totals.input_tokens,
            total_output_tokens=totals.output_tokens,
            record_count=totals.record_count,
        ),
        by_agent=tuple(agent_spendings),
        by_department=tuple(dept_spendings),
        budget_total_monthly=budget_monthly,
        budget_used_percent=used_pct,
        alert_level=alert,
    )

    logger.info(
        BUDGET_SUMMARY_BUILT,
        total_cost_usd=totals.cost,
        record_count=totals.record_count,
        agent_count=len(agent_spendings),
        department_count=len(dept_spendings),
        alert_level=alert.value,
    )

    return summary

async def get_category_breakdown(
    self,
    *,
    agent_id: str | None = None,
    task_id: str | None = None,
    start: datetime | None = None,
    end: datetime | None = None,
) -> CategoryBreakdown:
    """Build a per-category cost breakdown.

    Args:
        agent_id: Filter by agent.
        task_id: Filter by task.
        start: Inclusive lower bound on timestamp.
        end: Exclusive upper bound on timestamp.

    Returns:
        Category breakdown of cost, tokens, and call counts.

    Raises:
        ValueError: If ``start >= end``.
    """
    _validate_time_range(start, end)
    logger.debug(
        BUDGET_CATEGORY_BREAKDOWN_QUERIED,
        agent_id=agent_id,
        task_id=task_id,
        start=start,
        end=end,
    )
    snapshot = await self._snapshot()
    filtered = _filter_records(
        snapshot,
        agent_id=agent_id,
        task_id=task_id,
        start=start,
        end=end,
    )
    return build_category_breakdown(filtered)

async def get_orchestration_ratio(
    self,
    *,
    agent_id: str | None = None,
    task_id: str | None = None,
    start: datetime | None = None,
    end: datetime | None = None,
    thresholds: OrchestrationAlertThresholds | None = None,
) -> OrchestrationRatio:
    """Compute the orchestration overhead ratio.

    Args:
        agent_id: Filter by agent.
        task_id: Filter by task.
        start: Inclusive lower bound on timestamp.
        end: Exclusive upper bound on timestamp.
        thresholds: Optional custom alert thresholds.

    Returns:
        Orchestration ratio with alert level.

    Raises:
        ValueError: If ``start >= end``.
    """
    breakdown = await self.get_category_breakdown(
        agent_id=agent_id,
        task_id=task_id,
        start=start,
        end=end,
    )
    result = compute_orchestration_ratio(
        breakdown,
        thresholds=thresholds,
    )
    logger.debug(
        BUDGET_ORCHESTRATION_RATIO_QUERIED,
        agent_id=agent_id,
        task_id=task_id,
        ratio=result.ratio,
        alert_level=result.alert_level.value,
    )
    if result.alert_level != OrchestrationAlertLevel.NORMAL:
        logger.warning(
            BUDGET_ORCHESTRATION_RATIO_ALERT,
            agent_id=agent_id,
            task_id=task_id,
            ratio=result.ratio,
            alert_level=result.alert_level.value,
        )
    return result

async def get_budget_utilization_pct(self) -> float | None:
    """Return monthly budget utilization as a percentage (0--100+).

    Returns ``None`` when disabled (``total_monthly <= 0``) or
    when the cost query fails (graceful degradation).
    """
    cfg = self._budget_config
    if cfg.total_monthly <= 0:
        return None
    try:
        period_start = billing_period_start(cfg.reset_day)
        monthly_cost = await self._cost_tracker.get_total_cost(
            start=period_start,
        )
    except MemoryError, RecursionError:
        raise
    except Exception:
        logger.exception(
            BUDGET_UTILIZATION_ERROR,
            reason="falling_back_to_none",
        )
        return None
    else:
        pct = monthly_cost / cfg.total_monthly * 100
        logger.debug(
            BUDGET_UTILIZATION_QUERIED,
            monthly_cost=monthly_cost,
            total_monthly=cfg.total_monthly,
            utilization_pct=pct,
        )
        return pct