This commit is contained in:
Andy Lo-A-Foe 2026-05-15 11:29:59 +03:00 committed by GitHub
commit 0dd24117ed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 189 additions and 4 deletions

View file

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

View file

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

View file

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

View file

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

View file

@ -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
}
@ -493,6 +512,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,

View file

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

View file

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

View file

@ -581,13 +581,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