fix(runner): save the session when the caller cancels the turn

[zalbiraw]

Summary

A turn that the caller cancels is now saved to the session store. Before, it was lost, including the user's message.

Runner.Run saves the session in a deferred call when a run ends. That save used the run's context. A caller stops a turn by canceling that context, so the save ran on an already-canceled context and failed. In ai-agent on Postgres, every canceled A2A turn logged session save failed after consumer stopped: begin tx: context canceled, and the session kept only the previous turn. A battle test on k3d showed this for a cancel on the same pod and for one from another pod.

Change

• The deferred save runs on context.WithoutCancel(ctx), so it keeps the context's values but not its cancellation. A 10-second timeout (sessionSaveTimeout) bounds it.
• A canceled turn is saved the same way a failed turn is: the user's message, plus the messages the agent finished. Partial streamed text isn't in the session in either case.
• Nothing else changes. The per-message saves during the run still use the run's context, and the deferred save still doesn't call yield after the caller stops reading.

Testing

• TestRun_CanceledTurnIsSaved cancels a turn during a model call, once with a caller that keeps reading and once with a caller that stops. It checks that the turn is saved once, on a context the cancel didn't reach, and that the user's message is kept. It fails without the fix.
• go test -race -short ./runner/ passes.
• go test -short ./... passes, except two Schema Registry tests in adapter/a2a/kvstore. Their Redpanda test container exits on start on my machine, the same as on main.
• golangci-lint run --new-from-rev=origin/main ./runner/... reports nothing, apart from exhaustruct_v5. That is the renamed exhaustruct, which .golangci.yaml disables; golangci-lint 2.14 reports it under the new name and the repo pins 2.6.1.

cloudv2 picks this up through a module bump in the ADP autoscaling stack, which adds an ai-agent test against Postgres.

🤖 Generated with Claude Code

[redpanda-ai-merge-bot]

🟡 AI review (dry run): would have approved — no approval posted

Field	Value
Config version	10
Prompt version	agent.md@ee051ac95fdf
Model	claude-opus-5-5
Engine	agent-action (agent-action@base-action@756cc22e19660d20e8cc9496b4f242475a7f7790)
AI verdict	approve (confidence 0.88, reviewed fully: True, scope: localized)
Files / lines (all)	2 / 100
CI at this commit	passed
Reviewable diff	4499 chars (limit 120000)
Reviewable files / lines	2 / 100
Generated (CI-verified) files / lines	0 / 0
Test / source lines	85 / 15 (tests changed with source: True)
Dependency update	False
Run	https://github.com/redpanda-data/ai-sdk-go/actions/runs/37451559595

Review summary

Small, focused fix: Runner.Run's deferred final session save now runs on context.WithTimeout(context.WithoutCancel(ctx), 10s) instead of the run's own (possibly already canceled) ctx, so canceled turns (and the user's message) are persisted. Context values are kept and the save is bounded. Per-message saves still use the run ctx, and the consumerStopped guard around yield is unchanged. go.mod targets Go 1.26, so WithoutCancel is available. The new table test covers both a caller that keeps reading and one that breaks, and checks that exactly one save happens on a non-canceled context with the user message. CI is green.

Concerns

• After a cancel, the iterator can now block for up to sessionSaveTimeout (10s) on a slow store before returning. That is acceptable, but it is a small behaviour change for callers that expect a fast return after cancel.
• The timeout is a hard-coded constant and can't be configured. Consider a runner option later if 10s doesn't suit every store.

Evidence cited by the engine

• runner/runner.go:210 — Deferred save uses WithoutCancel(ctx) with a bounded timeout and cancel is deferred, so nothing leaks
• runner/runner.go:217 — The consumerStopped guard still stops yield from being called after the consumer stops, so it can't panic
• runner/runner.go:242 — Incremental per-message saves still use the run ctx, as the PR says
• runner/runner_test.go:600 — New test cancels mid model call in both caller modes and checks one save on a non-canceled ctx that keeps the user message

Shadow engines (recorded for comparison; did not decide)

Engine	Verdict	Confidence	Reviewed fully	Scope
single-call	comment	0.55	True	localized

Generated by the AI review gate. The bot only posts an approval; the author merges. This is an audit record; the branch ruleset is what enforces it.

[zalbiraw]

@claude review this

[claude]

Claude finished @zalbiraw's task in 1m 56s —— View job

---

Review of PR #248: save the session when the caller cancels the turn

• Read the diff
• Read the surrounding Runner.Run code and how llmagent handles a cancel
• Run the tests. I couldn't: go test needs approval in this job's tool permissions. CI at this commit is green.
• Post the review

Verdict: looks good to merge. The fix is small and correct, and the test fails without it. I found no blocking issues, only the nits below.

Correctness

• The fix is right. context.WithTimeout(context.WithoutCancel(ctx), sessionSaveTimeout) at runner/runner.go:210 keeps the request-scoped values (trace/span, tenant, logger fields) and drops the cancel. defer cancel() releases the timer.

• No data race on sess. The agent's iterator is driven synchronously by the range at runner/runner.go:228. By the time the deferred save runs, agent.Run has returned and nothing else writes to sess.Messages.

• A canceled turn leaves a session the next turn can use. Saving more state after a cancel could leave a dangling tool_use. The agent already prevents that:

  • If the cancel lands during tool execution, interruptedToolParts (agent/llmagent/llmagent.go:930) gives every request a synthetic "interrupted" result before the tool message is appended.
  • The recovery path at llmagent.go:942 repairs anything older that is still incomplete.

  So the new saves don't add a malformed-history risk.

• The yield guard still holds. consumerStopped still guards the deferred yield, so the "caller stops reading" case can't panic. Your second subtest covers it.

Behaviour changes worth knowing (not blocking)

1. Returning after a cancel can now take up to 10s. If the store is slow or unreachable, the iterator blocks in the deferred save before it returns. A cancel used to return almost at once, because the save failed fast. For A2A cancel handlers that wait for Run to finish, this is what you want. If a caller needs a fast cancel path, it is new latency. Consider one line in the Run doc comment (the "Error Handling" section near runner.go:121) saying the final save survives a cancel and is bounded by sessionSaveTimeout.

2. An incremental save can still race the cancel. The per-message save at runner.go:242 still uses ctx. If the cancel arrives right after the agent emits a MessageEvent:

   • that save fails and the caller gets ErrSessionSave;
   • the deferred save then succeeds.

   The caller sees a save error even though the session was persisted. The window is small and the outcome is correct, so it is fine to leave. Note it if you get reports of spurious ErrSessionSave on cancel.

3. The timeout is hard-coded. 10s is a sensible default. A WithSessionSaveTimeout option can wait until someone needs it.

Test (runner/runner_test.go:598)

The test is well built: it uses a table, calls t.Parallel(), guards state with a mutex, copies state.Messages so later appends don't change it, and covers both caller modes. Two optional ways to make it check more of the design:

• Assert that the timeout bound is applied, so a later refactor can't drop it quietly:

  _, hasDeadline := ctx.Deadline()
  // record and assert hasDeadline == true

• Assert that context values survive. Put a value on the run's ctx and check it in saveFunc. That is the reason for WithoutCancel over context.Background(), and right now nothing tests it.

Add these test assertions →

Nits

• The comment at runner.go:205-208 is clear. "the turn and the user's message would be lost" describes the old bug, not the current code, so it could read "would otherwise be lost". Purely cosmetic.