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