From cbb6a58aaa9b5801e0226eeaa9c5ad95b29e6a7f Mon Sep 17 00:00:00 2001
From: townwish <townwish2023@outlook.com>
Date: Fri, 27 Mar 2026 17:05:02 +0800
Subject: [PATCH] fix(ui): keep stream delta order stable in TUI

---
 internal/app/run.go      |  1 +
 internal/app/run_test.go | 17 +++++++---
 ui/app.go                | 12 ++-----
 ui/app_train_test.go     | 67 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 83 insertions(+), 14 deletions(-)

diff --git a/internal/app/run.go b/internal/app/run.go
index 5b86d6c..48f77e2 100644
--- a/internal/app/run.go
+++ b/internal/app/run.go
@@ -198,6 +198,7 @@ func (a *Application) runTask(description string) {
 		}
 	})
 	if errors.Is(err, context.Canceled) {
+		emit(model.Event{Type: model.TaskDone})
 		persistSnapshot()
 		return
 	}
diff --git a/internal/app/run_test.go b/internal/app/run_test.go
index 9715888..859b29c 100644
--- a/internal/app/run_test.go
+++ b/internal/app/run_test.go
@@ -90,14 +90,23 @@ func TestInterruptTokenCancelsActiveTask(t *testing.T) {
 		t.Fatal("timed out waiting for task cancellation")
 	}
 
-	deadline := time.After(200 * time.Millisecond)
+	deadline := time.NewTimer(300 * time.Millisecond)
+	defer deadline.Stop()
+
+	foundTaskDone := false
 	for {
 		select {
 		case ev := <-app.EventCh:
-			if ev.Type == model.ToolError && strings.Contains(strings.ToLower(ev.Message), "canceled") {
-				t.Fatalf("expected interrupt cancellation to stay silent, got tool error %q", ev.Message)
+			switch ev.Type {
+			case model.TaskDone:
+				foundTaskDone = true
+			case model.ToolError:
+				t.Fatalf("expected no ToolError after interrupt, got %q", ev.Message)
+			}
+		case <-deadline.C:
+			if !foundTaskDone {
+				t.Fatal("timed out waiting for TaskDone after interrupt")
 			}
-		case <-deadline:
 			return
 		}
 	}
diff --git a/ui/app.go b/ui/app.go
index fe11b7a..b42edbf 100644
--- a/ui/app.go
+++ b/ui/app.go
@@ -244,7 +244,7 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			updated.updateViewport()
 			m = updated
 		}
-		return m, a.ensureWaitForEvent(cmd)
+		return m, cmd
 
 	case tea.MouseMsg:
 		var cmd tea.Cmd
@@ -286,15 +286,6 @@ func (a App) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	return a, tea.Batch(cmds...)
 }
 
-// ensureWaitForEvent wraps a cmd to always include waitForEvent,
-// so the UI keeps listening for backend events after key presses.
-func (a App) ensureWaitForEvent(cmd tea.Cmd) tea.Cmd {
-	if cmd == nil {
-		return a.waitForEvent
-	}
-	return tea.Batch(cmd, a.waitForEvent)
-}
-
 // chatWidth returns the width available for the chat panel.
 // In the stacked train layout the viewport is full-width.
 func (a App) chatWidth() int {
@@ -814,6 +805,7 @@ func (a App) handleEvent(ev model.Event) (tea.Model, tea.Cmd) {
 
 	case model.TaskDone:
 		a.state = a.state.WithThinking(false)
+		a.state = a.commitStreamingAgent()
 
 	case model.AgentThinking:
 		a.state = a.state.WithThinking(true)
diff --git a/ui/app_train_test.go b/ui/app_train_test.go
index 7743603..d4df018 100644
--- a/ui/app_train_test.go
+++ b/ui/app_train_test.go
@@ -557,6 +557,73 @@ func TestCtrlCSendsInterruptTokenForActiveTask(t *testing.T) {
 	}
 }
 
+func TestTaskDoneDispatchesQueuedInputWhenTrainNotBusy(t *testing.T) {
+	userCh := make(chan string, 1)
+	app := New(nil, userCh, "test", ".", "", "demo-model", 4096)
+	app.bootActive = false
+	app.trainView.Active = true
+	app.queuedInputs = []string{"continue"}
+	app.trainView.Runs = []model.TrainRunState{{
+		ID:    "primary",
+		Phase: model.TrainPhaseReady,
+	}}
+	app.trainView.ActiveRunID = "primary"
+
+	next, _ := app.handleEvent(model.Event{Type: model.TaskDone})
+	app = next.(App)
+
+	select {
+	case msg := <-userCh:
+		if msg != "continue" {
+			t.Fatalf("expected queued input to auto-dispatch after TaskDone, got %q", msg)
+		}
+	default:
+		t.Fatal("expected queued input to be auto-dispatched after TaskDone")
+	}
+	if got := len(app.queuedInputs); got != 0 {
+		t.Fatalf("expected queued input to be consumed, got %d items", got)
+	}
+}
+
+func TestKeyUpdateDoesNotForceEventResubscribe(t *testing.T) {
+	app := New(nil, nil, "test", ".", "", "demo-model", 4096)
+	app.bootActive = false
+
+	_, cmd := app.Update(tea.KeyMsg{Type: tea.KeyCtrlC})
+	if cmd != nil {
+		t.Fatal("expected key update to return only key-handler cmd without forced waitForEvent")
+	}
+}
+
+func TestTaskDoneCommitsStreamingAgentBeforeNextTurn(t *testing.T) {
+	app := New(nil, nil, "test", ".", "", "demo-model", 4096)
+	app.bootActive = false
+
+	next, _ := app.handleEvent(model.Event{Type: model.AgentReplyDelta, Message: "old"})
+	app = next.(App)
+
+	next, _ = app.handleEvent(model.Event{Type: model.TaskDone})
+	app = next.(App)
+
+	next, _ = app.handleEvent(model.Event{Type: model.UserInput, Message: "继续"})
+	app = next.(App)
+	next, _ = app.handleEvent(model.Event{Type: model.AgentReplyDelta, Message: "new"})
+	app = next.(App)
+
+	if got := len(app.state.Messages); got != 3 {
+		t.Fatalf("expected 3 messages (old agent, user, new agent), got %d: %#v", got, app.state.Messages)
+	}
+	if got := app.state.Messages[0].Content; got != "old" {
+		t.Fatalf("expected first agent message to stay unchanged after interrupt, got %q", got)
+	}
+	if app.state.Messages[0].Streaming {
+		t.Fatal("expected interrupted streaming message to be committed on TaskDone")
+	}
+	if got := app.state.Messages[2].Content; got != "new" {
+		t.Fatalf("expected new turn delta to start a fresh agent message, got %q", got)
+	}
+}
+
 func TestToolErrorClearsThinkingIndicator(t *testing.T) {
 	app := New(nil, nil, "test", ".", "", "demo-model", 4096)
 	app.bootActive = false