feat(agent): enhance event routing for reasoning and answers in streaming

- Introduced new tests to ensure correct routing of reasoning and answer content during streaming, addressing UX issues where answers appeared under the Thinking phase before transitioning. - Implemented a ThinkStreamSplitter to separate inline reasoning from answer text, allowing for real-time updates to the UI without duplication. - Updated the event emission logic to prevent re-emission of answers that were already streamed live, ensuring a seamless user experience. - Enhanced the ChatResponse struct to include flags for streamed answers, improving the handling of final answer events. - Added comprehensive unit tests for the ThinkStreamSplitter to validate its functionality across various scenarios.
2026-06-03 22:14:35 +08:00
parent 710d759b8e
commit 876fd0276f
6 changed files with 457 additions and 33 deletions
@@ -219,6 +219,125 @@ func TestStreamThinkingToEventBus_PropagatesFinishReason(t *testing.T) {
 	}
 }

+// TestStreamThinkingToEventBus_RoutesReasoningAndAnswerSeparately is the
+// regression guard for the "answer first shows under Thinking, then jumps to
+// the answer area" UX bug. A natural-stop response that carries reasoning in
+// the dedicated reasoning channel (ResponseTypeThinking) plus plain answer
+// content (ResponseTypeAnswer) must route the reasoning to thought events and
+// the answer live to final-answer events — never the reverse.
+func TestStreamThinkingToEventBus_RoutesReasoningAndAnswerSeparately(t *testing.T) {
+	mock := &mockChat{
+		responses: []mockResponse{
+			{chunks: []types.StreamResponse{
+				{ResponseType: types.ResponseTypeThinking, Content: "let me reason"},
+				{ResponseType: types.ResponseTypeThinking, Content: "", Done: true},
+				{ResponseType: types.ResponseTypeAnswer, Content: "The answer "},
+				{ResponseType: types.ResponseTypeAnswer, Content: "is 42.", Done: true, FinishReason: "stop"},
+			}},
+		},
+	}
+
+	engine := newTestEngine(t, mock)
+	var thoughts, answers string
+	engine.eventBus.On(event.EventAgentThought, func(_ context.Context, evt event.Event) error {
+		if d, ok := evt.Data.(event.AgentThoughtData); ok {
+			thoughts += d.Content
+		}
+		return nil
+	})
+	engine.eventBus.On(event.EventAgentFinalAnswer, func(_ context.Context, evt event.Event) error {
+		if d, ok := evt.Data.(event.AgentFinalAnswerData); ok {
+			answers += d.Content
+		}
+		return nil
+	})
+
+	resp, err := engine.streamThinkingToEventBus(context.Background(),
+		emptyMessages(), emptyTools(), 0, "sess-1")
+	require.NoError(t, err)
+
+	assert.Equal(t, "let me reason", thoughts, "reasoning_content must stream to thought events")
+	assert.Equal(t, "The answer is 42.", answers, "plain answer content must stream live to final-answer events")
+	assert.True(t, resp.AnswerStreamed, "AnswerStreamed must be set when answer text was streamed live")
+	assert.NotEmpty(t, resp.AnswerEventID, "AnswerEventID must identify the live answer stream")
+}
+
+// TestStreamThinkingToEventBus_SplitsInlineThinkBlock verifies that models which
+// embed reasoning inline as <think>…</think> in the content channel still have
+// their reasoning routed to thought events and only the real answer streamed to
+// the final-answer area.
+func TestStreamThinkingToEventBus_SplitsInlineThinkBlock(t *testing.T) {
+	mock := &mockChat{
+		responses: []mockResponse{
+			{chunks: []types.StreamResponse{
+				{ResponseType: types.ResponseTypeAnswer, Content: "<think>hidden reasoning</think>Visible answer.",
+					Done: true, FinishReason: "stop"},
+			}},
+		},
+	}
+
+	engine := newTestEngine(t, mock)
+	var thoughts, answers string
+	engine.eventBus.On(event.EventAgentThought, func(_ context.Context, evt event.Event) error {
+		if d, ok := evt.Data.(event.AgentThoughtData); ok {
+			thoughts += d.Content
+		}
+		return nil
+	})
+	engine.eventBus.On(event.EventAgentFinalAnswer, func(_ context.Context, evt event.Event) error {
+		if d, ok := evt.Data.(event.AgentFinalAnswerData); ok {
+			answers += d.Content
+		}
+		return nil
+	})
+
+	_, err := engine.streamThinkingToEventBus(context.Background(),
+		emptyMessages(), emptyTools(), 0, "sess-1")
+	require.NoError(t, err)
+
+	assert.Equal(t, "hidden reasoning", thoughts, "inline <think> content must route to thought events")
+	assert.Equal(t, "Visible answer.", answers, "answer outside <think> must stream to final-answer events")
+}
+
+// TestExecuteLoop_NaturalStop_DoesNotDuplicateAnswer ensures the natural-stop
+// branch does not re-emit the full answer (it was already streamed live), so
+// the final-answer content appears exactly once instead of streaming under
+// Thinking and then "jumping" to a duplicate answer block.
+func TestExecuteLoop_NaturalStop_DoesNotDuplicateAnswer(t *testing.T) {
+	mock := &mockChat{
+		responses: []mockResponse{
+			{chunks: []types.StreamResponse{
+				{ResponseType: types.ResponseTypeAnswer, Content: "Hello "},
+				{ResponseType: types.ResponseTypeAnswer, Content: "world", Done: true, FinishReason: "stop"},
+			}},
+		},
+	}
+
+	engine := newTestEngine(t, mock)
+	var answerContent string
+	var doneCount int
+	engine.eventBus.On(event.EventAgentFinalAnswer, func(_ context.Context, evt event.Event) error {
+		if d, ok := evt.Data.(event.AgentFinalAnswerData); ok {
+			answerContent += d.Content
+			if d.Done {
+				doneCount++
+			}
+		}
+		return nil
+	})
+
+	state := &types.AgentState{}
+	_, err := engine.executeLoop(context.Background(), state, "test query",
+		emptyMessages(), emptyTools(), "sess-1", "msg-1")
+	require.NoError(t, err)
+
+	assert.True(t, state.IsComplete)
+	assert.Equal(t, "Hello world", state.FinalAnswer)
+	assert.Equal(t, "Hello world", answerContent,
+		"answer content must be emitted exactly once (streamed live, not re-emitted by the natural-stop branch)")
+	assert.GreaterOrEqual(t, doneCount, 1, "a Done marker must close the answer stream")
+}
+
 func TestStreamFinalAnswerToEventBus_EmitsDoneWhenProviderEndsWithEmptyChunk(t *testing.T) {
 	mock := &mockChat{
 		responses: []mockResponse{
@@ -132,18 +132,33 @@ func (e *AgentEngine) analyzeResponse(
 			"answer_len": len(response.Content),
 		})

-		// Emit answer as final answer event (thinking events were already streamed)
-		answerID := generateEventID("answer")
-		if response.Content != "" {
-			e.eventBus.Emit(ctx, event.Event{
-				ID:        answerID,
-				Type:      event.EventAgentFinalAnswer,
-				SessionID: sessionID,
-				Data: event.AgentFinalAnswerData{
-					Content: response.Content,
-					Done:    false,
-				},
-			})
+		// Emit the final answer. The answer text reaches the UI by one of two
+		// paths:
+		//   (a) Already streamed live during the think phase — the common case
+		//       now that plain assistant content is routed straight to
+		//       EventAgentFinalAnswer (response.AnswerStreamed). Re-emitting the
+		//       full content here would render it twice and produce the
+		//       end-of-stream "jump from Thinking to Answer" the user reported,
+		//       so we only close the existing stream with a Done marker on the
+		//       same event ID.
+		//   (b) Not streamed live (e.g. the content only surfaced in the
+		//       accumulated result) — emit the full content, then Done.
+		var answerID string
+		if response.AnswerStreamed && response.AnswerEventID != "" {
+			answerID = response.AnswerEventID
+		} else {
+			answerID = generateEventID("answer")
+			if response.Content != "" {
+				e.eventBus.Emit(ctx, event.Event{
+					ID:        answerID,
+					Type:      event.EventAgentFinalAnswer,
+					SessionID: sessionID,
+					Data: event.AgentFinalAnswerData{
+						Content: response.Content,
+						Done:    false,
+					},
+				})
+			}
 		}
 		e.eventBus.Emit(ctx, event.Event{
 			ID:        answerID,
@@ -138,6 +138,59 @@ func (e *AgentEngine) streamThinkingToEventBus(
 	thinkingID := generateEventID("thinking")
 	answerID := generateEventID("answer")

+	// Routing state shared across chunk callbacks:
+	//   - splitter separates inline <think>…</think> reasoning from answer text
+	//     in the plain `content` channel (models that don't use reasoning_content).
+	//   - thinkingOpen tracks whether the thought stream still needs a Done marker.
+	//   - answerStreamed records that user-facing answer text was sent live to
+	//     the final-answer area, so the natural-stop branch only emits Done.
+	splitter := agenttools.NewThinkStreamSplitter()
+	thinkingOpen := false
+	answerStreamed := false
+
+	emitThought := func(content string, done bool) {
+		if content == "" && !done {
+			return
+		}
+		emittedEventTypes["thought_chunk"]++
+		e.eventBus.Emit(ctx, event.Event{
+			ID:        thinkingID,
+			Type:      event.EventAgentThought,
+			SessionID: sessionID,
+			Data: event.AgentThoughtData{
+				Content:   content,
+				Iteration: iteration,
+				Done:      done,
+			},
+		})
+	}
+	// closeThinking emits the thought Done marker once, used right before the
+	// first answer chunk so the UI flips the thinking card to "completed"
+	// instead of leaving it spinning while the answer streams.
+	closeThinking := func() {
+		if thinkingOpen {
+			emitThought("", true)
+			thinkingOpen = false
+		}
+	}
+	emitAnswer := func(content string) {
+		if content == "" {
+			return
+		}
+		closeThinking()
+		answerStreamed = true
+		emittedEventTypes["final_answer_chunk"]++
+		e.eventBus.Emit(ctx, event.Event{
+			ID:        answerID,
+			Type:      event.EventAgentFinalAnswer,
+			SessionID: sessionID,
+			Data: event.AgentFinalAnswerData{
+				Content: content,
+				Done:    false,
+			},
+		})
+	}
+
 	llmResult, err := e.streamLLMToEventBus(
 		ctx,
 		messages,
@@ -166,16 +219,7 @@ func (e *AgentEngine) streamThinkingToEventBus(
 			// Handle final_answer tool's streaming answer content
 			if chunk.ResponseType == types.ResponseTypeAnswer {
 				if source, _ := chunk.Data["source"].(string); source == "final_answer_tool" {
-					emittedEventTypes["final_answer_chunk"]++
-					e.eventBus.Emit(ctx, event.Event{
-						ID:        answerID,
-						Type:      event.EventAgentFinalAnswer,
-						SessionID: sessionID,
-						Data: event.AgentFinalAnswerData{
-							Content: chunk.Content,
-							Done:    false,
-						},
-					})
+					emitAnswer(chunk.Content)
 					return
 				}
 			}
@@ -204,18 +248,40 @@ func (e *AgentEngine) streamThinkingToEventBus(
 				}
 			}

+			// reasoning_content (separate thinking channel, e.g. DeepSeek V4) →
+			// thought area. Forward the Done marker the provider sends when it
+			// transitions from reasoning to answer.
+			if chunk.ResponseType == types.ResponseTypeThinking {
+				if chunk.Content != "" {
+					thinkingOpen = true
+					emitThought(chunk.Content, false)
+				} else if chunk.Done && thinkingOpen {
+					closeThinking()
+				}
+				return
+			}
+
+			// Plain content channel. This is the model's user-facing answer when
+			// it stops naturally (without the final_answer tool). Split out any
+			// inline <think> reasoning so the thought goes to the thought area
+			// and the genuine answer streams live into the final-answer area —
+			// fixing the "answer first shows under Thinking, then jumps" UX.
 			if chunk.Content != "" {
-				emittedEventTypes["thought_chunk"]++
-				e.eventBus.Emit(ctx, event.Event{
-					ID:        thinkingID, // Same ID for all chunks in this stream
-					Type:      event.EventAgentThought,
-					SessionID: sessionID,
-					Data: event.AgentThoughtData{
-						Content:   chunk.Content,
-						Iteration: iteration,
-						Done:      chunk.Done,
-					},
-				})
+				thinkPart, answerPart := splitter.Feed(chunk.Content)
+				if thinkPart != "" {
+					thinkingOpen = true
+					emitThought(thinkPart, false)
+				}
+				emitAnswer(answerPart)
+			}
+			if chunk.Done {
+				thinkPart, answerPart := splitter.Flush()
+				if thinkPart != "" {
+					thinkingOpen = true
+					emitThought(thinkPart, false)
+				}
+				emitAnswer(answerPart)
+				closeThinking()
 			}
 		},
 	)
@@ -243,6 +309,10 @@ func (e *AgentEngine) streamThinkingToEventBus(
 		ReasoningContent: llmResult.ReasoningContent,
 		ToolCalls:        llmResult.ToolCalls,
 		FinishReason:     finishReason,
+		AnswerStreamed:   answerStreamed,
+	}
+	if answerStreamed {
+		resp.AnswerEventID = answerID
 	}
 	if llmResult.Usage != nil {
 		resp.Usage = *llmResult.Usage
@@ -0,0 +1,106 @@
+package tools
+
+import "strings"
+
+// thinkOpenTag / thinkCloseTag are the inline reasoning markers some models
+// (DeepSeek, Qwen, …) embed directly in their `content` field instead of the
+// separate `reasoning_content` channel.
+const (
+	thinkOpenTag  = "<think>"
+	thinkCloseTag = "</think>"
+)
+
+// ThinkStreamSplitter incrementally separates inline <think>…</think> reasoning
+// from user-facing answer text as content arrives chunk-by-chunk. It is the
+// streaming counterpart to StripThinkBlocks: where StripThinkBlocks operates on
+// a fully-accumulated string, this splitter routes each chunk live so the
+// thinking portion can stream into the "thought" UI area and the answer portion
+// into the "final answer" area without waiting for the whole response.
+//
+// Tag boundaries that straddle two chunks (e.g. "<thi" at the end of one chunk
+// and "nk>" at the start of the next) are handled by buffering the trailing
+// bytes that could still become a tag prefix until the next Feed call. Call
+// Flush at end-of-stream to drain any buffered remainder.
+//
+// The splitter is NOT safe for concurrent use; create one per stream.
+type ThinkStreamSplitter struct {
+	inThink bool
+	// pending holds trailing bytes from a previous Feed that may be the start
+	// of a (possibly split) <think>/</think> tag and therefore cannot yet be
+	// classified as think or answer text.
+	pending string
+}
+
+// NewThinkStreamSplitter returns a splitter ready to receive content chunks.
+func NewThinkStreamSplitter() *ThinkStreamSplitter {
+	return &ThinkStreamSplitter{}
+}
+
+// Feed consumes one content chunk and returns the portions that are now
+// unambiguously thinking text and answer text respectively. Either return value
+// may be empty. Bytes that could still be part of a tag spanning into the next
+// chunk are buffered internally and surface on a later Feed or on Flush.
+func (sp *ThinkStreamSplitter) Feed(s string) (thinkOut, answerOut string) {
+	if s == "" {
+		return "", ""
+	}
+	sp.pending += s
+
+	var think, answer strings.Builder
+	for {
+		if sp.inThink {
+			if idx := strings.Index(sp.pending, thinkCloseTag); idx >= 0 {
+				think.WriteString(sp.pending[:idx])
+				sp.pending = sp.pending[idx+len(thinkCloseTag):]
+				sp.inThink = false
+				continue
+			}
+			safe, hold := holdBackPartialTag(sp.pending, thinkCloseTag)
+			think.WriteString(safe)
+			sp.pending = hold
+			return think.String(), answer.String()
+		}
+
+		if idx := strings.Index(sp.pending, thinkOpenTag); idx >= 0 {
+			answer.WriteString(sp.pending[:idx])
+			sp.pending = sp.pending[idx+len(thinkOpenTag):]
+			sp.inThink = true
+			continue
+		}
+		safe, hold := holdBackPartialTag(sp.pending, thinkOpenTag)
+		answer.WriteString(safe)
+		sp.pending = hold
+		return think.String(), answer.String()
+	}
+}
+
+// Flush drains any buffered remainder at end-of-stream. An unterminated <think>
+// block is treated as thinking text; anything else is answer text.
+func (sp *ThinkStreamSplitter) Flush() (thinkOut, answerOut string) {
+	rest := sp.pending
+	sp.pending = ""
+	if rest == "" {
+		return "", ""
+	}
+	if sp.inThink {
+		return rest, ""
+	}
+	return "", rest
+}
+
+// holdBackPartialTag splits s into the part that is safe to emit now and a
+// trailing suffix that is a proper prefix of tag (and so might complete into a
+// real tag on the next chunk). When s ends with no such prefix, the whole
+// string is safe and hold is empty.
+func holdBackPartialTag(s, tag string) (safe, hold string) {
+	maxK := len(tag) - 1
+	if maxK > len(s) {
+		maxK = len(s)
+	}
+	for k := maxK; k >= 1; k-- {
+		if strings.HasSuffix(s, tag[:k]) {
+			return s[:len(s)-k], s[len(s)-k:]
+		}
+	}
+	return s, ""
+}
@@ -0,0 +1,101 @@
+package tools
+
+import "testing"
+
+// feedAll runs the splitter over a sequence of chunks and concatenates the
+// think / answer outputs (including the final Flush), mirroring how the agent
+// stream router consumes it.
+func feedAll(sp *ThinkStreamSplitter, chunks []string) (think, answer string) {
+	for _, c := range chunks {
+		tk, ans := sp.Feed(c)
+		think += tk
+		answer += ans
+	}
+	tk, ans := sp.Flush()
+	think += tk
+	answer += ans
+	return think, answer
+}
+
+func TestThinkStreamSplitter(t *testing.T) {
+	tests := []struct {
+		name       string
+		chunks     []string
+		wantThink  string
+		wantAnswer string
+	}{
+		{
+			name:       "no think tags - all answer",
+			chunks:     []string{"Hello ", "world"},
+			wantThink:  "",
+			wantAnswer: "Hello world",
+		},
+		{
+			name:       "single think block then answer",
+			chunks:     []string{"<think>reasoning</think>The answer is 42."},
+			wantThink:  "reasoning",
+			wantAnswer: "The answer is 42.",
+		},
+		{
+			name:       "open tag split across chunks",
+			chunks:     []string{"<thi", "nk>secret</think>visible"},
+			wantThink:  "secret",
+			wantAnswer: "visible",
+		},
+		{
+			name:       "close tag split across chunks",
+			chunks:     []string{"<think>think part</thi", "nk>answer part"},
+			wantThink:  "think part",
+			wantAnswer: "answer part",
+		},
+		{
+			name:       "think content streamed across multiple chunks",
+			chunks:     []string{"<think>a", "b", "c</think>", "done"},
+			wantThink:  "abc",
+			wantAnswer: "done",
+		},
+		{
+			name:       "answer before think block",
+			chunks:     []string{"prefix <think>mid</think> suffix"},
+			wantThink:  "mid",
+			wantAnswer: "prefix  suffix",
+		},
+		{
+			name:       "two think blocks",
+			chunks:     []string{"<think>one</think>A<think>two</think>B"},
+			wantThink:  "onetwo",
+			wantAnswer: "AB",
+		},
+		{
+			name:       "unterminated think treated as think on flush",
+			chunks:     []string{"<think>still thinking"},
+			wantThink:  "still thinking",
+			wantAnswer: "",
+		},
+		{
+			name:       "literal less-than in answer is preserved",
+			chunks:     []string{"if a < ", "b then"},
+			wantThink:  "",
+			wantAnswer: "if a < b then",
+		},
+		{
+			name:       "empty chunks are no-ops",
+			chunks:     []string{"", "answer", ""},
+			wantThink:  "",
+			wantAnswer: "answer",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sp := NewThinkStreamSplitter()
+			think, answer := feedAll(sp, tt.chunks)
+			if think != tt.wantThink {
+				t.Errorf("think = %q, want %q", think, tt.wantThink)
+			}
+			if answer != tt.wantAnswer {
+				t.Errorf("answer = %q, want %q", answer, tt.wantAnswer)
+			}
+		})
+	}
+}
@@ -54,6 +54,19 @@ type ChatResponse struct {
 	ToolCalls        []LLMToolCall `json:"tool_calls,omitempty"`
 	FinishReason     string        `json:"finish_reason,omitempty"`
 	Usage            TokenUsage    `json:"usage"`
+
+	// AnswerStreamed reports whether the user-facing answer text was already
+	// streamed live to the final-answer UI area during this round (i.e. the
+	// model answered with plain content rather than via the final_answer
+	// tool). When true, the natural-stop branch must only emit the closing
+	// Done marker for AnswerEventID instead of re-emitting the whole answer —
+	// otherwise the answer would render twice and "jump" at end of stream.
+	// Transient, never persisted.
+	AnswerStreamed bool `json:"-"`
+	// AnswerEventID is the EventBus event ID under which the live answer
+	// chunks were streamed, so the natural-stop branch can close the same
+	// stream with a Done marker. Empty when AnswerStreamed is false.
+	AnswerEventID string `json:"-"`
 }

 // Response type