diff --git a/pkg/agent/adapters/messagebus.go b/pkg/agent/adapters/messagebus.go index ccae7e8bc..7284c5972 100644 --- a/pkg/agent/adapters/messagebus.go +++ b/pkg/agent/adapters/messagebus.go @@ -34,3 +34,7 @@ func (a *messageBusAdapter) PublishOutboundMedia(ctx context.Context, msg bus.Ou func (a *messageBusAdapter) InboundChan() <-chan bus.InboundMessage { return a.inner.InboundChan() } + +func (a *messageBusAdapter) GetStreamer(ctx context.Context, channel, chatID string) (bus.Streamer, bool) { + return a.inner.GetStreamer(ctx, channel, chatID) +} diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 5749149c1..4a90f35ec 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -560,11 +560,33 @@ func (al *AgentLoop) runAgentLoop( newTurnContext(opts.Dispatch.InboundContext, opts.Dispatch.RouteResult, opts.Dispatch.SessionScope), ) ts := newTurnState(agent, opts, turnScope) + + // Acquire streamer if channel supports streaming + shouldStream := opts.SendResponse || opts.AllowInterimPicoPublish + if al.bus != nil && shouldStream && opts.Dispatch.Channel() != "" { + if streamer, ok := al.bus.GetStreamer(ctx, opts.Dispatch.Channel(), opts.Dispatch.ChatID()); ok { + ts.setStreamer(streamer) + logger.DebugCF("agent", "Streaming enabled for turn", map[string]any{ + "channel": opts.Dispatch.Channel(), + "chat_id": opts.Dispatch.ChatID(), + }) + } + } + pipeline := NewPipeline(al) result, err := al.runTurn(ctx, ts, pipeline) if err != nil { + if ts.getStreamer() != nil { + ts.cancelStreamer(ctx) + } return "", err } + + // Handle streamer cleanup on abort or error + if ts.getStreamer() != nil && (result.status == TurnEndStatusAborted || result.status == TurnEndStatusError) { + ts.cancelStreamer(ctx) + } + if result.status == TurnEndStatusAborted { return "", nil } @@ -579,7 +601,8 @@ func (al *AgentLoop) runAgentLoop( } } - if opts.SendResponse && result.finalContent != "" { + // Only publish via bus if not already streamed + if opts.SendResponse && result.finalContent != "" && !ts.wasStreamed() { agentID, sessionKey, scope := outboundTurnMetadata( agent.ID, opts.Dispatch.SessionKey, diff --git a/pkg/agent/agent_outbound.go b/pkg/agent/agent_outbound.go index f4a01adfd..f8155be9d 100644 --- a/pkg/agent/agent_outbound.go +++ b/pkg/agent/agent_outbound.go @@ -223,6 +223,12 @@ func (al *AgentLoop) publishPicoToolCallInterim( "chat_id": ts.chatID, "error": err.Error(), }) + } else if err == nil { + logger.InfoCF("agent", "Published pico tool calls", map[string]any{ + "channel": ts.channel, + "chat_id": ts.chatID, + "tool_count": len(visibleToolCalls), + }) } } diff --git a/pkg/agent/interfaces/interfaces.go b/pkg/agent/interfaces/interfaces.go index 2efec05e1..3c4a9669d 100644 --- a/pkg/agent/interfaces/interfaces.go +++ b/pkg/agent/interfaces/interfaces.go @@ -23,6 +23,10 @@ type MessageBus interface { // InboundChan returns the channel for receiving inbound messages. InboundChan() <-chan bus.InboundMessage + + // GetStreamer returns a Streamer for the given channel+chatID if the channel + // supports streaming. + GetStreamer(ctx context.Context, channel, chatID string) (bus.Streamer, bool) } // ChannelManager manages channel lifecycle and provides channel access. diff --git a/pkg/agent/pipeline_llm.go b/pkg/agent/pipeline_llm.go index 3a3c496f6..daa1d561f 100644 --- a/pkg/agent/pipeline_llm.go +++ b/pkg/agent/pipeline_llm.go @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/constants" runtimeevents "github.com/sipeed/picoclaw/pkg/events" "github.com/sipeed/picoclaw/pkg/logger" @@ -144,7 +145,7 @@ func (p *Pipeline) CallLLM( }) // LLM call closure with fallback support - callLLM := func(messagesForCall []providers.Message, toolDefsForCall []providers.ToolDefinition) (*providers.LLMResponse, error) { + callLLM := func(messagesForCall []providers.Message, toolDefsForCall []providers.ToolDefinition, streamer bus.Streamer) (*providers.LLMResponse, error) { providerCtx, providerCancel := context.WithCancel(turnCtx) ts.setProviderCancel(providerCancel) defer func() { @@ -155,6 +156,19 @@ func (p *Pipeline) CallLLM( al.activeRequests.Add(1) defer al.activeRequests.Done() + // Use streaming if available (provider handles tool calls in stream) + useStreaming := streamer != nil + if sp, ok := exec.activeProvider.(providers.StreamingProvider); ok && useStreaming { + onChunk := func(accumulated string) { + if err := streamer.Update(providerCtx, accumulated); err != nil { + logger.DebugCF("agent", "Streaming update failed", map[string]any{ + "error": err.Error(), + }) + } + } + return sp.ChatStream(providerCtx, messagesForCall, toolDefsForCall, exec.llmModel, exec.llmOpts, onChunk) + } + if len(exec.activeCandidates) > 1 && p.Fallback != nil { fbResult, fbErr := p.Fallback.Execute( providerCtx, @@ -194,7 +208,12 @@ func (p *Pipeline) CallLLM( backoffSecs = 2 } for retry := 0; retry <= maxRetries; retry++ { - exec.response, err = callLLM(exec.callMessages, exec.providerToolDefs) + // Only stream on first attempt to avoid duplicate content + var callStreamer bus.Streamer + if retry == 0 { + callStreamer = ts.getStreamer() + } + exec.response, err = callLLM(exec.callMessages, exec.providerToolDefs, callStreamer) if err == nil { break } @@ -490,6 +509,10 @@ func (p *Pipeline) CallLLM( return ControlContinue, nil } exec.finalContent = responseContent + // Finalize streaming if active + if ts.getStreamer() != nil && exec.finalContent != "" { + ts.finalizeStreamer(turnCtx, exec.finalContent) + } logger.InfoCF("agent", "LLM response without tool calls (direct answer)", map[string]any{ "agent_id": ts.agent.ID, diff --git a/pkg/agent/turn_state.go b/pkg/agent/turn_state.go index ae058e49d..f0b9b516b 100644 --- a/pkg/agent/turn_state.go +++ b/pkg/agent/turn_state.go @@ -237,6 +237,11 @@ type turnState struct { // Back-reference to the owning AgentLoop (set for SubTurns only, used for hard abort cascade) al *AgentLoop + + // Streaming support + streamer bus.Streamer + streamerOnce sync.Once + streamerFinalized bool } // ============================================================================= @@ -810,3 +815,40 @@ func turnStateFromContext(ctx context.Context) *turnState { func TurnStateFromContext(ctx context.Context) *turnState { return turnStateFromContext(ctx) } + +// Streamer management methods + +func (ts *turnState) setStreamer(s bus.Streamer) { + ts.mu.Lock() + defer ts.mu.Unlock() + ts.streamer = s +} + +func (ts *turnState) getStreamer() bus.Streamer { + ts.mu.RLock() + defer ts.mu.RUnlock() + return ts.streamer +} + +func (ts *turnState) finalizeStreamer(ctx context.Context, content string) { + ts.streamerOnce.Do(func() { + if ts.streamer != nil { + _ = ts.streamer.Finalize(ctx, content) + ts.streamerFinalized = true + } + }) +} + +func (ts *turnState) wasStreamed() bool { + ts.mu.RLock() + defer ts.mu.RUnlock() + return ts.streamerFinalized +} + +func (ts *turnState) cancelStreamer(ctx context.Context) { + ts.streamerOnce.Do(func() { + if ts.streamer != nil { + ts.streamer.Cancel(ctx) + } + }) +} diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index d1de8f4d5..8a51e56b6 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -91,6 +91,56 @@ func (pc *picoConn) close() { } } +// picoStreamer implements channels.Streamer for real-time token streaming. +type picoStreamer struct { + channel *PicoChannel + chatID string + messageID string + content string + mu sync.Mutex + finalized bool + lastUpdateAt time.Time + throttleInterval time.Duration + minGrowth int +} + +func (s *picoStreamer) Update(ctx context.Context, content string) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.finalized { + return nil + } + + now := time.Now() + growth := len(content) - len(s.content) + + // Skip if not enough growth AND not enough time elapsed + if growth < s.minGrowth && time.Since(s.lastUpdateAt) < s.throttleInterval { + return nil + } + + s.content = content + s.lastUpdateAt = now + 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() + if s.finalized { + return nil + } + s.finalized = true + s.content = content + return s.channel.EditMessage(ctx, s.chatID, s.messageID, content) +} + +func (s *picoStreamer) Cancel(ctx context.Context) { + s.mu.Lock() + defer s.mu.Unlock() + s.finalized = true +} + // PicoChannel implements the native Pico Protocol WebSocket channel. // It serves as the reference implementation for all optional capability interfaces. type PicoChannel struct { @@ -679,6 +729,38 @@ func (c *PicoChannel) handleMediaDownload(w http.ResponseWriter, r *http.Request http.ServeContent(w, r, filename, info.ModTime(), file) } +// BeginStream implements channels.StreamingCapable. +func (c *PicoChannel) BeginStream(ctx context.Context, chatID string) (channels.Streamer, error) { + if !c.IsRunning() { + return nil, channels.ErrNotRunning + } + if !c.config.Streaming { + return nil, fmt.Errorf("streaming disabled in config") + } + + msgID := uuid.New().String() + outMsg := newMessage(TypeMessageCreate, map[string]any{ + PayloadKeyContent: "", + PayloadKeyThought: false, + "message_id": msgID, + }) + + sessionID := strings.TrimPrefix(chatID, "pico:") + outMsg.SessionID = sessionID + + if err := c.broadcastToSession(chatID, outMsg); err != nil { + return nil, err + } + + return &picoStreamer{ + channel: c, + chatID: chatID, + messageID: msgID, + throttleInterval: 100 * time.Millisecond, + minGrowth: 20, + }, nil +} + // broadcastToSession sends a message to all connections with a matching session. func (c *PicoChannel) broadcastToSession(chatID string, msg PicoMessage) error { // chatID format: "pico:" diff --git a/pkg/config/config.go b/pkg/config/config.go index 0a347e3db..ef7d79336 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -580,13 +580,14 @@ func (c *WeixinSettings) SetToken(token string) { } type PicoSettings struct { - Token SecureString `json:"token,omitzero" yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_PICO_TOKEN"` + Token SecureString `json:"token,omitzero" yaml:"token,omitempty" env:"PICOCLAW_CHANNELS_PICO_TOKEN"` AllowTokenQuery bool `json:"allow_token_query,omitempty" yaml:"-"` AllowOrigins []string `json:"allow_origins,omitempty" yaml:"-"` PingInterval int `json:"ping_interval,omitempty" yaml:"-"` ReadTimeout int `json:"read_timeout,omitempty" yaml:"-"` WriteTimeout int `json:"write_timeout,omitempty" yaml:"-"` MaxConnections int `json:"max_connections,omitempty" yaml:"-"` + Streaming bool `json:"streaming,omitempty" yaml:"streaming,omitempty"` } // SetToken sets the Pico token and marks it as dirty for security saving