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 // --------------------------------------------------------------------------