context.Context in a Go agent: cancellation, deadlines, and causes

An agent run makes many slow calls: a model call that takes half a minute, a tool that runs a test suite, another model call, and so on for dozens of turns. At any point, the person who started it may want it stopped, a call may hang, or the process may be shutting down. In Go, all three travel the same way, through a context.Context passed into every function that waits on something.

This page goes through the ways an agent uses one, each with code from the companion repository, where it is tested. Several of the tools are recent additions to Go, so here is what each needs:

The companion code needs Go 1.27; with an older toolchain, go downloads the right one when it reads the module’s go.mod.

Without a context

Take the context away and two things go wrong.

The first is Ctrl-C. By default, a Go program that gets an interrupt exits on the spot, and its deferred functions never run. For an agent, that means a model call dropped partway, which the provider may bill anyway; a tool stopped in the middle of an action, such as writing a file or filing an issue; no checkpoint saved, so the run can’t resume; and no report of what it spent or why it ended. Stopped with a plain kill, the agent’s process gets SIGTERM alone, so a go test or a tool server it started can keep running on its own.

The second is that every kind of waiting needs its own way to stop. You can catch the signal yourself, but then something has to reach everything that is blocked: an HTTP request, a child process to kill, a timer to stop, each of your own goroutines waiting on a channel. Each has a different mechanism, and before Go 1.7 brought context into the standard library, each library had its own; net/http still carries one of them, the Cancel field on a request, now deprecated in favor of the request’s context. Threading a different stop signal through every layer, with no way to say why it fired, is the job context does once. Every blocking call agrees on it, from net/http and os/exec to the model providers’ SDKs and your own code, and it carries the reason with it.

The rest of this page is that alternative, piece by piece.

How a context works

context.Context is an interface with four methods:

• Done() returns a channel that is closed when the context ends, for use in a select.
• Err() is nil while the context is live, then context.Canceled or context.DeadlineExceeded.
• Deadline() reports when the context will end on its own, if it has a deadline.
• Value(key) returns a value stored in the context, or nil.

A context is never changed. Functions such as context.WithCancel and context.WithTimeout take a parent and return a new child, with a function that cancels it, so the contexts in a program form a tree with context.Background() at the root. Cancellation travels down the tree and never up:

Two more rules complete the picture. The first cancellation wins: once a context is done, later calls to its cancel function change nothing, including its cause. And context.Cause always has an answer: the cause given when the context was cancelled, the parent’s cause when cancellation came from above, and Err() itself when nobody gave one.

context.Background() is the root for main, for init, and for anything that starts a whole operation. context.TODO() is the same empty context under another name, a marker for code that should be passed a context and isn’t yet; linters and readers can find it.

One context, from main down

The context starts in main, and the first thing it is good for is Ctrl-C:

func main() {
	// Ctrl-C cancels ctx, and the loop stops before its next turn.
	ctx, stop := signal.NotifyContext(context.Background(),
		os.Interrupt)
	defer stop()
	if err := run(ctx, os.Stdout); err != nil {
		log.Fatal(err)
	}
}

signal.NotifyContext returns a context that is cancelled when the process gets an interrupt. From here the context is passed down, as the first parameter of every function that does I/O or can block: the loop, the model client, each tool. Nothing below main makes a fresh context.Background(), because a fresh context cuts the chain: whatever ran under it would ignore Ctrl-C.

Passing it as a parameter, never storing it in a struct, keeps each call’s context its own. A client that stored a context would cancel every later request with the first one’s.

A server starts its chains elsewhere. In an HTTP handler, r.Context() is the request’s context, and it ends when the client disconnects or the handler returns. An agent run started from a handler derives its context from it, so a person who closes the browser tab stops the run; the browser streaming example is built that way.

A deadline for each model call

A model call that hangs should fail on its own, without waiting for a person to notice. The client gives every request its own deadline, derived from the run’s context:

// send posts a request to the API and returns the decoded response.
func (c *Client) send(ctx context.Context,
	wire chatRequest) (chatResponse, error) {
	ctx, cancel := context.WithTimeoutCause(ctx, c.timeout(),
		ErrCallTimeout)
	defer cancel()
	resp, err := c.post(ctx, wire)
	if err != nil {
		return chatResponse{}, err
	}
	defer resp.Body.Close()
	return readResponse(ctx, resp)
}

The derived context ends at whichever comes first: the caller’s cancellation, or the call’s own timeout. defer cancel() releases the timer when the call returns early, which is most of the time; without it, the timer lives until it fires.

The run has limits of its own above this, a turn limit and a cost budget, so a deadline for the whole run is rarely needed. When it is, wrap the run’s context the same way, and each call’s shorter deadline still applies inside it. context.WithDeadline does the same with a time instead of a duration.

Code can also read the time left: deadline, ok := ctx.Deadline() reports the deadline, if there is one, and time.Until(deadline) how long remains. An agent can use it to skip starting work that can’t finish in time, such as a model call that usually takes 40 seconds when 10 remain, and stop with a clear reason instead of a timeout partway through.

Say why it stopped

A cancelled context reports one of two errors, context.Canceled or context.DeadlineExceeded, and neither says which deadline passed. “Deadline exceeded” from a run with a call timeout, a tool timeout, and a test timeout is a puzzle.

context.WithTimeoutCause and context.WithCancelCause attach a cause, an error of your own, and context.Cause reads it back. Above, the call’s timeout carries ErrCallTimeout. One detail makes this harder than it looks: when a request is cancelled, Go’s HTTP/1 transport returns the cause, and its HTTP/2 transport returns only the kind of cancellation. The client joins both into the error, so callers can check either one:

// callError makes the error from a cancelled call carry both the kind
// of cancellation and its cause. The HTTP/1 transport reports only the
// cause, and the HTTP/2 transport only the kind.
func callError(ctx context.Context, err error) error {
	if ctx.Err() == nil {
		return err
	}
	if !errors.Is(err, ctx.Err()) {
		err = fmt.Errorf("%w: %w", ctx.Err(), err)
	}
	if cause := context.Cause(ctx); !errors.Is(err, cause) {
		err = fmt.Errorf("%w: %w", err, cause)
	}
	return err
}

With both in the chain, the program can say in plain words why a run ended, checking the specific causes before the general kinds:

// whyStopped says in a few words why a run ended, from its error.
func whyStopped(err error) string {
	switch {
	case err == nil:
		return "the model answered"
	case errors.Is(err, agent.ErrMaxTurns):
		return "the turn limit"
	case errors.Is(err, agent.ErrBudgetExceeded):
		return "the cost budget"
	case errors.Is(err, openrouter.ErrStreamIdle): // before Canceled
		return "the stream went quiet, worth retrying"
	case errors.Is(err, context.Canceled):
		return "cancelled"
	case errors.Is(err, openrouter.ErrCallTimeout):
		return "the model call timed out, worth retrying"
	case errors.Is(err, context.DeadlineExceeded):
		return "a deadline passed"
	case errors.Is(err, agent.ErrBadReply):
		return "a reply the loop can't use"
	}
	if e, ok := errors.AsType[*openrouter.APIError](err); ok {
		if e.Retryable() {
			return fmt.Sprintf("API error %d, worth retrying", e.Status)
		}
		return fmt.Sprintf("API error %d", e.Status)
	}
	return "an error"
}

“The model call timed out, worth retrying” and “cancelled” call for different responses, and only the cause tells them apart.

Between turns, and the calls that didn’t run

Every model call notices cancellation on its own, because the context reaches the HTTP request. The loop also checks between turns, so it doesn’t start a call it would cancel at once:

// beforeTurn checks what can stop a run before it spends another call:
// cancellation, and the cost budget.
func beforeTurn(ctx context.Context, res Result, lim Limits) error {
	if err := stopped(ctx); err != nil {
		return err
	}
	if lim.MaxCost > 0 && res.Usage.Cost >= lim.MaxCost {
		return fmt.Errorf("%w: spent $%.6f of $%.6f",
			ErrBudgetExceeded, res.Usage.Cost, lim.MaxCost)
	}
	return nil
}

// stopped returns the error for a run whose ctx is done, carrying the
// kind of cancellation and, when it says more, the cause. It returns
// nil while ctx isn't done.
func stopped(ctx context.Context) error {
	if ctx.Err() == nil {
		return nil
	}
	err := fmt.Errorf("stopped before turn: %w", ctx.Err())
	if cause := context.Cause(ctx); !errors.Is(err, cause) {
		err = fmt.Errorf("%w: %w", err, cause)
	}
	return err
}

When cancellation arrives partway through a reply that asked for several tools, some calls have run and some haven’t. Every tool call in a conversation needs a matching result, or the next request to the model is rejected, so the calls that didn’t run get a result that says why:

// runTools answers every call in a reply, in order, passing each tool
// message to yield as it is made. Once ctx is done, the calls that
// haven't run get a "cancelled" result instead, so every tool call in
// the conversation still has its answer. It reports false if yield did.
func runTools(ctx context.Context, byName map[string]Tool,
	calls []ToolCall, yield func(Message) bool) ([]Message, bool) {
	var out []Message
	for _, call := range calls {
		msg := Message{Role: "tool", ToolCallID: call.ID, IsError: true}
		if ctx.Err() != nil {
			msg.Content = "cancelled before it ran: " +
				context.Cause(ctx).Error()
		} else {
			msg = runTool(ctx, byName, call)
		}
		out = append(out, msg)
		if !yield(msg) {
			return out, false
		}
	}
	return out, true
}

Waiting in your own code

The standard library’s blocking calls take a context and stop on their own. Code that waits on its own channels and timers has to watch the context itself, with a select on ctx.Done() beside whatever it is waiting for. The agent limits how many tool calls run at once with a buffered channel of slots, and waiting for a slot gives up when the context ends:

// Slots limits how many things run at once, such as the tool calls of
// one reply, or test runs across every agent in a program. Make one with
// NewSlots; a nil Slots would wait forever.
type Slots chan struct{}

// NewSlots returns Slots that let n things run at once.
func NewSlots(n int) Slots { return make(Slots, max(n, 1)) }

// Acquire waits for a free slot, and gives up when ctx is done,
// returning why. Every successful Acquire needs a Release.
func (s Slots) Acquire(ctx context.Context) error {
	select {
	case s <- struct{}{}:
		return nil
	case <-ctx.Done():
		return fmt.Errorf("wait for a slot: %w", context.Cause(ctx))
	}
}

// Release frees a slot that Acquire took.
func (s Slots) Release() { <-s }

A rate limiter that waits on a timer does the same, and stops the timer when it leaves early:

// Wait returns when the caller may start a call, or when ctx is done,
// with why. A turn given up this way is still used, which errs on the
// side of sending fewer requests.
func (l *Limiter) Wait(ctx context.Context) error {
	l.mu.Lock()
	now := time.Now()
	at := l.next
	if at.Before(now) {
		at = now
	}
	l.next = at.Add(l.Every)
	l.mu.Unlock()
	if !at.After(now) {
		return nil
	}
	t := time.NewTimer(at.Sub(now))
	defer t.Stop()
	select {
	case <-t.C:
		return nil
	case <-ctx.Done():
		return fmt.Errorf("wait for the rate limit: %w",
			context.Cause(ctx))
	}
}

Both return the cause, so a caller sees why the wait ended. Without the ctx.Done() case, a goroutine blocked on a full channel waits forever after its run is gone: a goroutine leak, which in a long-running service adds up one stopped run at a time. Every goroutine needs an owner and a way to stop, and for most of them the way to stop is a context.

Tools that start processes

A tool that runs go test starts a process, and a cancelled context has to reach it too. exec.CommandContext kills the process when the context is done:

	ctx, cancel := context.WithTimeoutCause(ctx, timeout,
		ErrTestsTimedOut)
	defer cancel()
	inner := fmt.Sprintf("-timeout=%s", timeout*3/4)
	cmd := exec.CommandContext(ctx, "go", "test", "-json", "-count=1",
		inner, "./...")
	cmd.Dir = dir
	cmd.Env = testEnv()
	cmd.WaitDelay = 5 * time.Second
	var stdout, stderr bytes.Buffer
	cmd.Stdout, cmd.Stderr = &stdout, &stderr
	err = cmd.Run()
	if ctx.Err() != nil {
		err := fmt.Errorf("go test: %w", ctx.Err())
		if cause := context.Cause(ctx); !errors.Is(err, cause) {
			err = fmt.Errorf("%w: %w", err, cause)
		}
		return "", err
	}
	_, exited := errors.AsType[*exec.ExitError](err)
	if err != nil && !exited {
		return "", fmt.Errorf("run_tests: %w", err)
	}
	if err != nil && stdout.Len() == 0 {
		return "", fmt.Errorf("run_tests: %s",
			strings.TrimSpace(stderr.String()))
	}
	return report(&stdout, stderr.String(), dir), nil
}

Three details here are easy to miss:

• go test gets its own -timeout, three quarters of the tool’s, so the tests usually stop themselves first and report which test hung.
• cmd.WaitDelay bounds how long Run waits for the output pipes after the process is killed; a test that started a child process can otherwise keep them open.
• The tool checks ctx.Err() before looking at the command’s error, because a killed process reports only “signal: killed”, and the cause says why it was killed.

A tool written in Go, rather than a process, must check its context itself, by passing it to every call that can block. A tool that ignores its context can’t be stopped: a timeout wrapper around it still has to wait until the tool returns, however late.

Some calls can’t take a context at all, such as a read from a connection or a pipe. context.AfterFunc runs a function when the context is done, and closing the connection there unblocks the read. The MCP client starts a server process and talks to it over pipes, and arranges exactly that while it connects:

stop := context.AfterFunc(ctx, func() { _ = c.Close() })
if err := c.connect(ctx); err != nil {
	stop()
	err = fmt.Errorf("connect to %s: %w", cmd.Path, err)
	return nil, errors.Join(err, c.Close())
}
return c, nil

The stop function it returns cancels the arrangement, for when the work finishes first.

A timer as a cause: the idle stream

A streamed reply can stop arriving without ending: the connection stays open and nothing comes. A deadline for the whole call doesn’t fit, because a long reply legitimately takes minutes. What fits is a timer that is reset every time a piece arrives, and cancels with its own cause when it fires:

// Stream sends the conversation and yields the reply in pieces as the
// model writes it. A stream that goes quiet for IdleTimeout is
// cancelled with ErrStreamIdle as the cause.
func (c *Client) Stream(ctx context.Context,
	req agent.Request) iter.Seq2[agent.Delta, error] {
	return func(yield func(agent.Delta, error) bool) {
		ctx, cancel := context.WithCancelCause(ctx)
		defer cancel(nil)
		idle := time.AfterFunc(c.idleTimeout(), func() {
			cancel(ErrStreamIdle)
		})
		defer idle.Stop()

time.AfterFunc runs the cancellation on the timer’s own goroutine. The timer is armed only while the program is waiting for bytes, by a wrapper around the response body:

// idleReader runs an idle timer only while it waits for the network,
// so the timer fires when the stream itself has gone quiet for the
// whole timeout, not when the program is busy with what it already
// has. If the timer has already fired, Reset runs it again, which is
// harmless: cancelling twice keeps the first cause.
type idleReader struct {
	r       io.Reader
	timer   *time.Timer
	timeout time.Duration
}

func (ir *idleReader) Read(p []byte) (int, error) {
	ir.timer.Reset(ir.timeout)
	n, err := ir.r.Read(p)
	ir.timer.Stop()
	return n, err //nolint:wrapcheck // Reader errors pass as is
}

Each Read resets the timer before it waits and stops it once data arrives, so the timeout measures how long the connection has been silent, and time the program spends handling an event never counts against it. The cause, ErrStreamIdle, is what lets the program report “the stream went quiet, worth retrying” instead of “cancelled”.

Work that must finish anyway

Some work should finish even when the run is being cancelled: saving the checkpoint that preserves what the run has done, or writing the record of a tool call that already changed something. context.WithoutCancel returns a context that keeps the parent’s values but never ends with it:

// save writes the run as it is now. A tool message whose call still
// has a marker is left out: the call's outcome is unknown, so on
// resume it goes to the resolution ladder instead of counting as
// answered. So is one that doesn't settle its call. The save ignores
// cancellation, since the save that keeps a stopping run's work must
// finish.
func (j *journal) save(ctx context.Context, phase Phase) error {
	if j.cp == nil {
		return nil
	}
	st := j.st
	st.Version, st.Phase = StateVersion, phase
	st.Messages = j.settledMessages()
	st.Usage, st.Turns = j.res.Usage, j.res.Turns
	st.Context, st.Final = j.res.Context, j.res.Final
	st.Pending = j.res.Pending
	err := j.cp.Save(context.WithoutCancel(ctx), st)
	if err != nil {
		return fmt.Errorf("save run: %w", err)
	}
	return nil
}

Use it narrowly, for short writes that must complete, and give the operation its own timeout if it can block, since nothing else will stop it.

A graceful shutdown goes one step further: in-flight work gets a grace period before it is cancelled. This function builds that from the same parts: a context detached with WithoutCancel, and a goroutine that cancels it with the original cause, after the grace period when the cause is a shutdown and at once for anything else:

// graced returns a context that ctx's cancellation reaches d late if
// its cause is ErrShutdown, and at once otherwise. stop ends the
// goroutine that watches ctx.
func graced(ctx context.Context, d time.Duration) (context.Context,
	func()) {
	inner, cancel := context.WithCancelCause(
		context.WithoutCancel(ctx))
	done := make(chan struct{})
	go func() {
		select {
		case <-done:
			return
		case <-ctx.Done():
		}
		cause := context.Cause(ctx)
		if errors.Is(cause, ErrShutdown) {
			select {
			case <-done:
				return
			case <-time.After(d):
			}
		}
		cancel(cause)
	}()
	return inner, func() { close(done); cancel(nil) }
}

The goroutine has an owner: the returned stop function ends it, so a run that finishes early doesn’t leave it waiting.

Testing timeouts without waiting

A test for a two-minute timeout shouldn’t take two minutes. testing/synctest gives the test its own fake clock, which moves forward only when every goroutine in the test is blocked, and then straight to the next timer:

// TestTimeoutSynctest checks the call deadline with testing/synctest,
// whose fake clock lets a two-minute timeout pass instantly.
func TestTimeoutSynctest(t *testing.T) {
	synctest.Test(t, func(t *testing.T) {
		c := &Client{HTTP: &http.Client{Transport: blockingTransport{}}}
		start := time.Now()
		_, err := c.Generate(t.Context(), agent.Request{})
		if !errors.Is(err, ErrCallTimeout) ||
			!errors.Is(err, context.DeadlineExceeded) {
			t.Errorf("err = %v", err)
		}
		if d := time.Since(start); d != 2*time.Minute {
			t.Errorf("the call gave up after %v, want 2m0s", d)
		}
	})
}

The test checks both the cause and the kind of error, and that the call gave up after exactly two minutes of fake time. It finishes instantly. t.Context() is a context that is cancelled when the test ends, so anything the test starts with it stops with the test.

Values in a context

A context can also carry values, and the rule is narrow: values that describe the request, which pass through layers that don’t know about them, never optional parameters. The agent uses one, the tool call a tool is running for, so that logging and tracing wrappers can name the call without every tool passing it along:

// callKey carries the tool call a tool is running for, so a wrapper can
// name it; Tool.Run sees only the arguments.
type callKey struct{}

// CallFrom returns the tool call the context's tool is running for.
func CallFrom(ctx context.Context) (ToolCall, bool) {
	c, ok := ctx.Value(callKey{}).(ToolCall)
	return c, ok
}

The key is an unexported type, callKey; since only this package can name that type, only this package can set or read the value, and CallFrom hands it out with its real type.

Mistakes to avoid

• context.Background() deep inside the code. It cuts the chain, and nothing under it can be cancelled. Take a ctx parameter instead.
• A context stored in a struct. Each call should get its own; a stored one outlives the request it belonged to.
• A forgotten cancel. Every WithTimeout, WithCancel, and their Cause forms returns one; defer it, or the timer and goroutine it holds linger. go vet reports a cancel that some path never calls.
• A timeout with no cause. “Context deadline exceeded” from somewhere in a run is hard to act on; a cause names the deadline.
• A tool that ignores its context. It can’t be stopped, and a run can’t finish shutting down while it runs.
• Parameters passed as context values. A value the function needs belongs in its signature; a context value is invisible to the caller and unchecked by the compiler.
• Comparing error strings. Check errors.Is(err, context.Canceled) and your own causes with errors.Is; the text of a cancellation error differs between transports.

Try it

From the code directory of a clone of the companion repository:

$ go test ./ch06/openrouter -run 'TestTimeoutSynctest|TestStreamIdleSynctest' -v
$ go test ./ch06/agent -run TestCancelPartwayThroughAReply -v

1. In TestTimeoutSynctest, change the client’s timeout and watch the test check the new duration, still instantly.
2. In ch06/openrouter/client.go, change WithTimeoutCause to WithTimeout and drop the cause: which tests fail, and what does whyStopped report for a call timeout now?
3. Write a tool that sleeps without checking its context, wrap it with a short timeout, and see when the run can finish.