From 9b430f3883a73b03a25de398e2592c3e2bbe7da0 Mon Sep 17 00:00:00 2001 From: Ed Zynda Date: Thu, 26 Feb 2026 13:32:38 +0300 Subject: [PATCH] fix: ESC cancel now actually cancels agent response context Streaming callbacks were not checking ctx.Err(), so the fantasy library kept processing after the user cancelled. Additionally, context.Canceled was surfaced as a StepErrorEvent, printing an error instead of silently ending the turn. Add ctx.Err() checks to all streaming callbacks so the fantasy stream loop breaks immediately on cancel. Introduce StepCancelledEvent so the TUI flushes partial content and returns to input without an error message. --- internal/agent/agent.go | 12 ++++++ internal/app/app.go | 7 ++++ internal/app/events.go | 6 +++ internal/ui/model.go | 10 +++++ internal/ui/model_test.go | 80 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 115 insertions(+) diff --git a/internal/agent/agent.go b/internal/agent/agent.go index a7caef69..64a26647 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -163,6 +163,9 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan // Text streaming callback OnTextDelta: func(id, text string) error { + if ctx.Err() != nil { + return ctx.Err() + } if onStreamingResponse != nil { onStreamingResponse(text) } @@ -171,6 +174,9 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan // Tool call complete - the tool has been parsed and is about to execute OnToolCall: func(tc fantasy.ToolCallContent) error { + if ctx.Err() != nil { + return ctx.Err() + } currentToolName = tc.ToolName currentToolArgs = tc.Input @@ -189,6 +195,9 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan // Tool result - tool execution completed OnToolResult: func(tr fantasy.ToolResultContent) error { + if ctx.Err() != nil { + return ctx.Err() + } // Notify tool execution finished if onToolExecution != nil { onToolExecution(tr.ToolName, false) @@ -205,6 +214,9 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan // Step callbacks for content that accompanies tool calls OnStepFinish: func(step fantasy.StepResult) error { + if ctx.Err() != nil { + return ctx.Err() + } // Check if step has text content alongside tool calls text := step.Content.Text() toolCalls := step.Content.ToolCalls() diff --git a/internal/app/app.go b/internal/app/app.go index 651a082f..2bd40ea8 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -270,6 +270,13 @@ func (a *App) runPrompt(prompt string) { result, err := a.executeStep(stepCtx, prompt, prog, nil) if err != nil { + if stepCtx.Err() != nil { + // Step was cancelled by the user (e.g. double-ESC). Send a + // cancellation event so the TUI can cut off the response + // cleanly without printing an error. + a.sendEvent(StepCancelledEvent{}) + return + } a.sendEvent(StepErrorEvent{Err: err}) return } diff --git a/internal/app/events.go b/internal/app/events.go index d791f1c3..be6c8f73 100644 --- a/internal/app/events.go +++ b/internal/app/events.go @@ -69,6 +69,12 @@ type StepErrorEvent struct { Err error } +// StepCancelledEvent is sent when an agent step is cancelled by the user +// (e.g. via double-ESC). The TUI should flush any partially streamed content, +// cut off the agent message where it was, and return to input state without +// displaying an error. +type StepCancelledEvent struct{} + // QueueUpdatedEvent is sent whenever the message queue length changes. // The TUI uses this to update the queue badge display. type QueueUpdatedEvent struct { diff --git a/internal/ui/model.go b/internal/ui/model.go index 615002ed..d9b85714 100644 --- a/internal/ui/model.go +++ b/internal/ui/model.go @@ -441,6 +441,16 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.state = stateInput m.canceling = false + case app.StepCancelledEvent: + // User cancelled the step (double-ESC). Flush any partial content, + // cut off the response where it was, and return to input with no error. + cmds = append(cmds, m.flushStreamContent()) + if m.stream != nil { + m.stream.Reset() + } + m.state = stateInput + m.canceling = false + case app.StepErrorEvent: // Flush streamed text, print the error, reset stream, return to input. cmds = append(cmds, m.flushStreamContent()) diff --git a/internal/ui/model_test.go b/internal/ui/model_test.go index 6d808105..91ec40e1 100644 --- a/internal/ui/model_test.go +++ b/internal/ui/model_test.go @@ -195,6 +195,86 @@ func TestStepError_nilErr(t *testing.T) { } } +// -------------------------------------------------------------------------- +// StepCancelledEvent +// -------------------------------------------------------------------------- + +// TestStateTransition_WorkingToInput_StepCancelled verifies that StepCancelledEvent +// transitions from stateWorking back to stateInput and resets the stream component. +func TestStateTransition_WorkingToInput_StepCancelled(t *testing.T) { + ctrl := &stubAppController{} + m, stream, _ := newTestAppModel(ctrl) + m.state = stateWorking + + m = sendMsg(m, app.StepCancelledEvent{}) + + if m.state != stateInput { + t.Fatalf("expected stateInput after StepCancelledEvent, got %v", m.state) + } + if stream.resetCalled != 1 { + t.Fatalf("expected stream.Reset() called once, got %d", stream.resetCalled) + } +} + +// TestStepCancelled_clearsCanceling verifies that StepCancelledEvent clears +// the canceling flag. +func TestStepCancelled_clearsCanceling(t *testing.T) { + ctrl := &stubAppController{} + m, _, _ := newTestAppModel(ctrl) + m.state = stateWorking + m.canceling = true + + m = sendMsg(m, app.StepCancelledEvent{}) + + if m.canceling { + t.Fatal("expected canceling=false after StepCancelledEvent") + } +} + +// TestStepCancelled_flushesStreamContent verifies that StepCancelledEvent +// flushes accumulated stream content via tea.Println (non-nil cmd). +func TestStepCancelled_flushesStreamContent(t *testing.T) { + ctrl := &stubAppController{} + m, stream, _ := newTestAppModel(ctrl) + m.state = stateWorking + stream.renderedContent = "partial assistant response" + + _, cmd := m.Update(app.StepCancelledEvent{}) + + if cmd == nil { + t.Fatal("expected non-nil cmd (tea.Println) on StepCancelledEvent with stream content") + } +} + +// TestStepCancelled_noStreamContent_noCmd verifies that StepCancelledEvent with +// no accumulated stream content produces a nil cmd (nothing to flush). +func TestStepCancelled_noStreamContent_noCmd(t *testing.T) { + ctrl := &stubAppController{} + m, _, _ := newTestAppModel(ctrl) + m.state = stateWorking + + _, cmd := m.Update(app.StepCancelledEvent{}) + + if cmd != nil { + t.Fatal("expected nil cmd on StepCancelledEvent with no stream content") + } +} + +// TestStepCancelled_noErrorPrinted verifies that StepCancelledEvent does NOT +// produce an error message (unlike StepErrorEvent). +func TestStepCancelled_noErrorPrinted(t *testing.T) { + ctrl := &stubAppController{} + m, _, _ := newTestAppModel(ctrl) + m.state = stateWorking + + // With no stream content, cmd should be nil (no flush, and no error print). + _, cmd := m.Update(app.StepCancelledEvent{}) + + if cmd != nil { + t.Fatal("expected nil cmd for StepCancelledEvent with no stream content — should not print error") + } +} + // -------------------------------------------------------------------------- // ESC cancel flow // --------------------------------------------------------------------------