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 }