Budget¶

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

    The in-memory append runs under ``_lock``.  After the lock
    is released, ``_update_project_aggregate`` is awaited to
    update the durable project cost aggregate when the record
    has a ``project_id`` and a repository is configured.
    Aggregate updates are best-effort: failures are logged at
    WARNING but do not affect the in-memory recording.

    When a ``BudgetConfig`` is attached, the incoming record's
    ``currency`` must match ``budget_config.currency``; mismatches
    raise :class:`MixedCurrencyAggregationError` at the ingestion
    boundary so downstream aggregators never see mixed-currency
    data in the first place.

    Idempotency: a bounded LRU of accepted ``cost_record.claim_id``
    values protects against double-bills from JetStream redelivery
    or in-process retries. Repeat submissions are no-ops and emit
    ``BUDGET_RECORD_DEDUPED`` at INFO. Eviction at the LRU
    capacity is best-effort; once a claim ages out, a re-submitted
    record with the same key is treated as fresh.

    Args:
        cost_record: Immutable cost record to store.

    Raises:
        MixedCurrencyAggregationError: If the record's currency
            does not match the configured ``budget.currency``.
    """
    # Currency check first -- it's synchronous, has no in-flight
    # state to roll back, and a mismatch is a hard caller-contract
    # violation that must surface BEFORE we reserve a claim_id
    # slot. Per-project same-currency invariant is enforced by
    # ``ProjectCostAggregateRepository.increment``; no tracker-
    # side pin is required there.
    if (
        self._budget_config is not None
        and cost_record.currency != self._budget_config.currency
    ):
        # Log at WARNING with full record + budget context BEFORE
        # raising so a downstream catch-and-translate cannot hide
        # repeated mismatches from operators.
        logger.warning(
            BUDGET_MIXED_CURRENCY_REJECTED,
            agent_id=cost_record.agent_id,
            task_id=cost_record.task_id,
            project_id=cost_record.project_id,
            record_currency=cost_record.currency,
            budget_currency=self._budget_config.currency,
        )
        msg = (
            f"Record currency {cost_record.currency!r} does not match "
            f"configured budget currency "
            f"{self._budget_config.currency!r}"
        )
        raise MixedCurrencyAggregationError(
            msg,
            currencies=frozenset(
                {
                    cost_record.currency,
                    self._budget_config.currency,
                }
            ),
            agent_id=cost_record.agent_id,
            task_id=cost_record.task_id,
            project_id=cost_record.project_id,
        )

    # Idempotency fast-path with in-flight reservation. ``_lock``
    # protects both ``_inflight_claims`` (set of in-flight reservations)
    # and ``_seen_claims`` (bounded LRU of finalised entries). We
    # check both states under the lock, then either dedupe (if
    # already seen / in flight) or reserve the claim in
    # ``_inflight_claims`` so a concurrent second call sees the
    # entry and dedupes immediately. Keeping in-flight separate
    # from the LRU prevents the capacity trim from popping a
    # still-running reservation, which would let a duplicate
    # slip past the membership check.
    async with self._get_lock():
        if (
            cost_record.claim_id in self._inflight_claims
            or cost_record.claim_id in self._seen_claims
        ):
            if cost_record.claim_id in self._seen_claims:
                self._seen_claims.move_to_end(cost_record.claim_id)
            logger.info(
                BUDGET_RECORD_DEDUPED,
                claim_id=cost_record.claim_id,
                agent_id=cost_record.agent_id,
                task_id=cost_record.task_id,
                provider=cost_record.provider,
                model=cost_record.model,
                cost=cost_record.cost,
            )
            return
        self._inflight_claims.add(cost_record.claim_id)

    # Run the durable aggregate update OUTSIDE the lock -- DB I/O
    # must not block concurrent in-memory readers/writers. Any
    # failure releases the in-flight reservation so a retry with
    # the same claim_id is not falsely deduped.
    try:
        await self._update_project_aggregate(cost_record)
    except BaseException:
        async with self._get_lock():
            self._inflight_claims.discard(cost_record.claim_id)
        raise

    async with self._get_lock():
        # Promote the reservation to a finalised LRU entry under
        # the lock so the membership check above never observes
        # a gap where the claim is in neither set. Eviction only
        # affects ``_seen_claims``, so still-running reservations
        # in ``_inflight_claims`` are untouched.
        self._inflight_claims.discard(cost_record.claim_id)
        self._records.append(cost_record)
        self._seen_claims[cost_record.claim_id] = None
        self._seen_claims.move_to_end(cost_record.claim_id)
        while len(self._seen_claims) > self._claim_lru_capacity:
            self._seen_claims.popitem(last=False)
        logger.info(
            BUDGET_RECORD_ADDED,
            agent_id=cost_record.agent_id,
            model=cost_record.model,
            cost=cost_record.cost,
        )

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 ``self._clock.now()`` so
            ``FakeClock`` injection deterministically controls the
            eviction cutoff.

    Returns:
        Number of records removed.
    """
    ref = now or self._clock.now()
    cutoff = ref - timedelta(hours=_COST_WINDOW_HOURS)
    async with self._get_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`` 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 the configured currency.

    Raises:
        ValueError: If both *start* and *end* are given and
            ``start >= end``.
    """
    query_start = time.perf_counter()
    try:
        _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
    finally:
        record_budget_query(
            query_type="total_cost",
            duration_sec=time.perf_counter() - query_start,
        )

async def get_agent_cost(
    self,
    agent_id: str,
    *,
    start: datetime | None = None,
    end: datetime | None = None,
) -> float:
    """Sum of ``cost`` 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 the configured currency for the agent.

    Raises:
        ValueError: If both *start* and *end* are given and
            ``start >= end``.
    """
    query_start = time.perf_counter()
    try:
        _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
    finally:
        record_budget_query(
            query_type="agent_cost",
            duration_sec=time.perf_counter() - query_start,
        )

async def get_project_cost(
    self,
    project_id: NotBlankStr,
    *,
    start: datetime | None = None,
    end: datetime | None = None,
) -> float:
    """Sum of ``cost`` for a single project.

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

    Returns:
        Rounded total cost in the configured currency for the project.

    Raises:
        ValueError: If both *start* and *end* are given and
            ``start >= end``.
    """
    query_start = time.perf_counter()
    try:
        _validate_time_range(start, end)
        logger.debug(
            BUDGET_PROJECT_COST_QUERIED,
            project_id=project_id,
            start=start,
            end=end,
        )
        snapshot = await self._snapshot()
        filtered = _filter_records(
            snapshot,
            project_id=project_id,
            start=start,
            end=end,
        )
        return _aggregate(filtered).cost
    finally:
        record_budget_query(
            query_type="project_cost",
            duration_sec=time.perf_counter() - query_start,
        )

async def get_project_records(
    self,
    project_id: NotBlankStr,
    *,
    start: datetime | None = None,
    end: datetime | None = None,
) -> tuple[CostRecord, ...]:
    """Return cost records for a specific project.

    Args:
        project_id: Project identifier to filter by.
        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_PROJECT_RECORDS_QUERIED,
        project_id=project_id,
        start=start,
        end=end,
    )
    snapshot = await self._snapshot()
    return _filter_records(
        snapshot,
        project_id=project_id,
        start=start,
        end=end,
    )

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

    Returns:
        Number of cost records.
    """
    async with self._get_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,
    )