From 1d4c9b60a3f7d7f33aa6b3ef672e73b43f5c20de Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Sun, 22 Feb 2026 16:34:42 +0900 Subject: [PATCH] feat: add real-time streaming preview of LLM responses Show LLM output progressively in the chat placeholder instead of waiting for the full response. Uses throttled (500ms) IsStatus messages routed through the existing EditStatus() path, so channels that support placeholder editing (Telegram, etc.) get live updates while unsupported channels are unaffected. Co-Authored-By: Claude Opus 4.6 --- pkg/agent/loop.go | 30 +++++++++++++- pkg/agent/loop_test.go | 94 ++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 119 insertions(+), 5 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 88d229ac1..aa0d5d853 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -1386,6 +1386,7 @@ func consumeStreamWithRepetitionDetection( ch <-chan protocoltypes.StreamEvent, cancelFn context.CancelFunc, checkInterval int, + onChunk func(accumulated string), ) (*providers.LLMResponse, bool, error) { var content strings.Builder var toolCalls []streamToolCallAcc @@ -1400,6 +1401,9 @@ func consumeStreamWithRepetitionDetection( if ev.ContentDelta != "" { content.WriteString(ev.ContentDelta) runesSinceLastCheck += utf8.RuneCountInString(ev.ContentDelta) + if onChunk != nil { + onChunk(content.String()) + } } if ev.FinishReason != "" { finishReason = ev.FinishReason @@ -1551,6 +1555,30 @@ func (al *AgentLoop) runLLMIteration( var response *providers.LLMResponse var err error + // Build onChunk callback for streaming preview. + // When sending responses to a real (non-internal) channel, publish + // throttled status updates so the user sees LLM output in real time. + var onChunk func(string) + if !constants.IsInternalChannel(opts.Channel) { + lastPublish := time.Time{} + onChunk = func(accumulated string) { + if time.Since(lastPublish) < 500*time.Millisecond { + return + } + lastPublish = time.Now() + display := utils.StripThinkBlocks(accumulated) + if strings.TrimSpace(display) == "" { + return + } + al.bus.PublishOutbound(bus.OutboundMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: display + " \u2589", + IsStatus: true, + }) + } + } + // doCall invokes a single LLM provider, using streaming with // early repetition detection when the provider supports it. opts_ := map[string]any{ @@ -1565,7 +1593,7 @@ func (al *AgentLoop) runLLMIteration( if err != nil { return nil, err } - resp, repetition, err := consumeStreamWithRepetitionDetection(ch, streamCancel, 1000) + resp, repetition, err := consumeStreamWithRepetitionDetection(ch, streamCancel, 1000, onChunk) if err != nil { return nil, err } diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 567cd0aea..0e42b98ad 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -2135,7 +2135,7 @@ func TestConsumeStream_NormalCompletion(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() - resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000) + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -2181,7 +2181,7 @@ func TestConsumeStream_DetectsRepetition(t *testing.T) { close(ch) }() - resp, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000) + resp, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000, nil) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -2219,7 +2219,7 @@ func TestConsumeStream_ToolCallAccumulation(t *testing.T) { _, cancel := context.WithCancel(context.Background()) defer cancel() - resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000) + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) if err != nil { t.Fatalf("unexpected error: %v", err) } @@ -2248,7 +2248,7 @@ func TestConsumeStream_StreamError(t *testing.T) { _, cancel := context.WithCancel(context.Background()) defer cancel() - _, _, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000) + _, _, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, nil) if err == nil { t.Fatal("expected error, got nil") } @@ -2256,3 +2256,89 @@ func TestConsumeStream_StreamError(t *testing.T) { t.Errorf("error = %q, want to contain %q", err.Error(), "read error") } } + +func TestConsumeStream_OnChunkCallback(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 8) + go func() { + ch <- protocoltypes.StreamEvent{ContentDelta: "Hello "} + ch <- protocoltypes.StreamEvent{ContentDelta: "world"} + ch <- protocoltypes.StreamEvent{ContentDelta: "!"} + ch <- protocoltypes.StreamEvent{FinishReason: "stop"} + close(ch) + }() + + _, cancel := context.WithCancel(context.Background()) + defer cancel() + + var chunks []string + onChunk := func(accumulated string) { + chunks = append(chunks, accumulated) + } + + resp, detected, err := consumeStreamWithRepetitionDetection(ch, cancel, 1000, onChunk) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if detected { + t.Fatal("expected detected=false") + } + if resp.Content != "Hello world!" { + t.Errorf("Content = %q, want %q", resp.Content, "Hello world!") + } + // onChunk should be called once per content delta (3 times) + if len(chunks) != 3 { + t.Fatalf("onChunk called %d times, want 3", len(chunks)) + } + if chunks[0] != "Hello " { + t.Errorf("chunks[0] = %q, want %q", chunks[0], "Hello ") + } + if chunks[1] != "Hello world" { + t.Errorf("chunks[1] = %q, want %q", chunks[1], "Hello world") + } + if chunks[2] != "Hello world!" { + t.Errorf("chunks[2] = %q, want %q", chunks[2], "Hello world!") + } +} + +func TestConsumeStream_OnChunkWithRepetitionDetection(t *testing.T) { + ch := make(chan protocoltypes.StreamEvent, 64) + cancelCalled := false + + ctx, cancel := context.WithCancel(context.Background()) + wrappedCancel := func() { + cancelCalled = true + cancel() + } + + repeatedChunk := strings.Repeat("abcdefghij", 50) // 500 chars per chunk + go func() { + for i := 0; i < 6; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: repeatedChunk} + } + for i := 0; i < 10; i++ { + ch <- protocoltypes.StreamEvent{ContentDelta: "more data"} + } + close(ch) + }() + + var chunkCount int + onChunk := func(accumulated string) { + chunkCount++ + } + + _, detected, err := consumeStreamWithRepetitionDetection(ch, wrappedCancel, 1000, onChunk) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !detected { + t.Fatal("expected repetition detection to trigger") + } + if !cancelCalled { + t.Error("expected cancelFn to be called") + } + // onChunk should have been called at least once before detection + if chunkCount == 0 { + t.Error("expected onChunk to be called at least once") + } + _ = ctx +}