diff --git a/pkg/agent/streaming_test.go b/pkg/agent/streaming_test.go new file mode 100644 index 000000000..4b4e8616b --- /dev/null +++ b/pkg/agent/streaming_test.go @@ -0,0 +1,244 @@ +package agent + +import ( + "context" + "os" + "testing" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" +) + +type streamingTestProvider struct { + chatCalled bool + chatStreamCalled bool + chunks []string + response *providers.LLMResponse +} + +func (p *streamingTestProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + p.chatCalled = true + if p.response != nil { + return p.response, nil + } + return &providers.LLMResponse{Content: "fallback"}, nil +} + +func (p *streamingTestProvider) ChatStream( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, + onChunk func(accumulated string), +) (*providers.LLMResponse, error) { + p.chatStreamCalled = true + for _, chunk := range p.chunks { + onChunk(chunk) + } + if p.response != nil { + return p.response, nil + } + return &providers.LLMResponse{Content: "streamed"}, nil +} + +func (p *streamingTestProvider) GetDefaultModel() string { + return "test-model" +} + +type recordingStreamer struct { + updates []string + finalized []string + canceled bool +} + +func (s *recordingStreamer) Update(ctx context.Context, content string) error { + s.updates = append(s.updates, content) + return nil +} + +func (s *recordingStreamer) Finalize(ctx context.Context, content string) error { + s.finalized = append(s.finalized, content) + return nil +} + +func (s *recordingStreamer) Cancel(ctx context.Context) { + s.canceled = true +} + +type recordingStreamDelegate struct { + streamer *recordingStreamer + ok bool +} + +func (d *recordingStreamDelegate) GetStreamer( + ctx context.Context, + channel, chatID string, +) (bus.Streamer, bool) { + if !d.ok { + return nil, false + } + return d.streamer, true +} + +func newStreamingTestLoop( + t *testing.T, + streamingEnabled bool, + provider *streamingTestProvider, + streamer *recordingStreamer, +) (*AgentLoop, func()) { + t.Helper() + + tmpDir, err := os.MkdirTemp("", "agent-streaming-test-*") + if err != nil { + t.Fatalf("MkdirTemp() error = %v", err) + } + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + StreamingEnabled: streamingEnabled, + }, + }, + } + + msgBus := bus.NewMessageBus() + msgBus.SetStreamDelegate(&recordingStreamDelegate{ + streamer: streamer, + ok: true, + }) + + return NewAgentLoop(cfg, msgBus, provider), func() { + os.RemoveAll(tmpDir) + } +} + +func TestProcessMessage_UsesStreamingProviderWhenEnabled(t *testing.T) { + provider := &streamingTestProvider{ + chunks: []string{"Hel", "Hello"}, + response: &providers.LLMResponse{ + Content: "Hello", + }, + } + streamer := &recordingStreamer{} + al, cleanup := newStreamingTestLoop(t, true, provider, streamer) + defer cleanup() + + response, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "pico", + ChatID: "pico:test-session", + SenderID: "user-1", + Content: "hello", + })) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + + if response != "Hello" { + t.Fatalf("response = %q, want %q", response, "Hello") + } + if !provider.chatStreamCalled { + t.Fatal("expected ChatStream to be used") + } + if provider.chatCalled { + t.Fatal("expected Chat() not to be used when streaming is enabled") + } + if len(streamer.updates) != 2 { + t.Fatalf("stream updates = %#v, want 2 accumulated updates", streamer.updates) + } + if len(streamer.finalized) != 1 || streamer.finalized[0] != "Hello" { + t.Fatalf("finalized = %#v, want final Hello", streamer.finalized) + } + if streamer.canceled { + t.Fatal("streamer should not be canceled on direct response") + } +} + +func TestProcessMessage_SkipsStreamingProviderWhenDisabled(t *testing.T) { + provider := &streamingTestProvider{ + response: &providers.LLMResponse{ + Content: "No stream", + }, + } + streamer := &recordingStreamer{} + al, cleanup := newStreamingTestLoop(t, false, provider, streamer) + defer cleanup() + + response, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "pico", + ChatID: "pico:test-session", + SenderID: "user-1", + Content: "hello", + })) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + + if response != "No stream" { + t.Fatalf("response = %q, want %q", response, "No stream") + } + if provider.chatStreamCalled { + t.Fatal("expected ChatStream not to be used when streaming is disabled") + } + if !provider.chatCalled { + t.Fatal("expected Chat() to be used when streaming is disabled") + } + if len(streamer.updates) != 0 || len(streamer.finalized) != 0 || streamer.canceled { + t.Fatalf("streamer should stay unused, got updates=%#v finalized=%#v canceled=%v", streamer.updates, streamer.finalized, streamer.canceled) + } +} + +func TestProcessMessage_CancelsStreamingWhenToolCallsFollow(t *testing.T) { + provider := &streamingTestProvider{ + chunks: []string{"Thinking"}, + response: &providers.LLMResponse{ + Content: "Thinking", + ToolCalls: []providers.ToolCall{ + { + ID: "call-1", + Name: "missing_tool", + Arguments: map[string]any{ + "q": "stream", + }, + }, + }, + }, + } + streamer := &recordingStreamer{} + al, cleanup := newStreamingTestLoop(t, true, provider, streamer) + defer cleanup() + + response, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "pico", + ChatID: "pico:test-session", + SenderID: "user-1", + Content: "hello", + })) + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + + if response == "" { + t.Fatal("expected follow-up tool result content after tool-call path") + } + if !provider.chatStreamCalled { + t.Fatal("expected ChatStream to be used") + } + if !streamer.canceled { + t.Fatal("expected streamer to be canceled when tool calls follow streamed text") + } + if len(streamer.finalized) != 0 { + t.Fatalf("finalized = %#v, want no finalize after tool calls", streamer.finalized) + } +} diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index d1de8f4d5..ef87ff641 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -107,6 +107,16 @@ type PicoChannel struct { deleteMessageFn func(context.Context, string, string) error } +type picoStreamer struct { + channel *PicoChannel + chatID string + messageID string + + mu sync.Mutex + started bool + closed bool +} + // NewPicoChannel creates a new Pico Protocol channel. func NewPicoChannel( bc *config.Channel, @@ -354,6 +364,18 @@ func (c *PicoChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri return []string{msgID}, nil } +func (c *PicoChannel) BeginStream(_ context.Context, chatID string) (channels.Streamer, error) { + if !c.IsRunning() { + return nil, channels.ErrNotRunning + } + + return &picoStreamer{ + channel: c, + chatID: chatID, + messageID: uuid.New().String(), + }, nil +} + // EditMessage implements channels.MessageEditor. func (c *PicoChannel) EditMessage(ctx context.Context, chatID string, messageID string, content string) error { return c.editMessage(ctx, chatID, messageID, content, nil) @@ -440,6 +462,74 @@ func (c *PicoChannel) FinalizeToolFeedbackMessage(ctx context.Context, msg bus.O return c.finalizeTrackedToolFeedbackMessage(ctx, msg.ChatID, msg.Content, c.editMessage, msg.ContextUsage) } +func (s *picoStreamer) Update(_ context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return nil + } + + if !s.started { + s.started = true + return s.channel.broadcastToSession(s.chatID, newMessage(TypeMessageCreate, map[string]any{ + "message_id": s.messageID, + PayloadKeyContent: content, + PayloadKeyThought: false, + PayloadKeyStreaming: true, + })) + } + + return s.channel.broadcastToSession(s.chatID, newMessage(TypeMessageUpdate, map[string]any{ + "message_id": s.messageID, + "content": content, + "streaming": true, + })) +} + +func (s *picoStreamer) Finalize(_ context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return nil + } + s.closed = true + + if !s.started { + return s.channel.broadcastToSession(s.chatID, newMessage(TypeMessageCreate, map[string]any{ + "message_id": s.messageID, + PayloadKeyContent: content, + PayloadKeyThought: false, + PayloadKeyStreaming: false, + })) + } + + return s.channel.broadcastToSession(s.chatID, newMessage(TypeMessageUpdate, map[string]any{ + "message_id": s.messageID, + "content": content, + "streaming": false, + })) +} + +func (s *picoStreamer) Cancel(_ context.Context) { + s.mu.Lock() + defer s.mu.Unlock() + + if s.closed { + return + } + s.closed = true + + if !s.started { + return + } + + _ = s.channel.broadcastToSession(s.chatID, newMessage(TypeMessageDelete, map[string]any{ + "message_id": s.messageID, + })) +} + // StartTyping implements channels.TypingCapable. func (c *PicoChannel) StartTyping(ctx context.Context, chatID string) (func(), error) { startMsg := newMessage(TypeTypingStart, nil) diff --git a/pkg/channels/pico/protocol.go b/pkg/channels/pico/protocol.go index 6e3a5ca89..6f4435264 100644 --- a/pkg/channels/pico/protocol.go +++ b/pkg/channels/pico/protocol.go @@ -22,10 +22,13 @@ const ( TypeError = "error" TypePong = "pong" + PicoTokenPrefix = "pico-" + PayloadKeyContent = "content" PayloadKeyThought = "thought" PayloadKeyKind = "kind" PayloadKeyToolCalls = "tool_calls" + PayloadKeyStreaming = "streaming" MessageKindThought = "thought" MessageKindToolCalls = "tool_calls" diff --git a/pkg/config/config.go b/pkg/config/config.go index dc9e88949..d0b3f9207 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -273,6 +273,7 @@ type AgentDefaults struct { MaxParallelTurns int `json:"max_parallel_turns,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_PARALLEL_TURNS"` // Max concurrent turns (0 or 1 = sequential) SubTurn SubTurnConfig `json:"subturn" envPrefix:"PICOCLAW_AGENTS_DEFAULTS_SUBTURN_"` ToolFeedback ToolFeedbackConfig `json:"tool_feedback,omitempty"` + StreamingEnabled bool `json:"streaming_enabled" env:"PICOCLAW_AGENTS_DEFAULTS_STREAMING_ENABLED"` SplitOnMarker bool `json:"split_on_marker" env:"PICOCLAW_AGENTS_DEFAULTS_SPLIT_ON_MARKER"` // split messages on <|[SPLIT]|> marker ContextManager string `json:"context_manager,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_CONTEXT_MANAGER"` ContextManagerConfig json.RawMessage `json:"context_manager_config,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_CONTEXT_MANAGER_CONFIG"` diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index be8c32495..13a6b567f 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -39,7 +39,8 @@ func DefaultConfig() *Config { MaxArgsLength: 300, SeparateMessages: false, }, - SplitOnMarker: false, + StreamingEnabled: true, + SplitOnMarker: false, }, }, Session: SessionConfig{ diff --git a/web/frontend/src/components/chat/assistant-message.tsx b/web/frontend/src/components/chat/assistant-message.tsx index 07a3c0abc..f3969ea37 100644 --- a/web/frontend/src/components/chat/assistant-message.tsx +++ b/web/frontend/src/components/chat/assistant-message.tsx @@ -9,12 +9,8 @@ import { } from "@tabler/icons-react" import { useState } from "react" import { useTranslation } from "react-i18next" -import ReactMarkdown from "react-markdown" -import rehypeHighlight from "rehype-highlight" -import rehypeRaw from "rehype-raw" -import rehypeSanitize from "rehype-sanitize" -import remarkGfm from "remark-gfm" +import { StreamingMarkdown } from "@/components/chat/streaming-markdown" import { Button } from "@/components/ui/button" import { formatMessageTime } from "@/hooks/use-pico-chat" import { cn } from "@/lib/utils" @@ -29,6 +25,8 @@ interface AssistantMessageProps { attachments?: ChatAttachment[] kind?: AssistantMessageKind toolCalls?: ChatToolCall[] + streamingEnabled?: boolean + isStreaming?: boolean timestamp?: string | number } @@ -37,6 +35,8 @@ export function AssistantMessage({ attachments = [], kind = "normal", toolCalls = [], + streamingEnabled = true, + isStreaming = false, timestamp = "", }: AssistantMessageProps) { const { t } = useTranslation() @@ -184,21 +184,16 @@ export function AssistantMessage({ )} {(!isCollapsedBlock || isExpanded) && !isToolCalls && hasText && ( -