fix: Progressive runs ignore cancellation while persisting synthetic follow-up/finalization turns

Author: sam-saffron-jarvis

cmd/progressive.go

@@ -19,9 +19,10 @@ const (
 	progressiveStopWhenDone    progressiveStopWhen = "done"
 	progressiveStopWhenTimeout progressiveStopWhen = "timeout"
 
-	progressiveDefaultFinalizeGrace = 5 * time.Minute
-	progressiveMaxFinalizeBudget    = 5 * time.Minute
-	progressiveMinFinalizeBudget    = 5 * time.Second
+	progressiveDefaultFinalizeGrace    = 5 * time.Minute
+	progressiveMaxFinalizeBudget       = 5 * time.Minute
+	progressiveMinFinalizeBudget       = 5 * time.Second
+	progressiveSyntheticMessageTimeout = 5 * time.Second
 )
 
 type askProgressiveOptions struct {
@@ -296,7 +297,10 @@ func runProgressiveSession(ctx context.Context, engine *llm.Engine, req llm.Requ
 		continueMsg := llm.UserText(expandProgressiveTemplate(opts.ContinueWith, mainCtx))
 		history = append(history, continueMsg)
 		if opts.OnSyntheticUserMessage != nil {
-			if err := opts.OnSyntheticUserMessage(context.Background(), continueMsg); err != nil {
+			msgCtx, cancel := progressiveSyntheticUserMessageContext(mainCtx)
+			err := opts.OnSyntheticUserMessage(msgCtx, continueMsg)
+			cancel()
+			if err != nil {
 				return buildProgressiveRunResult(opts.SessionID, exitReasonNatural, false, tracker.latest, lastText), err
 			}
 		}
@@ -425,7 +429,10 @@ func attemptProgressiveFinalization(parentCtx context.Context, engine *llm.Engin
 	finalPrompt := buildProgressiveFinalizePrompt(tracker.latest)
 	finalizeMsg := llm.UserText(finalPrompt)
 	if opts.OnSyntheticUserMessage != nil {
-		if err := opts.OnSyntheticUserMessage(context.Background(), finalizeMsg); err != nil {
+		msgCtx, cancel := progressiveSyntheticUserMessageContext(finalizeCtx)
+		err := opts.OnSyntheticUserMessage(msgCtx, finalizeMsg)
+		cancel()
+		if err != nil {
 			return false, ""
 		}
 	}
@@ -513,6 +520,15 @@ func progressiveHasRemainingBudget(parent context.Context, reserve time.Duration
 	return time.Until(deadline) > reserve
 }
 
+// progressiveSyntheticUserMessageContext caps synthetic user-message persistence
+// to a short window while still honoring the caller's cancellation/deadline.
+func progressiveSyntheticUserMessageContext(parent context.Context) (context.Context, context.CancelFunc) {
+	if parent == nil {
+		parent = context.Background()
+	}
+	return context.WithTimeout(parent, progressiveSyntheticMessageTimeout)
+}
+
 func progressiveFinalizeReserve(total time.Duration) time.Duration {
 	if total <= 0 {
 		return 0

cmd/progressive_test.go

@@ -3,10 +3,12 @@ package cmd
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"testing"
 	"time"
 
 	"github.com/samsaffron/term-llm/internal/llm"
+	toolpkg "github.com/samsaffron/term-llm/internal/tools"
 )
 
 type progressiveTestTool struct{}
@@ -175,6 +177,87 @@ func TestRunProgressiveSessionTimeoutDoesNotStartDetachedFinalizationAfterDeadli
 	}
 }
 
+func TestRunProgressiveSessionSyntheticContinuationPersistenceHonorsContext(t *testing.T) {
+	provider := llm.NewMockProvider("mock").WithCapabilities(llm.Capabilities{
+		ToolCalls:          true,
+		SupportsToolChoice: true,
+	})
+	provider.AddToolCall("progress-1", "update_progress", map[string]any{
+		"state": map[string]any{
+			"step": "draft",
+		},
+		"reason":  "milestone",
+		"message": "draft saved",
+	})
+	provider.AddTextResponse("draft answer")
+
+	engine := llm.NewEngine(provider, nil)
+	ctx, cancel := context.WithTimeout(context.Background(), 5050*time.Millisecond)
+	defer cancel()
+
+	done := make(chan error, 1)
+	go func() {
+		_, err := runProgressiveSession(ctx, engine, llm.Request{
+			Messages: []llm.Message{llm.UserText("Investigate X")},
+			MaxTurns: 8,
+			ToolChoice: llm.ToolChoice{
+				Mode: llm.ToolChoiceAuto,
+			},
+		}, progressiveRunOptions{
+			StopWhen: progressiveStopWhenTimeout,
+			OnSyntheticUserMessage: func(ctx context.Context, msg llm.Message) error {
+				<-ctx.Done()
+				return ctx.Err()
+			},
+		})
+		done <- err
+	}()
+
+	select {
+	case err := <-done:
+		if !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("runProgressiveSession error = %v, want deadline exceeded", err)
+		}
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("synthetic continuation persistence did not stop when its context expired")
+	}
+}
+
+func TestAttemptProgressiveFinalizationSyntheticPersistenceHonorsContext(t *testing.T) {
+	provider := llm.NewMockProvider("mock")
+	engine := llm.NewEngine(provider, nil)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
+	done := make(chan struct{})
+	go func() {
+		finalized, text := attemptProgressiveFinalization(ctx, engine, toolpkg.NewFinalizeProgressTool(), llm.Request{
+			Messages: []llm.Message{llm.UserText("Investigate X")},
+		}, nil, progressiveRunOptions{
+			OnSyntheticUserMessage: func(ctx context.Context, msg llm.Message) error {
+				<-ctx.Done()
+				return ctx.Err()
+			},
+		}, newProgressTracker(), 20*time.Millisecond, exitReasonTimeout)
+		if finalized {
+			t.Errorf("finalized = true, want false")
+		}
+		if text != "" {
+			t.Errorf("text = %q, want empty", text)
+		}
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		if len(provider.Requests) != 0 {
+			t.Fatalf("provider saw %d requests, want synthetic persistence failure to stop finalization before the model call", len(provider.Requests))
+		}
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("synthetic finalization persistence did not stop when its context expired")
+	}
+}
+
 func TestProgressiveFinalizationContextNaturalCompletionDetachesFromParent(t *testing.T) {
 	parent, cancelParent := context.WithCancel(context.Background())
 	finalizeCtx, finalizeCancel := progressiveFinalizationContext(parent, time.Second, exitReasonNatural)