From 15888963ea7c1b1c8715b9c49bd9c9934a70536c Mon Sep 17 00:00:00 2001 From: alanbulan <1308048843qq@gmail.com> Date: Tue, 7 Apr 2026 22:40:10 +0800 Subject: [PATCH] feat: stream provider output to compatible channels Prefer StreamingProvider responses in the agent loop when the active channel can consume incremental updates. Add Pico channel streaming primitives and message.delete handling so streamed turns can render progressively in the web UI while still preserving final LLM responses for normal turn completion. --- pkg/agent/loop.go | 59 +++++++++++++- pkg/agent/loop_test.go | 71 +++++++++++++++++ pkg/agent/mock_provider_test.go | 31 +++++++- pkg/channels/pico/pico.go | 72 +++++++++++++++-- pkg/channels/pico/pico_test.go | 83 ++++++++++++++++++++ pkg/channels/pico/protocol.go | 1 + pkg/providers/openai_compat/provider_test.go | 30 +++++++ web/frontend/src/features/chat/protocol.ts | 12 +++ 8 files changed, 350 insertions(+), 9 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 369928d78..3c9edf094 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -2004,6 +2004,61 @@ turnLoop: "tools_json": formatToolsForLog(providerToolDefs), }) + callProvider := func( + ctx context.Context, + provider providers.LLMProvider, + model string, + messagesForCall []providers.Message, + toolDefsForCall []providers.ToolDefinition, + ) (*providers.LLMResponse, error) { + if streamingProvider, ok := provider.(providers.StreamingProvider); ok { + if streamer, streamOK := al.bus.GetStreamer(ctx, ts.channel, ts.chatID); streamOK { + var finalized bool + response, err := streamingProvider.ChatStream( + ctx, + messagesForCall, + toolDefsForCall, + model, + llmOpts, + func(accumulated string) { + if accumulated == "" { + return + } + if err := streamer.Update(ctx, accumulated); err != nil { + logger.DebugCF("agent", "Streaming update failed", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "chat_id": ts.chatID, + "error": err.Error(), + }) + } + }, + ) + if err != nil { + streamer.Cancel(ctx) + return nil, err + } + if response != nil && response.Content != "" { + if err := streamer.Finalize(ctx, response.Content); err != nil { + logger.DebugCF("agent", "Streaming finalize failed", map[string]any{ + "agent_id": ts.agent.ID, + "channel": ts.channel, + "chat_id": ts.chatID, + "error": err.Error(), + }) + } else { + finalized = true + } + } + if !finalized && response != nil { + streamer.Cancel(ctx) + } + return response, nil + } + } + return provider.Chat(ctx, messagesForCall, toolDefsForCall, model, llmOpts) + } + callLLM := func(messagesForCall []providers.Message, toolDefsForCall []providers.ToolDefinition) (*providers.LLMResponse, error) { providerCtx, providerCancel := context.WithCancel(turnCtx) ts.setProviderCancel(providerCancel) @@ -2024,7 +2079,7 @@ turnLoop: if cp, ok := ts.agent.CandidateProviders[providers.ModelKey(provider, model)]; ok { candidateProvider = cp } - return candidateProvider.Chat(ctx, messagesForCall, toolDefsForCall, model, llmOpts) + return callProvider(ctx, candidateProvider, model, messagesForCall, toolDefsForCall) }, ) if fbErr != nil { @@ -2040,7 +2095,7 @@ turnLoop: } return fbResult.Response, nil } - return activeProvider.Chat(providerCtx, messagesForCall, toolDefsForCall, llmModel, llmOpts) + return callProvider(providerCtx, activeProvider, llmModel, messagesForCall, toolDefsForCall) } var response *providers.LLMResponse diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 3d04b81cc..0a21ed014 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -41,6 +41,38 @@ type fakeMediaChannel struct { sentMedia []bus.OutboundMediaMessage } +type fakeStreamer struct { + updates []string + finalized []string + cancelled bool +} + +func (s *fakeStreamer) Update(ctx context.Context, content string) error { + s.updates = append(s.updates, content) + return nil +} + +func (s *fakeStreamer) Finalize(ctx context.Context, content string) error { + s.finalized = append(s.finalized, content) + return nil +} + +func (s *fakeStreamer) Cancel(ctx context.Context) { + s.cancelled = true +} + +type fakeStreamingChannel struct { + fakeChannel + streamer *fakeStreamer +} + +func (f *fakeStreamingChannel) BeginStream(ctx context.Context, chatID string) (channels.Streamer, error) { + if f.streamer == nil { + f.streamer = &fakeStreamer{} + } + return f.streamer, nil +} + func (f *fakeMediaChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) ([]string, error) { f.sentMedia = append(f.sentMedia, msg) return nil, nil @@ -2626,6 +2658,45 @@ func TestProcessMessage_PublishesReasoningContentToReasoningChannel(t *testing.T } } +func TestProcessMessage_StreamsToChannelWhenProviderSupportsStreaming(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + msgBus := bus.NewMessageBus() + provider := &mockProvider{streamChunks: []string{"hel", "lo"}} + al := NewAgentLoop(cfg, msgBus, provider) + streamingChannel := &fakeStreamingChannel{fakeChannel: fakeChannel{id: "reason-chat"}, streamer: &fakeStreamer{}} + al.SetChannelManager(newStartedTestChannelManager(t, msgBus, nil, "pico", streamingChannel)) + + response, err := al.processMessage(context.Background(), bus.InboundMessage{ + Channel: "pico", + SenderID: "user1", + ChatID: "chat1", + Content: "hello", + }) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + if response != "hello" { + t.Fatalf("processMessage() response = %q, want %q", response, "hello") + } + if got := streamingChannel.streamer.updates; len(got) != 2 || got[0] != "hel" || got[1] != "hello" { + t.Fatalf("stream updates = %v, want [hel hello]", got) + } + if got := streamingChannel.streamer.finalized; len(got) != 1 || got[0] != "hello" { + t.Fatalf("stream finalized = %v, want [hello]", got) + } +} + func TestProcessHeartbeat_DoesNotPublishToolFeedback(t *testing.T) { tmpDir := t.TempDir() heartbeatFile := filepath.Join(tmpDir, "heartbeat-task.txt") diff --git a/pkg/agent/mock_provider_test.go b/pkg/agent/mock_provider_test.go index 4962810dc..ed0c1855d 100644 --- a/pkg/agent/mock_provider_test.go +++ b/pkg/agent/mock_provider_test.go @@ -6,7 +6,9 @@ import ( "github.com/sipeed/picoclaw/pkg/providers" ) -type mockProvider struct{} +type mockProvider struct { + streamChunks []string +} func (m *mockProvider) Chat( ctx context.Context, @@ -24,3 +26,30 @@ func (m *mockProvider) Chat( func (m *mockProvider) GetDefaultModel() string { return "mock-model" } + +func (m *mockProvider) ChatStream( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + opts map[string]any, + onChunk func(accumulated string), +) (*providers.LLMResponse, error) { + accumulated := "" + for _, chunk := range m.streamChunks { + accumulated += chunk + if onChunk != nil { + onChunk(accumulated) + } + } + if accumulated == "" { + accumulated = "Mock response" + if onChunk != nil { + onChunk(accumulated) + } + } + return &providers.LLMResponse{ + Content: accumulated, + ToolCalls: []providers.ToolCall{}, + }, nil +} diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index e22da1ba1..4ad406a42 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -23,12 +23,13 @@ import ( // picoConn represents a single WebSocket connection. type picoConn struct { - id string - conn *websocket.Conn - sessionID string - writeMu sync.Mutex - closed atomic.Bool - cancel context.CancelFunc // cancels per-connection goroutines (e.g. pingLoop) + id string + conn *websocket.Conn + sessionID string + writeMu sync.Mutex + closed atomic.Bool + cancel context.CancelFunc // cancels per-connection goroutines (e.g. pingLoop) + writeJSONFunc func(any) error } var allowedInlineImageMIMETypes = map[string]struct{}{ @@ -46,6 +47,9 @@ func (pc *picoConn) writeJSON(v any) error { } pc.writeMu.Lock() defer pc.writeMu.Unlock() + if pc.writeJSONFunc != nil { + return pc.writeJSONFunc(v) + } return pc.conn.WriteJSON(v) } @@ -72,6 +76,14 @@ type PicoChannel struct { cancel context.CancelFunc } +type picoStreamer struct { + channel *PicoChannel + chatID string + messageID string + content string + mu sync.Mutex +} + // NewPicoChannel creates a new Pico Protocol channel. func NewPicoChannel(cfg config.PicoConfig, messageBus *bus.MessageBus) (*PicoChannel, error) { if cfg.Token.String() == "" { @@ -255,6 +267,28 @@ func (c *PicoChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri return nil, c.broadcastToSession(msg.ChatID, outMsg) } +// BeginStream implements channels.StreamingCapable. +func (c *PicoChannel) BeginStream(ctx context.Context, chatID string) (channels.Streamer, error) { + if !c.IsRunning() { + return nil, channels.ErrNotRunning + } + + messageID := uuid.New().String() + outMsg := newMessage(TypeMessageCreate, map[string]any{ + "content": "", + "message_id": messageID, + }) + if err := c.broadcastToSession(chatID, outMsg); err != nil { + return nil, err + } + + return &picoStreamer{ + channel: c, + chatID: chatID, + messageID: messageID, + }, nil +} + // EditMessage implements channels.MessageEditor. func (c *PicoChannel) EditMessage(ctx context.Context, chatID string, messageID string, content string) error { outMsg := newMessage(TypeMessageUpdate, map[string]any{ @@ -264,6 +298,14 @@ func (c *PicoChannel) EditMessage(ctx context.Context, chatID string, messageID return c.broadcastToSession(chatID, outMsg) } +// DeleteMessage implements channels.MessageDeleter. +func (c *PicoChannel) DeleteMessage(ctx context.Context, chatID string, messageID string) error { + outMsg := newMessage(TypeMessageDelete, map[string]any{ + "message_id": messageID, + }) + return c.broadcastToSession(chatID, outMsg) +} + // StartTyping implements channels.TypingCapable. func (c *PicoChannel) StartTyping(ctx context.Context, chatID string) (func(), error) { startMsg := newMessage(TypeTypingStart, nil) @@ -276,6 +318,24 @@ func (c *PicoChannel) StartTyping(ctx context.Context, chatID string) (func(), e }, nil } +func (s *picoStreamer) Update(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.content = content + return s.channel.EditMessage(ctx, s.chatID, s.messageID, content) +} + +func (s *picoStreamer) Finalize(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + s.content = content + return s.channel.EditMessage(ctx, s.chatID, s.messageID, content) +} + +func (s *picoStreamer) Cancel(ctx context.Context) { + _ = s.channel.DeleteMessage(ctx, s.chatID, s.messageID) +} + // SendPlaceholder implements channels.PlaceholderCapable. // It sends a placeholder message via the Pico Protocol that will later be // edited to the actual response via EditMessage (channels.MessageEditor). diff --git a/pkg/channels/pico/pico_test.go b/pkg/channels/pico/pico_test.go index e712767ad..8951b0977 100644 --- a/pkg/channels/pico/pico_test.go +++ b/pkg/channels/pico/pico_test.go @@ -23,6 +23,7 @@ func newTestPicoChannel(t *testing.T) *PicoChannel { } ch.ctx = context.Background() + ch.SetRunning(true) return ch } @@ -122,6 +123,88 @@ func TestBroadcastToSession_TargetsOnlyRequestedSession(t *testing.T) { } } +func TestDeleteMessage_SendsDeleteEvent(t *testing.T) { + ch := newTestPicoChannel(t) + + messages := make([]PicoMessage, 0, 1) + conn := &picoConn{ + id: "delete", + sessionID: "sess-del", + writeJSONFunc: func(v any) error { + msg, ok := v.(PicoMessage) + if ok { + messages = append(messages, msg) + } + return nil + }, + } + ch.addConnForTest(conn) + + if err := ch.DeleteMessage(context.Background(), "pico:sess-del", "msg-1"); err != nil { + t.Fatalf("DeleteMessage() error = %v", err) + } + if len(messages) != 1 { + t.Fatalf("expected 1 delete message, got %d", len(messages)) + } + if messages[0].Type != TypeMessageDelete { + t.Fatalf("message type = %q, want %q", messages[0].Type, TypeMessageDelete) + } + if messages[0].Payload["message_id"] != "msg-1" { + t.Fatalf("message_id = %v, want msg-1", messages[0].Payload["message_id"]) + } +} + +func TestBeginStream_SendsCreateAndUpdatesMessage(t *testing.T) { + ch := newTestPicoChannel(t) + + messages := make([]PicoMessage, 0, 3) + conn := &picoConn{ + id: "stream", + sessionID: "sess-1", + writeJSONFunc: func(v any) error { + msg, ok := v.(PicoMessage) + if ok { + messages = append(messages, msg) + } + return nil + }, + } + ch.addConnForTest(conn) + + streamer, err := ch.BeginStream(context.Background(), "pico:sess-1") + if err != nil { + t.Fatalf("BeginStream() error = %v", err) + } + if len(messages) == 0 { + t.Fatal("expected initial message.create for stream placeholder") + } + createPayload := messages[0].Payload + messageID, _ := createPayload["message_id"].(string) + if messageID == "" { + t.Fatal("expected message_id in initial stream create payload") + } + if err := streamer.Update(context.Background(), "hello"); err != nil { + t.Fatalf("Update() error = %v", err) + } + if err := streamer.Finalize(context.Background(), "hello world"); err != nil { + t.Fatalf("Finalize() error = %v", err) + } + if len(messages) < 3 { + t.Fatalf("expected at least 3 messages, got %d", len(messages)) + } + updatePayload := messages[1].Payload + if updatePayload["message_id"] != messageID { + t.Fatalf("update message_id = %v, want %s", updatePayload["message_id"], messageID) + } + if updatePayload["content"] != "hello" { + t.Fatalf("update content = %v, want hello", updatePayload["content"]) + } + finalPayload := messages[2].Payload + if finalPayload["content"] != "hello world" { + t.Fatalf("final content = %v, want hello world", finalPayload["content"]) + } +} + func (c *PicoChannel) addConnForTest(pc *picoConn) { c.connsMu.Lock() defer c.connsMu.Unlock() diff --git a/pkg/channels/pico/protocol.go b/pkg/channels/pico/protocol.go index 3f8ba8643..4d582582d 100644 --- a/pkg/channels/pico/protocol.go +++ b/pkg/channels/pico/protocol.go @@ -12,6 +12,7 @@ const ( // TypeMessageCreate is sent from server to client. TypeMessageCreate = "message.create" TypeMessageUpdate = "message.update" + TypeMessageDelete = "message.delete" TypeMediaCreate = "media.create" TypeTypingStart = "typing.start" TypeTypingStop = "typing.stop" diff --git a/pkg/providers/openai_compat/provider_test.go b/pkg/providers/openai_compat/provider_test.go index d140d63d6..046baf8c8 100644 --- a/pkg/providers/openai_compat/provider_test.go +++ b/pkg/providers/openai_compat/provider_test.go @@ -764,6 +764,36 @@ func TestProviderChat_CustomHeadersInjected(t *testing.T) { } } +func TestProviderChatStream_ParsesTextDeltas(t *testing.T) { + var chunks []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hel\"},\"finish_reason\":null}]}\n\n")) + _, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"lo\"},\"finish_reason\":\"stop\"}]}\n\n")) + _, _ = w.Write([]byte("data: [DONE]\n\n")) + })) + defer server.Close() + + p := NewProvider("key", server.URL, "") + out, err := p.ChatStream( + t.Context(), + []Message{{Role: "user", Content: "hi"}}, + nil, + "gpt-4o", + nil, + func(accumulated string) { chunks = append(chunks, accumulated) }, + ) + if err != nil { + t.Fatalf("ChatStream() error = %v", err) + } + if out.Content != "hello" { + t.Fatalf("Content = %q, want %q", out.Content, "hello") + } + if len(chunks) != 2 || chunks[0] != "hel" || chunks[1] != "hello" { + t.Fatalf("chunks = %v, want [hel hello]", chunks) + } +} + func TestProviderChatStream_CustomHeadersInjected(t *testing.T) { var gotSource, gotAuth, gotUserAgent string diff --git a/web/frontend/src/features/chat/protocol.ts b/web/frontend/src/features/chat/protocol.ts index 7429aef01..1e362f726 100644 --- a/web/frontend/src/features/chat/protocol.ts +++ b/web/frontend/src/features/chat/protocol.ts @@ -61,6 +61,18 @@ export function handlePicoMessage( break } + case "message.delete": { + const messageId = payload.message_id as string + if (!messageId) { + break + } + + updateChatStore((prev) => ({ + messages: prev.messages.filter((msg) => msg.id !== messageId), + })) + break + } + case "typing.start": updateChatStore({ isTyping: true }) break