middleware

Documentation

Overview

Package middleware provides composable func(Handler) Handler batteries for the model call — the "LiteLLM for Go" layer at the semantic level (handlers see messages, tool calls, and token usage, not raw bytes). Attach with the agent.WithMiddleware option (and tool middleware with agent.WithToolMiddleware).

a, err := agent.New(model, journal, agent.WithTools(tools...), agent.WithMiddleware(
	middleware.Retry(3, middleware.WithRetryIf(middleware.Retryable)),
	middleware.RateLimit(middleware.NewRateLimiter(time.Second, 5)),
))

Per-run limits that must hold across a resume live on the agent, where the journal is: agent.WithTokenBudget and agent.WithMaxTurns.

Index

• func Cost(m *CostMeter, r Rates) agent.Middleware
• func ErrorSummary(err error) string
• func Hedge(delay time.Duration, backups ...agent.Model) agent.Middleware
• func RateLimit(r *RateLimiter) agent.Middleware
• func Retry(n int, opts ...RetryOption) agent.Middleware
• func Retryable(err error) bool
• func ToolCache() agent.ToolMiddleware
• func ToolLog(logf func(format string, args ...any), opts ...ToolLogOption) agent.ToolMiddleware
• func ToolRateLimit(r *RateLimiter) agent.ToolMiddleware
• func ToolRetry(n int, opts ...RetryOption) agent.ToolMiddleware
• type CostMeter
  • func (m *CostMeter) Snapshot() CostSnapshot
• type CostSnapshot
• type RateLimiter
  • func NewRateLimiter(interval time.Duration, burst int) *RateLimiter
• type Rates
  • func (r Rates) Cost(u agent.Usage) float64
• type RetryOption
  • func WithBackoff(base, max time.Duration) RetryOption
  • func WithRetryIf(pred func(error) bool) RetryOption
  • func WithTimeout(d time.Duration) RetryOption
• type ToolLogOption
  • func LogErrorText() ToolLogOption

Functions

func Cost

func Cost(m *CostMeter, r Rates) agent.Middleware

Cost is a model middleware that records token usage and computes running cost, wherever it sits in the chain.

The answer view counts the usage of each call's answer, the response the agent records (see agent.ModelCall.OnAnswer), exactly once per call: a Cost inside a Hedge does not count a losing target's response, and a Cost inside a Retry does not count a response a middleware above it rejected. The spend view counts every request the call sends, through a hook it adds to the call (see agent.ModelCallHook): each attempt of a Retry and each target of a Hedge below it, the usage a failed request reported before failing, and the discarded spend a replayed turn reports (agent.ModelAttempt.Discarded). Outside an agent, call the model through agent.CallModel, whose model handler runs the hooks and reports the answer too; a Cost handler called directly, on a call that belongs to no turn, counts the response it returns as the answer. A call a middleware kept and passes on after its turn is over counts no answer: the turn's is already counted. The caller reads the totals with m.Snapshot.

func ErrorSummary

func ErrorSummary(err error) string

ErrorSummary describes err by what failed, never by what its text says: the control-flow signal it is (pending approval, resume halted, saga aborted), the provider's HTTP status and error type or code (from *agent.APIError) or a rate limit, and the condition or category sentinel it wraps. An error's text often carries content: a provider error body that echoes the prompt, a tool error that embeds the call's arguments, a URL with a credential in its query. The summary carries none of it, so it is safe to put in a log line, a span, or a metric label when content capture is off. An error that wraps none of these is summarized by its Go type, as "unclassified error (*url.Error)". A nil error summarizes to "".

api error: status 400 (invalid_request_error): model
rate limited (retry after 2s): model
unknown tool: tool

func Hedge

func Hedge(delay time.Duration, backups ...agent.Model) agent.Middleware

Hedge races the model call against one or more backup models and returns the FIRST successful response, cancelling the rest. It cuts tail latency (a slow primary no longer blocks the turn, you take whichever target answers first) and adds provider-outage failover (a down or throttled primary does not stall the run). It is safe here in a way it is not in a bare service: the agent loop journals only the winning response, at most once, so the durable record stays exactly-once and replays deterministically regardless of which target won the race.

Scheduling:

• The primary (the wrapped model, i.e. the one passed to agent.New) fires immediately.
• If delay > 0, the backups fire only after the primary has been outstanding for delay, so you pay for a backup generation only when the primary is actually slow. A delay of 0 fires every target at once (pure redundancy, always pays for all of them).
• If a target returns an ERROR while backups are still pending, the backups fire immediately (fast failover), rather than waiting out the remaining delay on a fast failure.

The first target to return a nil-error result wins; every other in-flight call is cancelled via a derived context. If every target fails, Hedge returns the joined error. With no backups it is a pass-through, so it is safe to wire unconditionally and add backups later.

Every target goes through the middleware installed inside Hedge (after it in Use) and the agent's checks on each model response, backups included: a backup is the same call with its Model set to the backup, and nothing else changed. Every request a target sends counts wherever the counting middleware sits: a RateLimit takes a token per request (1 + the backups launched), and Cost and the run's token budget count each target's usage as spend. Only the turn's answer, the winner's response, counts as an answer, once, even with Cost inside Hedge (see agent.ModelCall.OnAnswer). Each request is numbered apart (agent.ModelCall.Attempt).

Hedge does not wait for the losers: it returns as soon as a target wins, and a loser runs until its model honors the cancellation. Its usage reaches Cost when it ends. The run counts it with its next model turn, or, when the run ends first, waits for it (for at most two seconds, and not past the run's context) and journals it in a spend record of its own. A loser that has not reached the model when the turn ends is not sent. A model that ignores the cancellation for longer keeps a loser, and the middleware inside Hedge on its path, running after the run has ended, and its usage is not in the run's spend.

Streaming: Hedge has no streaming code. The agent's model handler lets one request of a turn at a time stream to a streaming caller (Agent.Stream), the first to start: the others run without streaming. If that request's response wins, the caller has seen it live. If another target wins, the caller gets agent.TurnRestarted, which retracts what the first one streamed, and then the winner's response. Either way nothing reaches the caller after the turn is over, so a loser can never show the caller text the run does not record.

Hedge only races the model generation. It does not duplicate tool calls or any other side effect: those run in the agent loop, above this middleware, under the at-most-once journal and the Safety layer. Note it hedges latency and availability, not correctness: taking the first answer is not the same as requiring agreement across models (that is a governance-quorum concern, not this).

// Primary is Anthropic; if it is quiet for 800ms, also try OpenAI and take the first.
a, err := agent.New(anthropicModel, journal, agent.WithTools(tools...),
	agent.WithMiddleware(middleware.Hedge(800*time.Millisecond, openaiModel)))

func RateLimit

func RateLimit(r *RateLimiter) agent.Middleware

RateLimit throttles model requests through the shared limiter, blocking (respecting context cancellation) until a token is available. Proactive throttling complements Retry's reactive Retry-After backoff: it keeps you under the provider's ceiling instead of bouncing off it.

It takes a token for every request the call sends, wherever it sits in the chain: each attempt of a Retry below it and each target a Hedge below it launches waits for its own token, through a hook it adds to the call (see agent.ModelCallHook). Outside an agent, call the model through agent.CallModel, whose model handler runs the hooks too.

func Retry

func Retry(n int, opts ...RetryOption) agent.Middleware

Retry retries the model call up to n additional times on error, respecting context cancellation. Between attempts it sleeps with exponential backoff (full jitter, base×2^attempt, capped at max). When the error is *agent.RateLimited it sleeps for min(RetryAfter, max) instead.

Default backoff: base=200ms, max=10s. Override with WithBackoff.

n must be at least 0 (0 calls the model once). With n < 0 every call fails with an error wrapping agent.ErrConfig without calling the model (see checkRetryCount).

Each attempt passes the call on unchanged, so every attempt's requests run the hooks middleware outside Retry added (a RateLimit token, Cost's spend) and are numbered by the turn's shared counter (agent.ModelCall.Attempt).

Streaming: a streaming caller (Agent.Stream) sees each attempt's deltas live. When an attempt that streamed deltas fails and Retry calls the model again, the stream emits agent.TurnRestarted before the next attempt's deltas, so the caller can discard the failed attempt's partial text. Retry has no code for this: the agent's model handler does it.

func Retryable

func Retryable(err error) bool

Retryable is a ready-made classifier for WithRetryIf: it retries transient failures and fails fast on terminal ones. It retries *agent.RateLimited, a per-attempt timeout (context.DeadlineExceeded), and *agent.APIError with a 5xx or 408 status. It does not retry parent cancellation (context.Canceled), 4xx API errors (auth, validation), or failures the same request would repeat: agent.ErrConfig (a request the adapter refused to build), agent.ErrQuotaExhausted (used-up quota or credit, which no wait lifts), agent.ErrResponseTooLarge, and agent.ErrTruncatedToolArgs (a tool call cut off by the response's output token limit, which the same request hits again). It does retry agent.ErrStreamProtocol (a stream that broke its provider's event protocol, a fault of that one response) and agent.ErrToolUseIDReused (a fresh generation can issue new ids). Errors it cannot classify (e.g. raw network errors, a provider's mid-stream server error) are retried, since those are usually transient.

agent.WithMiddleware(middleware.Retry(3, middleware.WithRetryIf(middleware.Retryable)))

func ToolCache

func ToolCache() agent.ToolMiddleware

ToolCache memoizes successful tool results by (name, args) in a process-local map and short-circuits a repeat call — the tool never runs on a hit. It demonstrates the short-circuit power of tool middleware.

It caches only tools marked ReadOnly (the call's Spec.Safety); every other tool call runs, since two calls with the same arguments to a side effect are two effects, such as two separate charges. Even for a ReadOnly tool a hit returns the earlier result, so use it only where the output depends solely on the arguments. Errors are never cached. This is an in-memory, unbounded, per-instance cache — it is NOT the durable journal (which already dedupes each tool-use ID at-most-once); it dedupes DISTINCT calls with identical arguments within one process.

func ToolLog

func ToolLog(logf func(format string, args ...any), opts ...ToolLogOption) agent.ToolMiddleware

ToolLog logs one line per tool call — name, duration, and ok/error — via logf (pass log.Printf, t.Logf, or a structured logger's Printf-shaped method). Attach with the agent.WithToolMiddleware option.

a, err := agent.New(model, journal, agent.WithTools(tools...), agent.WithToolMiddleware(middleware.ToolLog(log.Printf)))

A failed call is logged by its ErrorSummary (category, condition, provider status), not its text: a tool's error text commonly embeds the call's arguments or a URL with a credential in it, and a log line is not the place for them. Pass LogErrorText to log the text the agent journals for the call instead (see agent.ToolCall.ErrorText).

func ToolRateLimit

func ToolRateLimit(r *RateLimiter) agent.ToolMiddleware

ToolRateLimit throttles tool calls through the shared limiter. A call that gives up waiting for a slot (its context was done: the run was cancelled, or the tool's WithTimeout deadline passed) never calls the tool, and fails with an error wrapping agent.ErrToolNotCalled, so the agent records a known failure rather than halt for an outcome that cannot have happened.

func ToolRetry

func ToolRetry(n int, opts ...RetryOption) agent.ToolMiddleware

ToolRetry retries a failing tool call up to n additional times with exponential backoff + jitter (honoring an *agent.RateLimited RetryAfter), respecting context cancellation. Reuses RetryOption/WithBackoff from the model-side Retry.

It retries only a tool that is retry-safe (the call's Spec.Safety is RetrySafe: ReadOnly or Idempotent). A tool that is not runs once, and its error is returned as is: a failed call to a side effect may still have taken effect, as when a payment gateway times out after charging, and running it again could repeat it.

Default backoff: base=200ms, max=10s. Override with WithBackoff.

n must be at least 0 (0 runs the tool once). With n < 0 every call fails with an error wrapping agent.ErrConfig without running the tool, whatever its safety.

a, err := agent.New(model, journal, agent.WithTools(tools...), agent.WithToolMiddleware(middleware.ToolRetry(3)))

Types

type CostMeter

type CostMeter struct {
	// contains filtered or unexported fields
}

CostMeter accumulates token usage and computed cost across model calls, in two views: the answers, the responses the calls returned, which a run records; and the spend, every request the calls sent, including failed attempts a Retry repeated and losing Hedge targets, which the provider bills too. Read both with Snapshot. It is safe for concurrent use.

func (*CostMeter) Snapshot

func (m *CostMeter) Snapshot() CostSnapshot

Snapshot returns the meter's totals, read together under one lock, so the answer and spend views are always of the same moment.

type CostSnapshot

type CostSnapshot struct {
	Answer, Spend       agent.Usage
	AnswerUSD, SpendUSD float64
}

CostSnapshot is a CostMeter's totals at one moment. Answer and AnswerUSD are the usage and USD cost of the responses the calls returned; Spend and SpendUSD those of every request the calls sent.

type RateLimiter

type RateLimiter struct {
	// contains filtered or unexported fields
}

RateLimiter is a token-bucket limiter shared across calls: it allows one token every `interval` with a burst up to `burst`. It is dependency-free (no background goroutine): tokens accrue lazily from elapsed time. Safe for concurrent use. Share one limiter across model and tool middleware to cap a whole agent, or use separate limiters per surface. Waiters are not served in arrival order: a call arriving as a token frees can take it ahead of one already waiting.

func NewRateLimiter

func NewRateLimiter(interval time.Duration, burst int) *RateLimiter

NewRateLimiter allows one call every `interval` (so rate = 1/interval), bursting up to `burst` calls. A burst < 1 is treated as 1. Starts full so the first `burst` calls proceed immediately. An interval <= 0 sets no limit: every call proceeds at once.

type Rates

type Rates struct {
	InputPer1M      float64
	OutputPer1M     float64
	CacheReadPer1M  float64
	CacheWritePer1M float64
}

Rates holds the billing rates for a model in USD per 1,000,000 tokens.

func (Rates) Cost

func (r Rates) Cost(u agent.Usage) float64

Cost returns the USD cost of the token usage u at these rates.

type RetryOption

type RetryOption func(*retryConfig)

RetryOption configures optional behavior of Retry.

func WithBackoff

func WithBackoff(base, max time.Duration) RetryOption

WithBackoff sets the initial backoff duration and the maximum cap. Between retries the sleep doubles each attempt (full jitter applied) and is capped at max. If the error is *agent.RateLimited the Retry-After hint overrides the computed backoff (capped at max).

func WithRetryIf

func WithRetryIf(pred func(error) bool) RetryOption

WithRetryIf classifies which errors are worth retrying: pred returns true to retry, false to fail fast. Use it to stop burning attempts (and tokens) on terminal errors such as HTTP 4xx auth/validation failures, while still retrying transient ones. The default (nil) retries every error. Context cancellation always stops the loop regardless. See Retryable for a ready-made classifier.

func WithTimeout

func WithTimeout(d time.Duration) RetryOption

WithTimeout bounds each individual attempt with its own deadline: the call gets a context that is cancelled after d, so a hung model or tool call fails that attempt (and is retried) instead of blocking the run forever. The parent context still governs overall cancellation. 0 (default) means no per-attempt timeout.

type ToolLogOption

type ToolLogOption func(*toolLogConfig)

ToolLogOption configures ToolLog.

func LogErrorText

func LogErrorText() ToolLogOption

LogErrorText makes ToolLog log a failed call's error text instead of its ErrorSummary: the text the agent journals for the call and sends to the model (agent.ToolCall.ErrorText), which is the agent's WithToolErrorRedactor text if one is set, with every URL in it redacted. It can still carry content (the call's arguments, a provider's echo of them), so enable it only where the log may hold that content.