From ac238f69030416a58ef9e02ff673fedf9701ed62 Mon Sep 17 00:00:00 2001
From: Levi Neely <lkn@darkstar.example.net>
Date: Sun, 26 Apr 2026 07:30:59 +0200
Subject: [PATCH] agent: roll back session on turn failure
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 pkg/agent/core.go | 21 +++++++++++++++++++--
 pkg/agent/loop.go | 18 ------------------
 2 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/pkg/agent/core.go b/pkg/agent/core.go
index a818478..8637bfa 100644
--- a/pkg/agent/core.go
+++ b/pkg/agent/core.go
@@ -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{
diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 2a903f3..d807221 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -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 {