fix(claude): stop a resumed turn on cancel once its deferred call settles

go/core/test/e2e/cancel_after_answer_test.go

@@ -0,0 +1,127 @@
+package e2e_test
+
+import (
+	"bytes"
+	"io"
+	"net/http"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	a2atype "github.com/a2aproject/a2a-go/v2/a2a"
+	a2agrpc "github.com/a2aproject/a2a-go/v2/a2agrpc/v1"
+	"github.com/a2aproject/a2a-go/v2/a2apb/v1/pbconv"
+	"google.golang.org/grpc/metadata"
+
+	adka2a "github.com/kagent-dev/kagent/go/adk/pkg/a2a"
+	kagenta2a "github.com/kagent-dev/kagent/go/api/a2a"
+)
+
+// Issue #3314: once an ask-user answer resumes a task, cancelling it must stop
+// the resumed turn, as it does for a task that never paused.
+func TestSessionCancelAfterAskUserAnswer(t *testing.T) {
+	t.Parallel()
+	forEachHarness(t, func(t *testing.T, harness testHarness) {
+		t.Parallel()
+		model := startAnswerBlockingMock(t, harness)
+		fixture := newInteractionFixture(t, harness, interactionTarget(t), model.url)
+		fixture.ctx = metadata.AppendToOutgoingContext(fixture.ctx, strings.ToLower(a2atype.SvcParamExtensions), adka2a.HITLExtensionURI)
+
+		// Given a task that asked a question and was suspended while waiting.
+		_, _, waiting := fixture.send(t, "Which database should we use for storage?")
+		if waiting.Status.State != a2atype.TaskStateInputRequired {
+			t.Fatalf("A2A task state = %s, want INPUT_REQUIRED", waiting.Status.State)
+		}
+		request := adka2a.GetAskUserRequest(waiting.Status.Message)
+		if request == nil {
+			t.Fatal("INPUT_REQUIRED task has no ask_user request")
+		}
+		assertActorSuspended(t, fixture)
+
+		// When the answer resumes it and the task is cancelled mid-turn.
+		model.holdNextCalls()
+		reply := adka2a.AttachHitlExtension(a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("PostgreSQL")), &kagenta2a.AskUserResponse{
+			Type: adka2a.HITLTypeAskUserResponse, ID: request.ID,
+			Answers: []kagenta2a.AskUserAnswer{{Answer: []string{"PostgreSQL"}}},
+		})
+		reply.TaskID, reply.ContextID = waiting.ID, waiting.ContextID
+		if _, err := a2agrpc.NewGRPCTransportFromClient(fixture.client).SendMessage(fixture.ctx, nil, &a2atype.SendMessageRequest{
+			Tenant: fixture.tenant, Message: reply, Config: &a2atype.SendMessageConfig{ReturnImmediately: true},
+		}); err != nil {
+			t.Fatalf("resume A2A task: %v", err)
+		}
+		select {
+		case <-model.held:
+		case <-time.After(time.Minute):
+			t.Fatal("resumed turn did not call the model")
+		}
+		cancelRequest, err := pbconv.ToProtoCancelTaskRequest(&a2atype.CancelTaskRequest{Tenant: fixture.tenant, ID: waiting.ID})
+		if err != nil {
+			t.Fatalf("build CancelTask request: %v", err)
+		}
+		canceledProto, err := fixture.client.CancelTask(fixture.ctx, cancelRequest)
+
+		// Assert that the resumed turn is cancelled, in the reply and in storage.
+		if err != nil {
+			t.Fatalf("cancel resumed task: %v", err)
+		}
+		canceled, err := pbconv.FromProtoTask(canceledProto)
+		if err != nil {
+			t.Fatalf("decode canceled task: %v", err)
+		}
+		if canceled.ID != waiting.ID || canceled.Status.State != a2atype.TaskStateCanceled {
+			t.Fatalf("canceled task = %s in %s, want %s in CANCELED", canceled.ID, canceled.Status.State, waiting.ID)
+		}
+		if persisted := getTask(t, fixture, waiting.ID); persisted.Status.State != a2atype.TaskStateCanceled {
+			t.Fatalf("persisted task state = %s, want CANCELED", persisted.Status.State)
+		}
+	})
+}
+
+// answerBlockingModel serves the ask-user mock until holdNextCalls, then holds
+// every later model call open, so the resumed turn stays in progress.
+type answerBlockingModel struct {
+	url  string
+	held <-chan struct{}
+	hold atomic.Bool
+}
+
+func (m *answerBlockingModel) holdNextCalls() { m.hold.Store(true) }
+
+func startAnswerBlockingMock(t *testing.T, harness testHarness) *answerBlockingModel {
+	t.Helper()
+	fixture := "mocks/invoke_golang_hitl_ask_user.json"
+	if harness.name == codexE2EHarness || harness.name == claudeE2EHarness {
+		fixture = "mocks/invoke_native_ask_user.json"
+	}
+	upstream := mockLLMProxy(t, startMockLLMServer(t, interactionMocks, fixture))
+	held, release := make(chan struct{}), make(chan struct{})
+	var heldOnce, releaseOnce sync.Once
+	model := &answerBlockingModel{held: held}
+	recorder := startModelRecorder(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+		r.Body = io.NopCloser(bytes.NewReader(body))
+		// Hold only the model call that carries the answer, so the replay has
+		// settled its question and the turn is doing ordinary work.
+		if !model.hold.Load() || !bytes.Contains(body, []byte("PostgreSQL")) {
+			upstream.ServeHTTP(w, r)
+			return
+		}
+		heldOnce.Do(func() { close(held) })
+		select {
+		case <-r.Context().Done():
+		case <-release:
+			http.Error(w, "test finished", http.StatusServiceUnavailable)
+		}
+	}), nil)
+	// Registered after the recorder, so it runs before the server closes.
+	t.Cleanup(func() { releaseOnce.Do(func() { close(release) }) })
+	model.url = reachableModelURL(t, recorder.URL)
+	return model
+}

go/harness/claude/internal/driver/deferred.go

@@ -407,6 +407,14 @@ func (p *deferredTurn) Resume(ctx context.Context, response runtime.InputRespons
 		logging.FromContext(ctx).ErrorContext(ctx, "failed to start claude replay; input remains waiting", "error", err)
 		return runtime.Outcome{Pending: p}, nil
 	}
+	// run returns the bare context.Canceled only when the replay settled its
+	// call before cancellation stopped the rest of the turn.
+	if err == context.Canceled {
+		if err := p.recordSettledCancellation(); err != nil {
+			return runtime.Outcome{}, &runtime.CancellationError{Err: err}
+		}
+		return runtime.Outcome{}, context.Canceled
+	}
 	if err != nil || outcome.Pending != nil {
 		return outcome, err
 	}
@@ -463,6 +471,10 @@ func (p *deferredTurn) Cancel(ctx context.Context, sink runtime.ModelCallSink) (
 			returnErr = p.abandon(returnErr)
 		}
 	}()
+	if err == context.Canceled {
+		// The stop decision settled before this request's own context ended.
+		return p.recordSettledCancellation()
+	}
 	if err != nil {
 		return err
 	}
@@ -482,6 +494,24 @@ func (p *deferredTurn) Cancel(ctx context.Context, sink runtime.ModelCallSink) (
 	return saveDeferred(b.directory, p.driver.config.MaxEventBytes, current)
 }
 
+// recordSettledCancellation tombstones a replay whose call consumed its
+// decision before cancellation stopped the process, so a restored harness
+// does not read the consumed call as an unfinished replay.
+func (p *deferredTurn) recordSettledCancellation() error {
+	b := p.driver.config.ApprovalBroker
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	current, err := p.current()
+	if err != nil {
+		return err
+	}
+	if current.Phase != phaseExecuting && current.Phase != phaseReleased {
+		return fmt.Errorf("Claude did not consume its deferred decision")
+	}
+	current.Phase, current.Canceling = phaseCanceled, true
+	return saveDeferred(b.directory, p.driver.config.MaxEventBytes, current)
+}
+
 // abandon prevents a failed task from being restored as pending input.
 func (p *deferredTurn) abandon(cause error) error {
 	b := p.driver.config.ApprovalBroker

go/harness/claude/internal/driver/deferred_cancel_test.go

@@ -0,0 +1,121 @@
+package driver
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/kagent-dev/kagent/go/harness/runtime"
+)
+
+// resultSink reports each tool result to onResult after recording it.
+type resultSink struct {
+	*recordingSink
+	onResult func(runtime.ToolResult)
+}
+
+func (s resultSink) ToolResult(event runtime.ToolResult) error {
+	if err := s.recordingSink.ToolResult(event); err != nil {
+		return err
+	}
+	s.onResult(event)
+	return nil
+}
+
+func TestDeferredReplayCancelsAfterItsCallSettles(t *testing.T) {
+	t.Run("Given a replay whose call has its result, when the turn is canceled during later work, assert that it stops promptly and records the cancellation", func(t *testing.T) {
+		b := newTestBroker(t)
+		r := testRecord()
+		require.NoError(t, saveDeferred(b.directory, 1<<20, &r))
+		// The replay returns the deferred call's result, then keeps working, as
+		// Claude does while it waits on its next model call or runs a tool.
+		dir, executable := writeFakeClaude(t, `#!/bin/sh
+printf '%s\n' '{"type":"system","subtype":"init","session_id":"11111111-1111-4111-8111-111111111111"}'
+printf '%s\n' '{"type":"user","message":{"content":[{"type":"tool_result","tool_use_id":"call","content":"ok","is_error":false}]}}'
+exec sleep 30
+`)
+		d := NewProcessDriver(ProcessConfig{Executable: executable, Workspace: dir, ApprovalBroker: b, MaxEventBytes: 1 << 20, MaxStderrBytes: 1024, InterruptGrace: 50 * time.Millisecond})
+		_, _, pending, err := d.RestorePending()
+		require.NoError(t, err)
+
+		ctx, cancel := context.WithCancel(t.Context())
+		defer cancel()
+		sink := resultSink{recordingSink: &recordingSink{}, onResult: func(result runtime.ToolResult) {
+			if result.ID != r.Tool.ID {
+				return
+			}
+			// Claude's hooks consume the decision before the call runs.
+			b.mu.Lock()
+			saved, err := loadDeferred(b.directory, 1<<20)
+			if err == nil {
+				saved.Phase = phaseReleased
+				err = saveDeferred(b.directory, 1<<20, saved)
+			}
+			b.mu.Unlock()
+			require.NoError(t, err)
+			cancel()
+		}}
+		started := time.Now()
+		_, err = pending.Resume(ctx, &runtime.ApprovalDecision{ID: r.RequestID, Approved: true}, sink)
+
+		assertSettledCancellation(t, b, err, started)
+	})
+
+	t.Run("Given a replay whose call has not returned yet, when the turn is canceled and the result then arrives, assert that it stops promptly and records the cancellation", func(t *testing.T) {
+		b := newTestBroker(t)
+		r := testRecord()
+		require.NoError(t, saveDeferred(b.directory, 1<<20, &r))
+		// The result arrives a second after start, after cancellation began cleanup.
+		dir, executable := writeFakeClaude(t, `#!/bin/sh
+printf '%s\n' '{"type":"system","subtype":"init","session_id":"11111111-1111-4111-8111-111111111111"}'
+sleep 1
+printf '%s\n' '{"type":"user","message":{"content":[{"type":"tool_result","tool_use_id":"call","content":"ok","is_error":false}]}}'
+exec sleep 30
+`)
+		d := NewProcessDriver(ProcessConfig{Executable: executable, Workspace: dir, ApprovalBroker: b, MaxEventBytes: 1 << 20, MaxStderrBytes: 1024, InterruptGrace: 50 * time.Millisecond})
+		_, _, pending, err := d.RestorePending()
+		require.NoError(t, err)
+
+		ctx, cancel := context.WithCancel(t.Context())
+		defer cancel()
+		go func() {
+			time.Sleep(200 * time.Millisecond)
+			// Claude's hooks consume the decision before the call runs.
+			b.mu.Lock()
+			saved, err := loadDeferred(b.directory, 1<<20)
+			if err == nil {
+				saved.Phase = phaseReleased
+				err = saveDeferred(b.directory, 1<<20, saved)
+			}
+			b.mu.Unlock()
+			if err != nil {
+				t.Errorf("mark decision consumed: %v", err)
+			}
+			cancel()
+		}()
+		started := time.Now()
+		_, err = pending.Resume(ctx, &runtime.ApprovalDecision{ID: r.RequestID, Approved: true}, &recordingSink{})
+
+		assertSettledCancellation(t, b, err, started)
+	})
+}
+
+func assertSettledCancellation(t *testing.T, b *ApprovalBroker, err error, started time.Time) {
+	t.Helper()
+	require.ErrorIs(t, err, context.Canceled)
+	var cleanupErr *runtime.CancellationError
+	require.False(t, errors.As(err, &cleanupErr), "cancellation reported a cleanup failure: %v", err)
+	require.Less(t, time.Since(started), 5*time.Second)
+	saved, err := loadDeferred(b.directory, 1<<20)
+	require.NoError(t, err)
+	require.Equal(t, phaseCanceled, saved.Phase)
+	restored := NewProcessDriver(ProcessConfig{ApprovalBroker: b, MaxEventBytes: 1 << 20})
+	taskID, contextID, restoredPending, err := restored.RestorePending()
+	require.NoError(t, err)
+	require.Empty(t, taskID)
+	require.Empty(t, contextID)
+	require.Nil(t, restoredPending)
+}

go/harness/claude/internal/driver/process.go

@@ -71,6 +71,9 @@ type processSession struct {
 	stopOnce  sync.Once
 	canceled  bool
 	seed      *DeferredTool
+	// seedSettled records that the replayed deferred call has its result in
+	// the transcript, so the rest of the turn is ordinary work.
+	seedSettled bool
 }
 
 // processStartError proves that no native process ran for this attempt.
@@ -157,7 +160,9 @@ func (d *ProcessDriver) run(ctx context.Context, turn runtime.Turn, sink runtime
 	defer func() {
 		if !started && returnErr != nil {
 			returnErr = &processStartError{err: returnErr}
-		} else if seed != nil && ctx.Err() != nil && returnErr != nil {
+		} else if seed != nil && ctx.Err() != nil && returnErr != nil && !errors.Is(returnErr, context.Canceled) {
+			// context.Canceled means a settled replay stopped as an ordinary
+			// turn; Resume records that cancellation itself.
 			var cancellationErr *runtime.CancellationError
 			if !errors.As(returnErr, &cancellationErr) {
 				returnErr = &runtime.CancellationError{Err: returnErr}
@@ -347,8 +352,10 @@ func sendPrompt(ctx context.Context, stdin io.WriteCloser, message []byte, gate
 
 func (d *ProcessDriver) consume(ctx context.Context, session *processSession, sink runtime.EventSink) (outcome runtime.Outcome, returnErr error) {
 	cleaning := false
+	// canceledErr is the caller's cancellation, kept while cleanup replaces ctx.
+	var canceledErr error
 	defer func() {
-		if cleaning && returnErr != nil {
+		if cleaning && returnErr != nil && returnErr != canceledErr {
 			returnErr = &runtime.CancellationError{Err: returnErr}
 		}
 	}()
@@ -379,6 +386,14 @@ func (d *ProcessDriver) consume(ctx context.Context, session *processSession, si
 						}
 						session.sessionID = item.event.SessionID
 					}
+					if session.seed != nil && item.event.Kind == EventToolActivity && item.event.ToolPhase == "completed" && item.event.ToolID == session.seed.ID {
+						session.seedSettled = true
+						if cleaning {
+							// Cancellation arrived before the result did; stop now,
+							// as a settled replay would have.
+							return runtime.Outcome{}, canceledErr
+						}
+					}
 					if outcome != nil {
 						session.terminal = outcome
 					}
@@ -408,13 +423,16 @@ func (d *ProcessDriver) consume(ctx context.Context, session *processSession, si
 			return *session.terminal, nil
 		case <-ctx.Done():
 			session.canceled = true
-			if session.seed == nil || cleaning {
+			// Once the replayed call has its result, the process stops like any
+			// other turn. Waiting for a hook could take as long as the model call
+			// or tool now running, which is longer than cleanup allows.
+			if session.seed == nil || cleaning || session.seedSettled {
 				return runtime.Outcome{}, ctx.Err()
 			}
 			// A deferred replay must settle its transcript before cancellation is
 			// published. Deny subsequent tools and drain this process; killing it
 			// here would leave its journal reserved forever.
-			cleaning = true
+			cleaning, canceledErr = true, ctx.Err()
 			if err := d.cancelReplay(session.seed); err != nil {
 				return runtime.Outcome{}, err
 			}