From 610a3862b1023aa45dec48611e26abc36f3bbebb Mon Sep 17 00:00:00 2001 From: Arunkumar Eli Date: Wed, 11 Nov 2026 07:07:52 +0000 Subject: [PATCH] fix(claude): stop a resumed turn on cancel once its deferred call settles Canceling a task that an approval or ask-user answer had resumed waited up to 30s for Claude to reach a hook and stop by itself. When Claude was in a model call or a long tool after the replayed call returned, no hook fired, so CancelTask failed with "runtime cancellation cleanup failed: context deadline exceeded" and the session could not be deleted. Once the replayed call has its result, the rest of the turn is ordinary work, so the driver now stops it like any other turn and records the replay as canceled, so a restored harness does not report the session as abandoned. A cancel that arrives before the result still waits for Claude to settle it. Fixes #3314 Signed-off-by: Arunkumar Eli --- go/core/test/e2e/cancel_after_answer_test.go | 127 ++++++++++++++++++ go/harness/claude/internal/driver/deferred.go | 30 +++++ .../internal/driver/deferred_cancel_test.go | 121 +++++++++++++++++ go/harness/claude/internal/driver/process.go | 26 +++- 4 files changed, 300 insertions(+), 4 deletions(-) create mode 100644 go/core/test/e2e/cancel_after_answer_test.go create mode 100644 go/harness/claude/internal/driver/deferred_cancel_test.go diff --git a/go/core/test/e2e/cancel_after_answer_test.go b/go/core/test/e2e/cancel_after_answer_test.go new file mode 100644 index 0000000000..64d16d8cf7 --- /dev/null +++ b/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 +} diff --git a/go/harness/claude/internal/driver/deferred.go b/go/harness/claude/internal/driver/deferred.go index 376b35b77f..0ebae2db8c 100644 --- a/go/harness/claude/internal/driver/deferred.go +++ b/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 diff --git a/go/harness/claude/internal/driver/deferred_cancel_test.go b/go/harness/claude/internal/driver/deferred_cancel_test.go new file mode 100644 index 0000000000..e04917cb6a --- /dev/null +++ b/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) +} diff --git a/go/harness/claude/internal/driver/process.go b/go/harness/claude/internal/driver/process.go index cce04fab15..5b2404ead3 100644 --- a/go/harness/claude/internal/driver/process.go +++ b/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 }