From 876fd0276f16d0f9cdccedb7d87993a7d0d97084 Mon Sep 17 00:00:00 2001 From: wizardchen Date: Wed, 3 Jun 2026 22:14:35 +0800 Subject: [PATCH] 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. --- internal/agent/engine_test.go | 119 ++++++++++++++++++++++ internal/agent/observe.go | 39 ++++--- internal/agent/think.go | 112 ++++++++++++++++---- internal/agent/tools/think_stream.go | 106 +++++++++++++++++++ internal/agent/tools/think_stream_test.go | 101 ++++++++++++++++++ internal/types/chat.go | 13 +++ 6 files changed, 457 insertions(+), 33 deletions(-) create mode 100644 internal/agent/tools/think_stream.go create mode 100644 internal/agent/tools/think_stream_test.go diff --git a/internal/agent/engine_test.go b/internal/agent/engine_test.go index 8871a55b..0b3da547 100644 --- a/internal/agent/engine_test.go +++ b/internal/agent/engine_test.go @@ -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 … 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: "hidden reasoningVisible 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 content must route to thought events") + assert.Equal(t, "Visible answer.", answers, "answer outside 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{ diff --git a/internal/agent/observe.go b/internal/agent/observe.go index 64ce4211..d0b41178 100644 --- a/internal/agent/observe.go +++ b/internal/agent/observe.go @@ -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, diff --git a/internal/agent/think.go b/internal/agent/think.go index 67454500..f82446e0 100644 --- a/internal/agent/think.go +++ b/internal/agent/think.go @@ -138,6 +138,59 @@ func (e *AgentEngine) streamThinkingToEventBus( thinkingID := generateEventID("thinking") answerID := generateEventID("answer") + // Routing state shared across chunk callbacks: + // - splitter separates inline … 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 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 diff --git a/internal/agent/tools/think_stream.go b/internal/agent/tools/think_stream.go new file mode 100644 index 00000000..934935fe --- /dev/null +++ b/internal/agent/tools/think_stream.go @@ -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 = "" + thinkCloseTag = "" +) + +// ThinkStreamSplitter incrementally separates inline … 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. "" 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) / 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 +// 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, "" +} diff --git a/internal/agent/tools/think_stream_test.go b/internal/agent/tools/think_stream_test.go new file mode 100644 index 00000000..a61ffae7 --- /dev/null +++ b/internal/agent/tools/think_stream_test.go @@ -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{"reasoningThe answer is 42."}, + wantThink: "reasoning", + wantAnswer: "The answer is 42.", + }, + { + name: "open tag split across chunks", + chunks: []string{"secretvisible"}, + wantThink: "secret", + wantAnswer: "visible", + }, + { + name: "close tag split across chunks", + chunks: []string{"think partanswer part"}, + wantThink: "think part", + wantAnswer: "answer part", + }, + { + name: "think content streamed across multiple chunks", + chunks: []string{"a", "b", "c", "done"}, + wantThink: "abc", + wantAnswer: "done", + }, + { + name: "answer before think block", + chunks: []string{"prefix mid suffix"}, + wantThink: "mid", + wantAnswer: "prefix suffix", + }, + { + name: "two think blocks", + chunks: []string{"oneAtwoB"}, + wantThink: "onetwo", + wantAnswer: "AB", + }, + { + name: "unterminated think treated as think on flush", + chunks: []string{"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) + } + }) + } +} diff --git a/internal/types/chat.go b/internal/types/chat.go index e13166ea..5c539277 100644 --- a/internal/types/chat.go +++ b/internal/types/chat.go @@ -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