A pluggable pre-call gate for token throughput.
Oban caps the number of concurrent LLM jobs, but never the tokens they send per minute — so a burst of steps (e.g. compare mode fanning one flow out across many models) can trip a provider's input-tokens-per-minute (ITPM) or output-tokens-per-minute (OTPM) ceiling even while the queue is within its concurrency limit. This behaviour is the seam where a host inserts a token accountant.
Baton.LLMStep calls acquire/3 immediately before every model call and
reconcile/3 immediately after every call it gated — with the real usage on
success, and with zero actuals when the call failed outright (a 429, a
timeout, a transport error), so the reservation is credited back in full
rather than leaking out of a tight bucket:
acquire/3is handed the model, an estimate of the input tokens the call will spend, and the client opts. It either reserves that budget and returns{:ok, opts}— withoptspossibly rewritten (an implementation may set:rate_account, and in a fallback design may rewrite:provider/:model) — or returns{:snooze, seconds}to park the job attempt-free (no retry consumed, the Oban slot is freed) until budget is available. The snooze reusesBaton.LLMStep's existing 429 snooze contract. Output-token budget is not part ofacquire/3— the call's ownopts[:max_tokens]is already the caller's requested output ceiling, so an implementation that tracks OTPM can reserve against it directly without Baton computing anything new.reconcile/3is handed the account that was reserved, an estimate map, and an actual map — each%{input: non_neg_integer(), output: non_neg_integer()}— so the implementation can correct both the input and output estimates' drift against the real spend. It is bookkeeping —Baton.LLMStepswallows its errors so it can never fail a paid job.
The default implementation (Baton.RateLimiter.Noop) reserves nothing and is
wired unless a host configures its own:
config :baton, rate_limiter: MyApp.LLM.RateLimiterEstimates vs. actuals
acquire/3's est_tokens is a pre-call input estimate (Baton derives it
from the prompt character count — see Baton.LLMStep), because the true
input-token count is only known once the provider replies. reconcile/3
closes that loop for both dimensions: estimate.input is est_tokens from
the paired acquire/3 call, and estimate.output is the call's
opts[:max_tokens] (0 when unset) — the reservation an OTPM-tracking
implementation would have made against its own output bucket, using data it
already had at acquire/3 time. actual.input and actual.output come from
the response's reported usage. An implementation credits (or, on an
under-estimate, further debits) each dimension's own bucket independently —
the two are unrelated budgets, not a combined total.
Summary
Types
Opaque account identifier an implementation associates with a reservation.
acquire/3 may stash it in the returned opts under :rate_account;
Baton.LLMStep threads that value back into reconcile/3.
Token counts split by direction — the shape of both reconcile/3 arguments
after the account.
Callbacks
Reserve budget for an upcoming call, or ask the caller to snooze.
Gate a batch submission, or ask the caller to snooze.
Reconcile a gated call's estimated input/output spend against its actual
input/output tokens. Always returns :ok; treated as best-effort by
Baton.LLMStep.
Types
@type account() :: term()
Opaque account identifier an implementation associates with a reservation.
acquire/3 may stash it in the returned opts under :rate_account;
Baton.LLMStep threads that value back into reconcile/3.
@type token_counts() :: %{input: non_neg_integer(), output: non_neg_integer()}
Token counts split by direction — the shape of both reconcile/3 arguments
after the account.
Callbacks
@callback acquire( model :: String.t() | nil, est_tokens :: non_neg_integer(), opts :: keyword() ) :: {:ok, keyword()} | {:snooze, pos_integer()}
Reserve budget for an upcoming call, or ask the caller to snooze.
{:ok, opts} proceeds with the (possibly rewritten) opts. {:snooze, secs}
parks the job attempt-free for secs seconds and re-runs the step later.
@callback acquire_batch(opts :: keyword()) :: {:ok, keyword()} | {:snooze, pos_integer()}
Gate a batch submission, or ask the caller to snooze.
Optional — when a limiter doesn't export it, Baton.LLMStep proceeds as if
it returned {:ok, opts}.
Batch traffic is metered differently from live traffic: providers draw it
from a separate pool, so there is nothing to reserve against the ITPM/OTPM
budgets acquire/3 protects and no reconcile/3 counterpart here. What can
still be exceeded is the rate of submissions (and any cap on batches in
flight), which is what this callback is for. Polling is ungated — it is cheap
and asking a limiter about it would only add latency to a step that is
already waiting hours.
@callback reconcile(account(), estimate :: token_counts(), actual :: token_counts()) :: :ok
Reconcile a gated call's estimated input/output spend against its actual
input/output tokens. Always returns :ok; treated as best-effort by
Baton.LLMStep.
A call that failed without usable usage (rate-limited, timed out, transport
error) reconciles with %{input: 0, output: 0} — an implementation doing
estimate - actual bucket arithmetic credits the full reservation back
with no special-casing.