opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: stream provider output to compatible channels

Browse source 
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.
This commit is contained in:
alanbulan 2026-04-07 22:40:10 +08:00
parent 6ce0306c66
commit 15888963ea
8 changed files with 350 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
pkg/agent/loop.go
View file

@ -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

71
pkg/agent/loop_test.go
View file

@ -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")

31
pkg/agent/mock_provider_test.go
View file

@ -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
}

72
pkg/channels/pico/pico.go
View file

@ -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).

83
pkg/channels/pico/pico_test.go
View file

@ -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()

1
pkg/channels/pico/protocol.go
View file

@ -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"

30
pkg/providers/openai_compat/provider_test.go
View file

@ -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

12
web/frontend/src/features/chat/protocol.ts
View file

@ -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