TL;DR: Before an agent starts an expensive AI step, price the bounded request, compare that estimate with available budget, not the stored balance, and atomically reserve the estimate. During a production API-key rotation, attach both the logical account and the credential version to that reservation. Settle actual usage against the same record after the call. This keeps the service live while preventing overlapping agent turns, retries, and old/new keys from spending or attributing the same budget twice.
The important trade-off is conservative admission versus useful throughput. A padded estimate rejects some work that might have fit; a hopeful estimate admits work whose final charge can cross the cap. For a marketplace account platform, I would favor a documented ceiling for each step and make exceptions explicit. OTP delivery taught this class of system a useful lesson: a request being accepted is not the same as the operation being finished, and retries are part of the protocol rather than an afterthought.
How should an agent estimate cost before an expensive step?
Suppose a seller account has a remaining AI budget of 12.00 units. Two workers each estimate a catalog-enrichment turn at 7.00. Both read 12.00, both pass, and together commit 14.00. The arithmetic was correct; the concurrency model was not.
The admission value therefore needs three components:
available = limit - settled_usage - open_reservations
The comparison and reservation must share one atomic transaction. A later network call cannot be part of that transaction, so the durable reservation becomes the bridge between local budget state and remote execution. Give it a stable operation ID. A retry with that ID must recover the existing decision instead of creating another hold.
Short version: reserve first.
There is another edge case hiding in estimation. An agent loop often knows the input size and the configured output ceiling, but not the eventual output size. Admission should use an upper bound derived from those known inputs and the applicable rate snapshot. It should not use the average bill from previous turns as if that were a ceiling. Store the estimate inputs and the rate snapshot identifier with the reservation; otherwise a later audit can reproduce neither the decision nor its arithmetic.
Use fixed-point decimal arithmetic for money-like units. Binary floating point is the wrong representation for a hard comparison at a boundary. The budget's unit also needs to be explicit: currency, internal credits, or another metered unit are different contracts even if each happens to display two decimal places.
Rotation changes identity, not the spending authority
A zero-downtime key rotation creates an overlap window. Requests signed by the retiring key may still be in flight while new requests use the replacement. If budget is partitioned by raw API key, each credential can appear to own a fresh allowance. Billing attribution then splits one marketplace account into two accidental spenders.
Model those identities separately. The logical billing account owns the limit. A credential version authenticates a request and becomes an attribution dimension. The reservation should record account_id, credential_version, and operation_id, but only account_id selects the budget bucket. This lets operators answer two different questions later: who was allowed to spend, and which credential authorized this particular attempt?
Do not overwrite the version on retry. An operation admitted under key-v17 remains attributed to key-v17, even if the worker resumes after traffic has moved to key-v18. Changing that field would make the audit trail describe the recovery worker rather than the admitted operation. New operations use the new version; old reservations drain naturally or expire under a defined policy.
The secret itself does not belong in the ledger, logs, traces, or idempotency key. Store an opaque version label. The OWASP Secrets Management Cheat Sheet treats rotation, expiration, revocation, and auditing as parts of a secret lifecycle; separating the credential value from its operational metadata follows that boundary and limits unnecessary exposure.
A minimal reservation gate in Python
The following example keeps storage behind a small interface. try_reserve must be implemented as one compare-and-write operation by the database. The rate calculation is deliberately injected: admission logic should not quietly fetch mutable pricing in the middle of a transaction.
from dataclasses import dataclass
from decimal import Decimal, ROUND_UP
from typing import Protocol
@dataclass(frozen=True)
class PlannedStep:
account_id: str
operation_id: str
credential_version: str
input_units: int
max_output_units: int
@dataclass(frozen=True)
class RateSnapshot:
snapshot_id: str
input_rate: Decimal
output_rate: Decimal
unit_scale: Decimal
@dataclass(frozen=True)
class Admission:
accepted: bool
reserved: Decimal
reason: str
class BudgetLedger(Protocol):
def existing_reservation(self, operation_id: str) -> Admission | None: ...
def try_reserve(
self,
*,
account_id: str,
operation_id: str,
credential_version: str,
estimate: Decimal,
rate_snapshot_id: str,
) -> bool: ...
def estimate_ceiling(step: PlannedStep, rates: RateSnapshot) -> Decimal:
raw = (
Decimal(step.input_units) * rates.input_rate
+ Decimal(step.max_output_units) * rates.output_rate
) / rates.unit_scale
return raw.quantize(Decimal("0.000001"), rounding=ROUND_UP)
def admit(
step: PlannedStep,
rates: RateSnapshot,
ledger: BudgetLedger,
) -> Admission:
previous = ledger.existing_reservation(step.operation_id)
if previous is not None:
return previous
estimate = estimate_ceiling(step, rates)
accepted = ledger.try_reserve(
account_id=step.account_id,
operation_id=step.operation_id,
credential_version=step.credential_version,
estimate=estimate,
rate_snapshot_id=rates.snapshot_id,
)
if not accepted:
return Admission(False, Decimal("0"), "insufficient available budget")
return Admission(True, estimate, "reserved")
The code does not read a balance and then write a reservation as separate calls. That omission is intentional. A repository API that exposes get_remaining() plus create_hold() invites the race described earlier; the safer interface makes the atomic invariant visible.
After the AI call, settlement records actual metered usage and releases the unused portion. If execution fails before usage is known, keep the reservation until reconciliation can establish the result or a documented expiry policy releases it. Blindly releasing on a client timeout can admit replacement work while the original request is still running.
Make attribution testable
Budget correctness deserves adversarial tests, not one happy-path unit test. Start two admissions against a limit that can fund only one and use a barrier so both contend at the storage boundary. Exactly one reservation should succeed. Repeat the same operation ID and confirm that the reserved total does not move. Then rotate from key-v17 to key-v18; admit one distinct operation under each and verify that both debit the same account bucket while retaining separate credential labels.
Reconciliation needs equal attention. Test actual usage below the estimate, exactly at it, and above it. The last case should produce an explicit overage state for investigation or policy handling, not a negative balance disguised by clamping it to zero. Also test a worker crash after reservation but before dispatch, a timeout after dispatch, and duplicate completion events.
Observability should follow the same identity model. Useful fields include the operation ID, account ID, credential version, rate snapshot ID, estimated amount, settled amount, and state transition. Never emit the key. Metrics can track reservation age and estimate error distributions without turning credentials into labels; raw credential values are both sensitive and disastrously high-cardinality.
The decision rule is now inspectable: an operation runs only if its worst bounded estimate can be reserved against the account's currently available allowance. Humans may approve a different policy for low-risk work, but the agent should not invent one mid-loop.
Roll out the gate without interrupting rotation
Deploy the ledger schema and observe estimates before enforcing rejection. During that phase, calculate what the decision would have been, record the rate snapshot, and compare estimates with settled usage. This validates units and attribution while existing traffic continues.
Next, enforce reservations for a narrow class of expensive steps, then expand by operation type. Keep both credential versions valid for the planned overlap, route new work to the replacement, and preserve the original credential version on every open reservation. Revoke the retiring secret only after its in-flight operations and reservations are accounted for under the rotation policy.
The compact operational sequence is: create the replacement credential, label it with a non-secret version, shift new admissions, drain old work, reconcile holds, and revoke the retired credential. The account owns the budget; the credential version explains the charge. Keeping those roles distinct is what makes cost control and live rotation compatible.
Top comments (0)