agent: roll back session on turn failure

Snapshot session state before each turn. On any error from run(),
restore to the pre-turn snapshot so the session is clean for the next
attempt. Remove recordInterruption — those messages were appended to
state that gets rolled back anyway, making them dead code.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Levi Neely 2026-04-26 07:30:59 +02:00
parent 0fba869ab7
commit ac238f6903
2 changed files with 19 additions and 20 deletions

View File

@ -11,6 +11,7 @@ import (
"path/filepath"
"runtime"
"runtime/debug"
"slices"
"strconv"
"strings"
"sync"
@ -733,6 +734,14 @@ func (s *agent) executeTurn(ctx context.Context, input string, handler EventHand
s.setState("thinking")
// Snapshot session state before this turn modifies it. Restored on failure
// so the session is clean for the next attempt.
snapSession := s.session
var snapMessages []backend.Message
if s.session != nil {
snapMessages = slices.Clone(s.session.messages)
}
if s.session == nil {
for _, msg := range s.startupMessages {
s.log.Debug("startup: %s", msg)
@ -831,8 +840,16 @@ func (s *agent) executeTurn(ctx context.Context, input string, handler EventHand
replyBuf.Reset()
s.setState("idle")
if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, ErrInterrupted) {
handler(Event{Role: "error", Content: err.Error()})
if err != nil {
if snapSession == nil {
s.session = nil
} else {
s.session.messages = snapMessages
}
if !errors.Is(err, context.Canceled) && !errors.Is(err, ErrInterrupted) {
handler(Event{Role: "error", Content: err.Error()})
}
return ""
}
stopResult := s.hooks.Run(ctx, HookPostTurn, map[string]string{

View File

@ -46,7 +46,6 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
break
}
if ctx.Err() != nil {
recordInterruption(state, "request")
return ctx.Err()
}
var rlErr *backend.RateLimitError
@ -112,7 +111,6 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
})
}
state.update(msg, results) //nolint:errcheck
recordInterruption(state, "stream")
return fmt.Errorf("step %d: stream interrupted: %w", step, ctx.Err())
}
@ -242,7 +240,6 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
}
if interrupted {
recordInterruption(state, "tools")
return ctx.Err()
}
@ -260,21 +257,6 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
return nil
}
// recordInterruption appends a note to history so the model knows what happened.
func recordInterruption(state state, phase string) {
var note string
switch phase {
case "stream":
note = "Note: generation was interrupted before completion. Treat the previous assistant output as partial."
case "tools":
note = "Note: tool execution was interrupted before completion. Tool outputs may be missing or cancelled."
case "request":
note = "Note: the previous assistant turn was interrupted before any response was received."
default:
note = "Note: the previous assistant turn was interrupted."
}
state.update(backend.Message{Role: "assistant", Content: note}, nil) //nolint:errcheck
}
func emit(cfg agentConfig, msg Event) {
if cfg.Output != nil {