Tracking and cutting token costs in Go
Go is usually the thing in front of the model rather than the thing calling it — which makes it the right place to enforce budgets, cancel abandoned work and expose the metrics.
The economics are language-independent — what tokens cost and where the money goes is the model. This page is the Go implementation, and Go usually occupies a specific position in these systems: the gateway, the proxy, the worker pool. The service that sits between many callers and a metered API.
That position is where cost control belongs, and Go is unusually well suited to it: contexts propagate cancellation, sync/atomic makes counters cheap, and errgroup bounds fan-out in one line.
Carry the budget in the context#
The idiomatic Go answer to "every call site needs this and I do not want it in every signature".
package budget
import (
"context"
"errors"
"sync/atomic"
)
var ErrExceeded = errors.New("budget exceeded")
// Budget is safe for concurrent use.
type Budget struct {
capMicroUSD int64
maxTurns int64
spentMicros atomic.Int64
turns atomic.Int64
}
func New(capMicroUSD int64, maxTurns int) *Budget {
return &Budget{capMicroUSD: capMicroUSD, maxTurns: int64(maxTurns)}
}
// Reserve is called before a request leaves. It is deliberately conservative:
// it charges the estimate up front and refunds the difference afterwards, so
// concurrent callers cannot all pass the check and then collectively overspend.
func (b *Budget) Reserve(estimateMicros int64) error {
if b.turns.Add(1) > b.maxTurns {
return errors.Join(ErrExceeded, errors.New("turn limit"))
}
if b.spentMicros.Add(estimateMicros) > b.capMicroUSD {
b.spentMicros.Add(-estimateMicros)
return ErrExceeded
}
return nil
}
func (b *Budget) Settle(estimateMicros, actualMicros int64) {
b.spentMicros.Add(actualMicros - estimateMicros)
}
func (b *Budget) Spent() int64 { return b.spentMicros.Load() }
type ctxKey struct{}
func With(ctx context.Context, b *Budget) context.Context {
return context.WithValue(ctx, ctxKey{}, b)
}
func From(ctx context.Context) (*Budget, bool) {
b, ok := ctx.Value(ctxKey{}).(*Budget)
return b, ok
}The reserve-then-settle pattern is the part worth copying. A naive "check, then call, then add" lets ten concurrent goroutines all pass the check against the same stale total and blow the cap tenfold. Reserving up front and refunding the difference is correct under concurrency, which in a Go gateway is the only case that matters.
Cancellation is a cost control#
If the caller goes away and you keep streaming from the provider, you keep paying. In Go this is nearly free to get right, and it is the most common thing generated gateway code omits.
func (h *Handler) Chat(w http.ResponseWriter, r *http.Request) {
ctx := r.Context() // cancelled when the client disconnects
b := budget.New(500_000, 12)
ctx = budget.With(ctx, b)
flusher, _ := w.(http.Flusher)
var usage llm.Usage
stream, err := h.llm.Stream(ctx, req) // ctx MUST reach the outbound request
if err != nil {
http.Error(w, err.Error(), http.StatusBadGateway)
return
}
defer func() { h.record(ctx, "chat", usage, ctx.Err() != nil) }()
for chunk, err := range stream {
if err != nil {
return
}
if _, werr := io.WriteString(w, chunk.Text); werr != nil {
return // client gone; ctx cancels upstream
}
usage = chunk.Usage
if flusher != nil {
flusher.Flush()
}
}
}The whole fix is that r.Context() reaches the outbound HTTP request. A handler that accepts a context and then builds its upstream call with context.Background() has silently disabled cancellation — contextcheck in golangci-lint catches most instances, and it is worth turning on for exactly this reason. See the Go failure modes.
Record on cancellation too. You still paid for input tokens and whatever output was generated before the abort, so dropping those records understates spend.
Bound the fan-out#
Unbounded go statements against a rate-limited API is the classic way a Go gateway turns a batch job into a 429 storm plus a retry bill.
func ClassifyAll(ctx context.Context, c llm.Client, inputs []string, limit int) ([]Category, error) {
g, ctx := errgroup.WithContext(ctx)
g.SetLimit(limit) // the whole fix
out := make([]Category, len(inputs))
for i, in := range inputs {
g.Go(func() error {
b, _ := budget.From(ctx)
est := int64(len(in)/4) * 3 // rough micros; conservative
if b != nil {
if err := b.Reserve(est); err != nil {
return fmt.Errorf("classify %d: %w", i, err)
}
}
res, usage, err := c.Complete(ctx, promptFor(in))
if b != nil {
b.Settle(est, cost(usage))
}
if err != nil {
return fmt.Errorf("classify %d: %w", i, err)
}
out[i] = res
return nil
})
}
return out, g.Wait()
}errgroup.WithContext also cancels every sibling on the first failure — so a budget exhaustion stops the remaining work rather than letting it run to completion and overspend.
Make the cache hit#
Same rule as everywhere: stable prefix first, volatile last. The Go-specific trap is map iteration order, which is deliberately randomised.
// BAD: tool definitions from a map iterate in a different order every run,
// so the prefix differs every call and the cache never hits.
for name, tool := range h.tools {
defs = append(defs, tool.Definition(name))
}
// GOOD: deterministic order
for _, name := range slices.Sorted(maps.Keys(h.tools)) {
defs = append(defs, h.tools[name].Definition(name))
}This one is genuinely easy to miss because nothing fails — the responses are correct, the tests pass, and the only symptom is a cache hit rate near zero. It is the Go flavour of a mistake every language has a version of.
func TestPromptPrefixIsStable(t *testing.T) {
a := BuildMessages("q1", []string{"doc-b", "doc-a"})
b := BuildMessages("q2", []string{"doc-a", "doc-b"})
if diff := cmp.Diff(a[:len(a)-1], b[:len(b)-1]); diff != "" {
t.Errorf("prefix is not stable, cache will miss (-a +b):\n%s", diff)
}
}Retries, bounded#
func withRetry(ctx context.Context, attempts int, fn func() error) error {
var err error
for n := 0; n < attempts; n++ {
if err = fn(); err == nil {
return nil
}
var apiErr *llm.APIError
if !errors.As(err, &apiErr) || !apiErr.Retryable() {
return err // never retry a 400
}
delay := time.Duration(1<<n)*time.Second + time.Duration(rand.N(500))*time.Millisecond
select {
case <-time.After(delay):
case <-ctx.Done():
return ctx.Err()
}
}
return err
}The select on ctx.Done() matters: a retry loop that sleeps without watching the context keeps a cancelled request alive, and you pay for the attempt that lands after the caller left.
Expose it as metrics#
Go services usually already have Prometheus. Cost is just another metric, and having it beside latency and error rate is what makes it a thing people look at.
var (
Tokens = promauto.NewCounterVec(prometheus.CounterOpts{
Name: "llm_tokens_total",
}, []string{"feature", "model", "kind"}) // kind: input|cached|output
CostMicros = promauto.NewCounterVec(prometheus.CounterOpts{
Name: "llm_cost_micro_usd_total",
}, []string{"feature", "model", "tenant"})
Calls = promauto.NewCounterVec(prometheus.CounterOpts{
Name: "llm_calls_total",
}, []string{"feature", "model", "outcome"}) // ok|error|aborted|budget
)
func Record(feature, model, tenant string, u llm.Usage, outcome string) {
Tokens.WithLabelValues(feature, model, "input").Add(float64(u.Input - u.Cached))
Tokens.WithLabelValues(feature, model, "cached").Add(float64(u.Cached))
Tokens.WithLabelValues(feature, model, "output").Add(float64(u.Output))
CostMicros.WithLabelValues(feature, model, tenant).Add(float64(cost(u)))
Calls.WithLabelValues(feature, model, outcome).Inc()
}The three queries worth putting on a dashboard:
# cache hit rate — should be high and flat. a drop means someone broke the prefix.
sum(rate(llm_tokens_total{kind="cached"}[1h]))
/ sum(rate(llm_tokens_total{kind=~"input|cached"}[1h]))
# cost per successful call, by feature
sum by (feature) (rate(llm_cost_micro_usd_total[1h]))
/ sum by (feature) (rate(llm_calls_total{outcome="ok"}[1h]))
# spend by tenant — find the customer who is 40% of the bill
topk(10, sum by (tenant) (increase(llm_cost_micro_usd_total[24h])))Keep tenant cardinality in mind. If you have a hundred thousand tenants, aggregate to a tier label in Prometheus and keep the per-tenant detail in your own log.
The four that usually do it in Go
- Make sure the request context reaches the outbound call. Cancellation is a cost control, and it is the most commonly dropped one.
- Reserve-then-settle budgets with atomics. Correct under concurrency, which is the only case a gateway has.
errgroup.SetLimiton every fan-out. Bounds the worst case.- Sort anything that comes from a map before it enters a prompt. Otherwise your cache never hits and nothing tells you.
Common questions#
Should the gateway enforce budgets, or the caller?#
The gateway, because it is the only place that sees all the traffic and cannot be bypassed. Callers should also have their own limits, but a budget enforced client-side is advice; one enforced at the boundary is a control.
Why reserve before the call rather than charge after?#
Because concurrent goroutines all pass a check against the same stale total and then collectively overspend. Reserving up front and refunding the difference on settle is the only version that is correct under concurrency, and concurrency is the normal case here.
Is context.Value the right place for a budget?#
For request-scoped values that cross API boundaries, yes — that is exactly what it is for. Keep the key unexported and provide typed With/From helpers, so nothing outside the package can put the wrong thing in.
Do I need a token counter in Go?#
Usually not. Go is typically the transport rather than the prompt builder, so a character-based estimate (len(s)/4) is enough for a reserve, and the usage the API returns is what you settle against. If Go is building your prompts, use the provider's counting endpoint for the boundary cases.
Get the Go agent pack
A battle-tested AGENTS.md, the review checklist, and the failure-mode cheat sheet for Go. One email, then occasional updates when the tooling shifts. No course pitch.
AGENTS.md now — no email needed.