diff --git a/AGENT_LOOP_IMPROVEMENTS.md b/AGENT_LOOP_IMPROVEMENTS.md index 338f5bde7..9bd38f27e 100644 --- a/AGENT_LOOP_IMPROVEMENTS.md +++ b/AGENT_LOOP_IMPROVEMENTS.md @@ -22,7 +22,7 @@ This document outlines a series of tasks to improve the main loop of the agentic * [x] **Concurrent Message Processing:** Evaluate introducing a worker pool or goroutines to process independent user requests concurrently without blocking the entire agent instance. * [x] **Background Summarization:** Offload `maybeSummarize` and context compression to an asynchronous worker. Instead of blocking the main thread, the worker outputs its execution trace and compressed context to a structured `/logs/{session}/{subagent}/` directory to provide a persistent audit trail while keeping the loop responsive. -* [ ] **Streaming Responses:** Implement streaming LLM token generation directly to the `bus.PublishOutbound` instead of waiting for full generation. +* [x] **Streaming Responses:** Implement streaming LLM token generation directly to the `bus.PublishOutbound` instead of waiting for full generation. ## Phase 4: Features diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index d3932fc88..3bc589cc2 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -53,6 +53,7 @@ type processOptions struct { EnableSummary bool // Whether to trigger summarization SendResponse bool // Whether to send response via bus NoHistory bool // If true, don't load session history (for heartbeat) + Stream bool // Whether to stream LLM generation } const ( diff --git a/pkg/agent/loop_execute_llm.go b/pkg/agent/loop_execute_llm.go index fb8023896..12b132af4 100644 --- a/pkg/agent/loop_execute_llm.go +++ b/pkg/agent/loop_execute_llm.go @@ -31,6 +31,29 @@ func (al *AgentLoop) executeLLMWithRetry( "prompt_cache_key": agent.ID, } + if opts.Stream { + llmOpts["stream_callback"] = func(ctx context.Context, chunk string, isReasoning bool) { + pubCtx, pubCancel := context.WithTimeout(ctx, 2*time.Second) + defer pubCancel() + + if err := al.bus.PublishOutboundStream(pubCtx, bus.OutboundStreamMessage{ + Channel: opts.Channel, + ChatID: opts.ChatID, + Content: chunk, + IsReasoning: isReasoning, + }); err != nil { + // Don't log full timeout errors for every dropped streaming chunk to avoid log spam, + // just drop the chunk. + if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, bus.ErrBusClosed) { + logger.DebugCF("agent", "Failed to publish stream chunk", map[string]any{ + "channel": opts.Channel, + "error": err.Error(), + }) + } + } + } + } + if agent.ThinkingLevel != ThinkingOff { if tc, ok := agent.Provider.(providers.ThinkingCapable); ok && tc.SupportsThinking() { llmOpts["thinking_level"] = string(agent.ThinkingLevel) diff --git a/pkg/agent/loop_process.go b/pkg/agent/loop_process.go index fff717c91..e24f8ad46 100644 --- a/pkg/agent/loop_process.go +++ b/pkg/agent/loop_process.go @@ -134,6 +134,7 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) DefaultResponse: defaultResponse, EnableSummary: true, SendResponse: false, + Stream: true, // Hardcoded for Phase 1 streaming implementation } // Medical Persona specific routing interception diff --git a/pkg/bus/bus.go b/pkg/bus/bus.go index 5a647d042..7d57a3623 100644 --- a/pkg/bus/bus.go +++ b/pkg/bus/bus.go @@ -11,22 +11,27 @@ import ( // ErrBusClosed is returned when publishing to a closed MessageBus. var ErrBusClosed = errors.New("message bus closed") +// ErrBusFull is returned when publishing to a full MessageBus. +var ErrBusFull = errors.New("message bus full") + const defaultBusBufferSize = 64 type MessageBus struct { - inbound chan InboundMessage - outbound chan OutboundMessage - outboundMedia chan OutboundMediaMessage - done chan struct{} - closed atomic.Bool + inbound chan InboundMessage + outbound chan OutboundMessage + outboundStream chan OutboundStreamMessage + outboundMedia chan OutboundMediaMessage + done chan struct{} + closed atomic.Bool } func NewMessageBus() *MessageBus { return &MessageBus{ - inbound: make(chan InboundMessage, defaultBusBufferSize), - outbound: make(chan OutboundMessage, defaultBusBufferSize), - outboundMedia: make(chan OutboundMediaMessage, defaultBusBufferSize), - done: make(chan struct{}), + inbound: make(chan InboundMessage, defaultBusBufferSize), + outbound: make(chan OutboundMessage, defaultBusBufferSize), + outboundStream: make(chan OutboundStreamMessage, defaultBusBufferSize), + outboundMedia: make(chan OutboundMediaMessage, defaultBusBufferSize), + done: make(chan struct{}), } } @@ -114,6 +119,37 @@ func (mb *MessageBus) SubscribeOutboundMedia(ctx context.Context) (OutboundMedia } } +func (mb *MessageBus) PublishOutboundStream(ctx context.Context, msg OutboundStreamMessage) error { + if mb.closed.Load() { + return ErrBusClosed + } + if err := ctx.Err(); err != nil { + return err + } + select { + case mb.outboundStream <- msg: + return nil + case <-mb.done: + return ErrBusClosed + case <-ctx.Done(): + return ctx.Err() + default: + // Non-blocking send. If the buffer is full, immediately drop the token. + return ErrBusFull + } +} + +func (mb *MessageBus) SubscribeOutboundStream(ctx context.Context) (OutboundStreamMessage, bool) { + select { + case msg, ok := <-mb.outboundStream: + return msg, ok + case <-mb.done: + return OutboundStreamMessage{}, false + case <-ctx.Done(): + return OutboundStreamMessage{}, false + } +} + func (mb *MessageBus) Close() { if mb.closed.CompareAndSwap(false, true) { close(mb.done) @@ -139,6 +175,15 @@ func (mb *MessageBus) Close() { } } doneOutbound: + for { + select { + case <-mb.outboundStream: + drained++ + default: + goto doneStream + } + } + doneStream: for { select { case <-mb.outboundMedia: diff --git a/pkg/bus/types.go b/pkg/bus/types.go index 12da3f1dd..476cc5ef0 100644 --- a/pkg/bus/types.go +++ b/pkg/bus/types.go @@ -36,6 +36,13 @@ type OutboundMessage struct { ReplyToMessageID string `json:"reply_to_message_id,omitempty"` } +type OutboundStreamMessage struct { + Channel string `json:"channel"` + ChatID string `json:"chat_id"` + Content string `json:"content"` + IsReasoning bool `json:"is_reasoning"` +} + // MediaPart describes a single media attachment to send. type MediaPart struct { Type string `json:"type"` // "image" | "audio" | "video" | "file"