From 3a68d268371d5328f008bc695c41585100e5549a Mon Sep 17 00:00:00 2001 From: SiYue-ZO <2835601846@qq.com> Date: Thu, 21 May 2026 01:14:42 +0800 Subject: [PATCH 1/3] Fix agent loop reload and panic cleanup stability --- pkg/agent/agent.go | 51 +++----- pkg/agent/agent_test.go | 106 +++++++++++++++++ pkg/agent/agent_utils.go | 50 ++++++++ pkg/agent/runtime_event_logger_test.go | 159 +++++++++++++++++++++++++ pkg/agent/turn_state.go | 12 +- 5 files changed, 343 insertions(+), 35 deletions(-) diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 0392829f..2723831f 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -117,6 +117,7 @@ const ( handledToolResponseSummary = "Requested output delivered via tool attachment." sessionKeyAgentPrefix = "agent:" pendingTurnPrefix = "pending-" + providerReloadGracePeriod = 30 * time.Second metadataKeyMessageKind = "message_kind" metadataKeyToolCalls = "tool_calls" metadataKeyOutboundKind = "outbound_kind" @@ -208,6 +209,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { // slot. The goroutine is spawned immediately so the main loop keeps // draining the inbound channel. The goroutine blocks on the semaphore. go func(m bus.InboundMessage) { + var releaseSession bool // Acquire semaphore slot (blocks if at capacity) select { case al.workerSem <- struct{}{}: @@ -215,7 +217,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { case <-ctx.Done(): // Context canceled while waiting for a slot — clean up the // placeholder to prevent session-level deadlock. - al.activeTurnStates.Delete(sessionKey) + al.releaseSessionTurnState(sessionKey, nil) return } @@ -224,16 +226,21 @@ func (al *AgentLoop) Run(ctx context.Context) error { // completes normally, clearActiveTurn deletes the real turnState and // this becomes a no-op (the key is already gone). defer func() { + if releaseSession { + al.releaseSessionTurnState(sessionKey, nil) + return + } if actual, ok := al.activeTurnStates.Load(sessionKey); ok { if ts, ok := actual.(*turnState); ok && strings.HasPrefix(ts.turnID, pendingTurnPrefix) { // Placeholder still present — runTurn never replaced it. - al.activeTurnStates.Delete(sessionKey) + al.releaseSessionTurnState(sessionKey, ts) } } }() defer func() { if r := recover(); r != nil { + releaseSession = true logger.RecoverPanicNoExit(r) logger.ErrorCF("agent", "Worker goroutine panicked", map[string]any{ @@ -251,7 +258,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { } if al.takePendingStop(sessionKey) { - al.activeTurnStates.Delete(sessionKey) + al.releaseSessionTurnState(sessionKey, nil) target := &continuationTarget{ SessionKey: sessionKey, Channel: m.Channel, @@ -365,37 +372,23 @@ func (al *AgentLoop) ReloadProviderAndConfig( return fmt.Errorf("config cannot be nil") } - // Create new registry with updated config and provider - // Wrap in defer/recover to handle any panics gracefully var registry *AgentRegistry - var panicErr error - done := make(chan struct{}, 1) - - go func() { + func() { defer func() { if r := recover(); r != nil { logger.RecoverPanicNoExit(r) - panicErr = fmt.Errorf("panic during registry creation: %v", r) logger.ErrorCF("agent", "Panic during registry creation", map[string]any{"panic": r}) + registry = nil } - close(done) }() - registry = NewAgentRegistry(cfg, provider) }() - - // Wait for completion or context cancellation - select { - case <-done: - if registry == nil { - if panicErr != nil { - return fmt.Errorf("registry creation failed: %w", panicErr) - } - return fmt.Errorf("registry creation failed (nil result)") + if registry == nil { + if err := ctx.Err(); err != nil { + return fmt.Errorf("context canceled during registry creation: %w", err) } - case <-ctx.Done(): - return fmt.Errorf("context canceled during registry creation: %w", ctx.Err()) + return fmt.Errorf("registry creation failed") } // Check context again before proceeding @@ -471,17 +464,7 @@ func (al *AgentLoop) ReloadProviderAndConfig( // This prevents blocking readers while closing if oldProvider, ok := extractProvider(oldRegistry); ok { if stateful, ok := oldProvider.(providers.StatefulProvider); ok { - // Give in-flight requests a moment to complete - // Use a reasonable timeout that balances cleanup vs resource usage - select { - case <-time.After(100 * time.Millisecond): - stateful.Close() - case <-ctx.Done(): - // Context canceled, close immediately but log warning - logger.WarnCF("agent", "Context canceled during provider cleanup, forcing close", - map[string]any{"error": ctx.Err()}) - stateful.Close() - } + al.closeReloadedProvider(ctx, stateful) } } diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index aaf3d1a8..58541140 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -13,6 +13,7 @@ import ( "slices" "strings" "sync" + "sync/atomic" "testing" "time" @@ -119,6 +120,31 @@ type recordingProvider struct { lastModel string } +type panicAfterStartProvider struct { + started chan struct{} + calls atomic.Int32 +} + +func (p *panicAfterStartProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + p.calls.Add(1) + select { + case <-p.started: + default: + close(p.started) + } + panic("provider panic after turn registration") +} + +func (p *panicAfterStartProvider) GetDefaultModel() string { + return "panic-after-start" +} + func (r *recordingProvider) Chat( ctx context.Context, messages []providers.Message, @@ -5710,3 +5736,83 @@ func (p *concurrentMockProvider) Chat( func (p *concurrentMockProvider) GetDefaultModel() string { return "test-model" } + +func TestRunWorkerPanicReleasesSessionTurnState(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = t.TempDir() + cfg.Agents.Defaults.MaxParallelTurns = 1 + + msgBus := bus.NewMessageBus() + provider := &panicAfterStartProvider{started: make(chan struct{})} + al := NewAgentLoop(cfg, msgBus, provider) + defer al.Close() + + runCtx, cancelRun := context.WithCancel(context.Background()) + runDone := make(chan error, 1) + go func() { + runDone <- al.Run(runCtx) + }() + defer func() { + cancelRun() + select { + case err := <-runDone: + if err != nil { + t.Fatalf("Run() error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for Run() to exit") + } + }() + + msg := bus.InboundMessage{ + Context: bus.InboundContext{ + Channel: "test", + ChatID: "panic-chat", + ChatType: "direct", + SenderID: "user1", + }, + Content: "trigger panic", + SessionKey: "panic-session", + } + route, _, err := al.resolveMessageRoute(msg) + if err != nil { + t.Fatalf("resolveMessageRoute() error = %v", err) + } + scopeKey := resolveScopeKey(al.allocateRouteSession(route, msg).SessionKey, msg.SessionKey) + + if err := msgBus.PublishInbound(context.Background(), msg); err != nil { + t.Fatalf("PublishInbound(first) error = %v", err) + } + + select { + case <-provider.started: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for first turn to start") + } + + deadline := time.Now().Add(2 * time.Second) + for { + if al.getActiveTurnState(scopeKey) == nil { + break + } + if time.Now().After(deadline) { + t.Fatal("session turn state remained stuck after worker panic") + } + time.Sleep(10 * time.Millisecond) + } + + if err := msgBus.PublishInbound(context.Background(), msg); err != nil { + t.Fatalf("PublishInbound(second) error = %v", err) + } + + deadline = time.Now().Add(2 * time.Second) + for { + if provider.calls.Load() >= 2 { + break + } + if time.Now().After(deadline) { + t.Fatal("second message did not start a new turn after panic cleanup") + } + time.Sleep(10 * time.Millisecond) + } +} diff --git a/pkg/agent/agent_utils.go b/pkg/agent/agent_utils.go index 432a9f24..e1da3bf6 100644 --- a/pkg/agent/agent_utils.go +++ b/pkg/agent/agent_utils.go @@ -13,6 +13,7 @@ import ( "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/commands" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/session" "github.com/sipeed/picoclaw/pkg/utils" @@ -584,6 +585,55 @@ func closeProviderIfStateful(provider providers.LLMProvider) { } } +func (al *AgentLoop) waitForActiveRequests(ctx context.Context, timeout time.Duration) bool { + done := make(chan struct{}) + go func() { + al.activeRequests.Wait() + close(done) + }() + + if timeout <= 0 { + select { + case <-done: + return true + case <-ctx.Done(): + return false + } + } + + timer := time.NewTimer(timeout) + defer timer.Stop() + + select { + case <-done: + return true + case <-timer.C: + return false + case <-ctx.Done(): + return false + } +} + +func (al *AgentLoop) closeReloadedProvider(ctx context.Context, provider providers.StatefulProvider) { + waitCtx := ctx + if waitCtx == nil { + waitCtx = context.Background() + } + + drained := al.waitForActiveRequests(waitCtx, providerReloadGracePeriod) + if !drained { + fields := map[string]any{"grace_period": providerReloadGracePeriod.String()} + if err := waitCtx.Err(); err != nil { + fields["error"] = err.Error() + logger.WarnCF("agent", "Provider reload interrupted while waiting for in-flight requests", fields) + } else { + logger.WarnCF("agent", "Provider reload grace period expired with in-flight requests still running", fields) + } + } + + provider.Close() +} + func makePendingTurnID(sessionKey string, seq uint64) string { return pendingTurnPrefix + sessionKey + "-" + fmt.Sprintf("%d", seq) } diff --git a/pkg/agent/runtime_event_logger_test.go b/pkg/agent/runtime_event_logger_test.go index 1c95b365..1da9803c 100644 --- a/pkg/agent/runtime_event_logger_test.go +++ b/pkg/agent/runtime_event_logger_test.go @@ -2,6 +2,7 @@ package agent import ( "context" + "errors" "sync/atomic" "testing" "time" @@ -9,6 +10,7 @@ import ( "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" runtimeevents "github.com/sipeed/picoclaw/pkg/events" + "github.com/sipeed/picoclaw/pkg/providers" ) func TestRuntimeEventLoggerFiltering(t *testing.T) { @@ -191,6 +193,163 @@ func TestReloadProviderAndConfigRefreshesRuntimeEventLogger(t *testing.T) { } } +type reloadBlockingProvider struct { + chatStarted chan struct{} + releaseChat chan struct{} + closeCalled chan struct{} +} + +func (p *reloadBlockingProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + select { + case <-p.chatStarted: + default: + close(p.chatStarted) + } + + select { + case <-p.releaseChat: + return &providers.LLMResponse{Content: "done"}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (p *reloadBlockingProvider) GetDefaultModel() string { + return "reload-blocking" +} + +func (p *reloadBlockingProvider) Close() { + select { + case <-p.closeCalled: + default: + close(p.closeCalled) + } +} + +func TestReloadProviderAndConfigWaitsForInFlightRequestsBeforeClosingOldProvider(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = t.TempDir() + + oldProvider := &reloadBlockingProvider{ + chatStarted: make(chan struct{}), + releaseChat: make(chan struct{}), + closeCalled: make(chan struct{}), + } + al := NewAgentLoop(cfg, bus.NewMessageBus(), oldProvider) + defer al.Close() + + msg := testInboundMessage(bus.InboundMessage{ + Channel: "test", + ChatID: "reload-chat", + SenderID: "user-1", + Content: "hold request open", + }) + + reqDone := make(chan error, 1) + go func() { + _, err := al.processMessage(context.Background(), msg) + reqDone <- err + }() + + select { + case <-oldProvider.chatStarted: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for in-flight provider request") + } + + reloadDone := make(chan error, 1) + go func() { + reloaded := config.DefaultConfig() + reloaded.Agents.Defaults.Workspace = cfg.Agents.Defaults.Workspace + reloadDone <- al.ReloadProviderAndConfig(context.Background(), &mockProvider{}, reloaded) + }() + + select { + case <-oldProvider.closeCalled: + t.Fatal("old provider closed before in-flight request completed") + case err := <-reloadDone: + t.Fatalf("reload returned early: %v", err) + case <-time.After(150 * time.Millisecond): + } + + close(oldProvider.releaseChat) + + select { + case err := <-reqDone: + if err != nil { + t.Fatalf("processMessage() error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for in-flight request to complete") + } + + select { + case err := <-reloadDone: + if err != nil { + t.Fatalf("ReloadProviderAndConfig() error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for reload to finish") + } + + select { + case <-oldProvider.closeCalled: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for old provider close") + } +} + +func TestWaitForActiveRequestsHonorsContextCancellation(t *testing.T) { + al := &AgentLoop{} + al.activeRequests.Add(1) + defer al.activeRequests.Done() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + if al.waitForActiveRequests(ctx, time.Second) { + t.Fatal("waitForActiveRequests() = true, want false on canceled context") + } +} + +func TestReloadProviderAndConfigReturnsCanceledErrorWhenRegistryCreationPanics(t *testing.T) { + cfg := config.DefaultConfig() + cfg.Agents.Defaults.Workspace = t.TempDir() + + al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) + defer al.Close() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + err := al.ReloadProviderAndConfig(ctx, &panicProviderForReloadTest{}, cfg) + if !errors.Is(err, context.Canceled) { + t.Fatalf("ReloadProviderAndConfig() error = %v, want context canceled", err) + } +} + +type panicProviderForReloadTest struct{} + +func (p *panicProviderForReloadTest) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + return &providers.LLMResponse{Content: "unused"}, nil +} + +func (p *panicProviderForReloadTest) GetDefaultModel() string { + panic("boom") +} + func TestCloseRuntimeEventLoggerSubscriptionWaitsForDrain(t *testing.T) { eventBus := runtimeevents.NewBus() defer func() { diff --git a/pkg/agent/turn_state.go b/pkg/agent/turn_state.go index ddd1eb89..2161490e 100644 --- a/pkg/agent/turn_state.go +++ b/pkg/agent/turn_state.go @@ -285,7 +285,17 @@ func (al *AgentLoop) registerActiveTurn(ts *turnState) { } func (al *AgentLoop) clearActiveTurn(ts *turnState) { - al.activeTurnStates.Delete(ts.sessionKey) + al.releaseSessionTurnState(ts.sessionKey, ts) +} + +func (al *AgentLoop) releaseSessionTurnState(sessionKey string, expected *turnState) { + if expected == nil { + al.activeTurnStates.Delete(sessionKey) + return + } + if actual, ok := al.activeTurnStates.Load(sessionKey); ok && actual == expected { + al.activeTurnStates.Delete(sessionKey) + } } func (al *AgentLoop) getActiveTurnState(sessionKey string) *turnState { From a7208983d04c41b7e574932b86c6261492fd1e8c Mon Sep 17 00:00:00 2001 From: SiYue-ZO <2835601846@qq.com> Date: Thu, 11 Jun 2026 13:55:56 +0800 Subject: [PATCH 2/3] feat(bus): add backpressure drop budget for audio streams only --- pkg/bus/bus.go | 225 +++++++++++++++++++++++++++++++--- pkg/bus/bus_test.go | 147 ++++++++++++++++++++++ pkg/bus/events.go | 51 ++++++++ pkg/channels/discord/voice.go | 8 +- pkg/events/kind.go | 4 + 5 files changed, 414 insertions(+), 21 deletions(-) diff --git a/pkg/bus/bus.go b/pkg/bus/bus.go index d076bd27..fc7ecb86 100644 --- a/pkg/bus/bus.go +++ b/pkg/bus/bus.go @@ -3,8 +3,10 @@ package bus import ( "context" "errors" + "fmt" "sync" "sync/atomic" + "time" runtimeevents "github.com/sipeed/picoclaw/pkg/events" "github.com/sipeed/picoclaw/pkg/logger" @@ -13,6 +15,10 @@ import ( // ErrBusClosed is returned when publishing to a closed MessageBus. var ErrBusClosed = errors.New("message bus closed") +// ErrBusBackpressure is returned when a publish attempt exceeds the configured +// backpressure wait budget and the message is dropped. +var ErrBusBackpressure = errors.New("message bus backpressure") + var ( ErrMissingInboundContext = errors.New("inbound message context is required") ErrMissingOutboundContext = errors.New("outbound message context is required") @@ -21,6 +27,41 @@ var ( const defaultBusBufferSize = 64 +const ( + defaultAudioPublishTimeout = 150 * time.Millisecond +) + +type publishPolicy struct { + stream string + // timeout is the backpressure drop budget. When positive, a full channel + // causes the message to be dropped after this duration. When zero, the + // publish blocks until context cancellation or bus close (no drop). + timeout time.Duration +} + +type streamStats struct { + dropped atomic.Uint64 + lastDropped atomic.Int64 + lastWaitNanos atomic.Int64 +} + +type MessageBusStats struct { + Inbound StreamStats `json:"inbound"` + Outbound StreamStats `json:"outbound"` + OutboundMedia StreamStats `json:"outbound_media"` + AudioChunks StreamStats `json:"audio_chunks"` + VoiceControls StreamStats `json:"voice_controls"` +} + +type StreamStats struct { + Depth int `json:"depth"` + Capacity int `json:"capacity"` + DroppedTotal uint64 `json:"dropped_total"` + LastDroppedAt time.Time `json:"last_dropped_at,omitempty"` + LastDropWait string `json:"last_drop_wait,omitempty"` + LastDropWaitMillis int64 `json:"last_drop_wait_ms,omitempty"` +} + // StreamDelegate is implemented by the channel Manager to provide streaming // capabilities to the agent loop without tight coupling. type StreamDelegate interface { @@ -62,8 +103,14 @@ type MessageBus struct { done chan struct{} closed atomic.Bool wg sync.WaitGroup + publishMu sync.Mutex streamDelegate atomic.Value // stores StreamDelegate eventPublisher atomic.Value // stores EventPublisher + inboundStats streamStats + outboundStats streamStats + mediaStats streamStats + audioStats streamStats + voiceStats streamStats } // EventPublisher is the minimal runtime event publisher used by MessageBus. @@ -83,13 +130,13 @@ func NewMessageBus() *MessageBus { } } -func publish[T any](ctx context.Context, mb *MessageBus, ch chan T, msg T) error { - // check bus closed before acquiring wg, to avoid unnecessary wg.Add and potential deadlock +func (mb *MessageBus) enterPublish(ctx context.Context) error { + mb.publishMu.Lock() + defer mb.publishMu.Unlock() + if mb.closed.Load() { return ErrBusClosed } - - // check again,before sending message, to avoid sending to closed channel select { case <-ctx.Done(): return ctx.Err() @@ -99,8 +146,63 @@ func publish[T any](ctx context.Context, mb *MessageBus, ch chan T, msg T) error } mb.wg.Add(1) + return nil +} + +func publish[T any]( + ctx context.Context, + mb *MessageBus, + ch chan T, + msg T, + policy publishPolicy, + stats *streamStats, + scope runtimeevents.Scope, +) error { + if ctx == nil { + ctx = context.Background() + } + if err := mb.enterPublish(ctx); err != nil { + return err + } defer mb.wg.Done() + // timeout == 0 means no backpressure drop budget; block until context + // cancellation or bus close. This is the default for critical streams + // (inbound, outbound, outboundMedia, voiceControl) where dropping + // messages silently is undesirable. + if policy.timeout > 0 { + timer := time.NewTimer(policy.timeout) + defer timer.Stop() + + select { + case ch <- msg: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + droppedTotal := stats.dropped.Add(1) + now := time.Now() + stats.lastDropped.Store(now.UnixNano()) + stats.lastWaitNanos.Store(policy.timeout.Nanoseconds()) + queueDepth := len(ch) + queueCap := cap(ch) + mb.publishDrop( + policy.stream, scope, "queue_full_timeout", + policy.timeout, queueDepth, queueCap, droppedTotal, + ) + logger.WarnCF("bus", "Dropped bus message due to backpressure", map[string]any{ + "stream": policy.stream, + "wait_ms": policy.timeout.Milliseconds(), + "queue_depth": queueDepth, + "queue_capacity": queueCap, + "dropped_total": droppedTotal, + }) + return fmt.Errorf("%w: %s queue full after %s", ErrBusBackpressure, policy.stream, policy.timeout) + case <-mb.done: + return ErrBusClosed + } + } + select { case ch <- msg: return nil @@ -117,8 +219,13 @@ func (mb *MessageBus) PublishInbound(ctx context.Context, msg InboundMessage) er mb.publishFailure("inbound", runtimeScopeFromInboundContext(msg.Context), ErrMissingInboundContext) return ErrMissingInboundContext } - if err := publish(ctx, mb, mb.inbound, msg); err != nil { - mb.publishFailure("inbound", runtimeScopeFromInboundContext(msg.Context), err) + if err := publish(ctx, mb, mb.inbound, msg, publishPolicy{ + stream: "inbound", + }, &mb.inboundStats, runtimeScopeFromInboundContext(msg.Context)); err != nil { + scope := runtimeScopeFromInboundContext(msg.Context) + if !errors.Is(err, ErrBusBackpressure) { + mb.publishFailure("inbound", scope, err) + } return err } return nil @@ -134,8 +241,13 @@ func (mb *MessageBus) PublishOutbound(ctx context.Context, msg OutboundMessage) mb.publishFailure("outbound", runtimeScopeFromInboundContext(msg.Context), ErrMissingOutboundContext) return ErrMissingOutboundContext } - if err := publish(ctx, mb, mb.outbound, msg); err != nil { - mb.publishFailure("outbound", runtimeScopeFromInboundContext(msg.Context), err) + if err := publish(ctx, mb, mb.outbound, msg, publishPolicy{ + stream: "outbound", + }, &mb.outboundStats, runtimeScopeFromInboundContext(msg.Context)); err != nil { + scope := runtimeScopeFromInboundContext(msg.Context) + if !errors.Is(err, ErrBusBackpressure) { + mb.publishFailure("outbound", scope, err) + } return err } return nil @@ -151,8 +263,13 @@ func (mb *MessageBus) PublishOutboundMedia(ctx context.Context, msg OutboundMedi mb.publishFailure("outbound_media", runtimeScopeFromInboundContext(msg.Context), ErrMissingOutboundMediaContext) return ErrMissingOutboundMediaContext } - if err := publish(ctx, mb, mb.outboundMedia, msg); err != nil { - mb.publishFailure("outbound_media", runtimeScopeFromInboundContext(msg.Context), err) + if err := publish(ctx, mb, mb.outboundMedia, msg, publishPolicy{ + stream: "outbound_media", + }, &mb.mediaStats, runtimeScopeFromInboundContext(msg.Context)); err != nil { + scope := runtimeScopeFromInboundContext(msg.Context) + if !errors.Is(err, ErrBusBackpressure) { + mb.publishFailure("outbound_media", scope, err) + } return err } return nil @@ -163,8 +280,14 @@ func (mb *MessageBus) OutboundMediaChan() <-chan OutboundMediaMessage { } func (mb *MessageBus) PublishAudioChunk(ctx context.Context, chunk AudioChunk) error { - if err := publish(ctx, mb, mb.audioChunks, chunk); err != nil { - mb.publishFailure("audio_chunk", runtimeScopeFromAudioChunk(chunk), err) + if err := publish(ctx, mb, mb.audioChunks, chunk, publishPolicy{ + stream: "audio_chunk", + timeout: defaultAudioPublishTimeout, + }, &mb.audioStats, runtimeScopeFromAudioChunk(chunk)); err != nil { + scope := runtimeScopeFromAudioChunk(chunk) + if !errors.Is(err, ErrBusBackpressure) { + mb.publishFailure("audio_chunk", scope, err) + } return err } return nil @@ -175,8 +298,13 @@ func (mb *MessageBus) AudioChunksChan() <-chan AudioChunk { } func (mb *MessageBus) PublishVoiceControl(ctx context.Context, ctrl VoiceControl) error { - if err := publish(ctx, mb, mb.voiceControls, ctrl); err != nil { - mb.publishFailure("voice_control", runtimeScopeFromVoiceControl(ctrl), err) + if err := publish(ctx, mb, mb.voiceControls, ctrl, publishPolicy{ + stream: "voice_control", + }, &mb.voiceStats, runtimeScopeFromVoiceControl(ctrl)); err != nil { + scope := runtimeScopeFromVoiceControl(ctrl) + if !errors.Is(err, ErrBusBackpressure) { + mb.publishFailure("voice_control", scope, err) + } return err } return nil @@ -204,15 +332,76 @@ func (mb *MessageBus) GetStreamer(ctx context.Context, channel, chatID, sessionK return nil, false } +func (mb *MessageBus) Stats() MessageBusStats { + if mb == nil { + return MessageBusStats{} + } + return MessageBusStats{ + Inbound: snapshotStreamStats(mb.inbound, &mb.inboundStats), + Outbound: snapshotStreamStats(mb.outbound, &mb.outboundStats), + OutboundMedia: snapshotStreamStats(mb.outboundMedia, &mb.mediaStats), + AudioChunks: snapshotStreamStats(mb.audioChunks, &mb.audioStats), + VoiceControls: snapshotStreamStats(mb.voiceControls, &mb.voiceStats), + } +} + +// HealthCheck returns a snapshot of queue depths and cumulative drop counts +// across all streams. It always reports ok=true: backpressure-induced drops are +// reflected in the message string for telemetry but do not affect the boolean. +// Callers that need readiness semantics (e.g. returning 503 when drops occur) +// should inspect Stats() directly and apply their own threshold logic. +func (mb *MessageBus) HealthCheck() (bool, string) { + stats := mb.Stats() + totalDropped := stats.Inbound.DroppedTotal + + stats.Outbound.DroppedTotal + + stats.OutboundMedia.DroppedTotal + + stats.AudioChunks.DroppedTotal + + stats.VoiceControls.DroppedTotal + message := fmt.Sprintf( + "in=%d/%d out=%d/%d media=%d/%d audio=%d/%d voice=%d/%d dropped=%d", + stats.Inbound.Depth, + stats.Inbound.Capacity, + stats.Outbound.Depth, + stats.Outbound.Capacity, + stats.OutboundMedia.Depth, + stats.OutboundMedia.Capacity, + stats.AudioChunks.Depth, + stats.AudioChunks.Capacity, + stats.VoiceControls.Depth, + stats.VoiceControls.Capacity, + totalDropped, + ) + + return true, message +} + +func snapshotStreamStats[T any](ch chan T, stats *streamStats) StreamStats { + snapshot := StreamStats{ + DroppedTotal: stats.dropped.Load(), + } + if ch != nil { + snapshot.Depth = len(ch) + snapshot.Capacity = cap(ch) + } + if unixNano := stats.lastDropped.Load(); unixNano > 0 { + snapshot.LastDroppedAt = time.Unix(0, unixNano) + } + if waitNanos := stats.lastWaitNanos.Load(); waitNanos > 0 { + wait := time.Duration(waitNanos) + snapshot.LastDropWait = wait.String() + snapshot.LastDropWaitMillis = wait.Milliseconds() + } + return snapshot +} + func (mb *MessageBus) Close() { mb.closeOnce.Do(func() { mb.publishCloseEvent(runtimeevents.KindBusCloseStarted, 0) - // notify all blocked publishers to exit - close(mb.done) - // because every publisher will check mb.closed before acquiring wg - // so we can be sure that new publishers will not be added new messages after this point + mb.publishMu.Lock() mb.closed.Store(true) + close(mb.done) + mb.publishMu.Unlock() // wait for all ongoing Publish calls to finish, ensuring all messages have been sent to channels or exited mb.wg.Wait() diff --git a/pkg/bus/bus_test.go b/pkg/bus/bus_test.go index a0a9e1e1..786d97fc 100644 --- a/pkg/bus/bus_test.go +++ b/pkg/bus/bus_test.go @@ -2,6 +2,7 @@ package bus import ( "context" + "errors" "sync" "testing" "time" @@ -459,6 +460,84 @@ func TestPublishAudioChunkSubscribe(t *testing.T) { } } +func TestPublishAudioChunk_BackpressureDropPublishesRuntimeEvent(t *testing.T) { + eventBus := runtimeevents.NewBus() + defer func() { + if err := eventBus.Close(); err != nil { + t.Errorf("event bus close failed: %v", err) + } + }() + + _, eventsCh, err := eventBus.Channel().OfKind(runtimeevents.KindBusMessageDropped).SubscribeChan( + t.Context(), + runtimeevents.SubscribeOptions{Name: "bus-drop-events", Buffer: 1}, + ) + if err != nil { + t.Fatalf("SubscribeChan failed: %v", err) + } + + mb := NewMessageBus() + defer mb.Close() + mb.SetEventPublisher(eventBus) + + for i := range defaultBusBufferSize * 4 { + if pubErr := mb.PublishAudioChunk(context.Background(), AudioChunk{ + SessionID: "voice-1", + SpeakerID: "speaker-1", + ChatID: "chat-1", + Channel: "discord", + Sequence: uint64(i), + Format: "opus", + Data: []byte{0x01}, + }); pubErr != nil { + t.Fatalf("fill failed at %d: %v", i, pubErr) + } + } + + err = mb.PublishAudioChunk(context.Background(), AudioChunk{ + SessionID: "voice-1", + SpeakerID: "speaker-1", + ChatID: "chat-1", + Channel: "discord", + Sequence: 999, + Format: "opus", + Data: []byte{0x01}, + }) + if !errors.Is(err, ErrBusBackpressure) { + t.Fatalf("PublishAudioChunk() error = %v, want %v", err, ErrBusBackpressure) + } + + evt := receiveBusRuntimeEvent(t, eventsCh) + if evt.Kind != runtimeevents.KindBusMessageDropped || + evt.Source.Name != "audio_chunk" || + evt.Severity != runtimeevents.SeverityWarn { + t.Fatalf("drop event = %+v", evt) + } + if evt.Scope.Channel != "discord" || evt.Scope.ChatID != "chat-1" { + t.Fatalf("drop event scope = %+v", evt.Scope) + } + if evt.Attrs["stream"] != "audio_chunk" || + evt.Attrs["reason"] != "queue_full_timeout" || + evt.Attrs["wait_ms"] != defaultAudioPublishTimeout.Milliseconds() || + evt.Attrs["queue_depth"] != defaultBusBufferSize*4 || + evt.Attrs["queue_capacity"] != defaultBusBufferSize*4 || + evt.Attrs["dropped_total"] != uint64(1) { + t.Fatalf("drop event attrs = %#v", evt.Attrs) + } + + stats := mb.Stats() + if stats.AudioChunks.DroppedTotal != 1 { + t.Fatalf("AudioChunks dropped = %d, want 1", stats.AudioChunks.DroppedTotal) + } + if stats.AudioChunks.Depth != defaultBusBufferSize*4 { + t.Fatalf("AudioChunks depth = %d, want %d", stats.AudioChunks.Depth, defaultBusBufferSize*4) + } + wantWaitMS := defaultAudioPublishTimeout.Milliseconds() + if stats.AudioChunks.LastDropWaitMillis != wantWaitMS { + t.Fatalf("AudioChunks last wait ms = %d, want %d", stats.AudioChunks.LastDropWaitMillis, wantWaitMS) + } +} + func TestPublishVoiceControlSubscribe(t *testing.T) { mb := NewMessageBus() defer mb.Close() @@ -728,6 +807,74 @@ func TestPublishInbound_FullBuffer(t *testing.T) { } } +// TestPublishInbound_FullBufferUsesBusBackpressureBudget exercises the generic +// publish() backpressure path directly (with a short 20ms timeout) rather than +// going through PublishInbound(). This avoids waiting for a long context timeout +// and keeps the test fast. Context validation and public-API wiring are covered +// by TestPublishInbound_FullBuffer and TestPublishInbound_ContextCancel. +func TestPublishInbound_FullBufferUsesBusBackpressureBudget(t *testing.T) { + mb := NewMessageBus() + defer mb.Close() + + ch := make(chan InboundMessage, 1) + ch <- InboundMessage{Content: "fill"} + + scope := runtimeevents.Scope{Channel: "test", ChatID: "chat-overflow"} + err := publish(context.Background(), mb, ch, InboundMessage{Content: "overflow"}, publishPolicy{ + stream: "inbound", + timeout: 20 * time.Millisecond, + }, &mb.inboundStats, scope) + if !errors.Is(err, ErrBusBackpressure) { + t.Fatalf("publish() error = %v, want %v", err, ErrBusBackpressure) + } + + stats := mb.Stats() + if stats.Inbound.DroppedTotal != 1 { + t.Fatalf("Inbound dropped = %d, want 1", stats.Inbound.DroppedTotal) + } + if stats.Inbound.LastDropWaitMillis != 20 { + t.Fatalf("Inbound last wait ms = %d, want 20", stats.Inbound.LastDropWaitMillis) + } +} + +func TestMessageBusHealthCheckIncludesQueueDepthAndDrops(t *testing.T) { + mb := NewMessageBus() + defer mb.Close() + + ok, msg := mb.HealthCheck() + if !ok { + t.Fatal("HealthCheck should remain ok for backpressure telemetry") + } + if msg == "" { + t.Fatal("HealthCheck message should not be empty") + } + + for i := range cap(mb.audioChunks) { + if err := mb.PublishAudioChunk(context.Background(), AudioChunk{ + Channel: "discord", + ChatID: "voice-room", + Sequence: uint64(i), + Data: []byte("fill"), + }); err != nil { + t.Fatalf("fill audio buffer at %d: %v", i, err) + } + } + _ = mb.PublishAudioChunk(context.Background(), AudioChunk{ + Channel: "discord", + ChatID: "voice-room", + Sequence: 999, + Data: []byte("overflow"), + }) + + stats := mb.Stats() + if stats.AudioChunks.Depth != cap(mb.audioChunks) { + t.Fatalf("audio depth = %d, want %d", stats.AudioChunks.Depth, cap(mb.audioChunks)) + } + if stats.AudioChunks.DroppedTotal != 1 { + t.Fatalf("audio dropped = %d, want 1", stats.AudioChunks.DroppedTotal) + } +} + func TestCloseIdempotent(t *testing.T) { mb := NewMessageBus() diff --git a/pkg/bus/events.go b/pkg/bus/events.go index 4640ed1f..787a7f64 100644 --- a/pkg/bus/events.go +++ b/pkg/bus/events.go @@ -1,6 +1,8 @@ package bus import ( + "time" + runtimeevents "github.com/sipeed/picoclaw/pkg/events" ) @@ -13,6 +15,15 @@ type busClosePayload struct { Drained int `json:"drained,omitempty"` } +type busMessageDroppedPayload struct { + Stream string `json:"stream"` + Reason string `json:"reason"` + WaitMS int64 `json:"wait_ms"` + QueueDepth int `json:"queue_depth"` + QueueCap int `json:"queue_capacity"` + DroppedTotal uint64 `json:"dropped_total"` +} + func (mb *MessageBus) publishFailure(stream string, scope runtimeevents.Scope, err error) { if mb == nil || err == nil { return @@ -38,6 +49,46 @@ func (mb *MessageBus) publishFailure(stream string, scope runtimeevents.Scope, e }) } +func (mb *MessageBus) publishDrop( + stream string, + scope runtimeevents.Scope, + reason string, + wait time.Duration, + queueDepth, queueCap int, + droppedTotal uint64, +) { + if mb == nil { + return + } + publisher, ok := mb.eventPublisher.Load().(EventPublisher) + if !ok || publisher == nil { + return + } + + publisher.PublishNonBlocking(runtimeevents.Event{ + Kind: runtimeevents.KindBusMessageDropped, + Source: runtimeevents.Source{Component: "bus", Name: stream}, + Scope: scope, + Severity: runtimeevents.SeverityWarn, + Payload: busMessageDroppedPayload{ + Stream: stream, + Reason: reason, + WaitMS: wait.Milliseconds(), + QueueDepth: queueDepth, + QueueCap: queueCap, + DroppedTotal: droppedTotal, + }, + Attrs: map[string]any{ + "stream": stream, + "reason": reason, + "wait_ms": wait.Milliseconds(), + "queue_depth": queueDepth, + "queue_capacity": queueCap, + "dropped_total": droppedTotal, + }, + }) +} + func (mb *MessageBus) publishCloseEvent(kind runtimeevents.Kind, drained int) { if mb == nil { return diff --git a/pkg/channels/discord/voice.go b/pkg/channels/discord/voice.go index 554b8ae7..50ed1ad3 100644 --- a/pkg/channels/discord/voice.go +++ b/pkg/channels/discord/voice.go @@ -298,9 +298,11 @@ func (c *DiscordChannel) receiveVoice(vc *discordgo.VoiceConnection, guildID str Data: p.Opus, } - ctx, cancel := context.WithTimeout(c.ctx, 100*time.Millisecond) - err := c.bus.PublishAudioChunk(ctx, chunk) - cancel() + // Pass the parent context directly; the bus applies its own + // audio drop budget internally. + // A 100ms caller-level timeout would fire before the bus's + // backpressure handling, masking drop counters and events. + err := c.bus.PublishAudioChunk(c.ctx, chunk) if err != nil { logger.ErrorCF("discord", "Failed to publish audio chunk", map[string]any{ "guild": guildID, diff --git a/pkg/events/kind.go b/pkg/events/kind.go index b9327e15..e12c4e62 100644 --- a/pkg/events/kind.go +++ b/pkg/events/kind.go @@ -68,6 +68,9 @@ const ( // KindBusPublishFailed is emitted when message bus publish fails. KindBusPublishFailed Kind = "bus.publish.failed" + // KindBusMessageDropped is emitted when a message is dropped due to + // backpressure (channel full for longer than the drop budget). + KindBusMessageDropped Kind = "bus.message.dropped" // KindBusCloseStarted is emitted when message bus close starts. KindBusCloseStarted Kind = "bus.close.started" // KindBusCloseCompleted is emitted when message bus close completes. @@ -133,6 +136,7 @@ var knownKinds = []Kind{ KindChannelMessageOutboundFailed, KindChannelRateLimited, KindBusPublishFailed, + KindBusMessageDropped, KindBusCloseStarted, KindBusCloseCompleted, KindBusCloseDrained, From 8290c72d44b18a336daf4c4f27b6b3d092d4ce97 Mon Sep 17 00:00:00 2001 From: SiYue-ZO <2835601846@qq.com> Date: Mon, 15 Jun 2026 01:32:13 +0800 Subject: [PATCH 3/3] fix(agent): replace WaitGroup with Cond-based counter and conditionalize panic cleanup Two concurrency bugs identified in PR #2904 review: 1. Replace sync.WaitGroup with sync.Cond-based activeReqCount to avoid the "WaitGroup is reused before previous Wait has returned" panic that occurs when Add(1) races with a goroutine-launched Wait(). 2. Make panic cleanup conditional: when runTurn panics, only delete the session's activeTurnStates entry if it still points to our placeholder. Previously, an unconditional delete could wipe a new message's slot claimed between the panic and the deferred cleanup. --- pkg/agent/agent.go | 23 +++++++-- pkg/agent/agent_init.go | 2 + pkg/agent/agent_utils.go | 66 +++++++++++++++++--------- pkg/agent/context_legacy.go | 4 +- pkg/agent/pipeline_llm.go | 4 +- pkg/agent/runtime_event_logger_test.go | 6 ++- 6 files changed, 72 insertions(+), 33 deletions(-) diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 2723831f..0ba8c370 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -69,8 +69,14 @@ type AgentLoop struct { activeTurnStates sync.Map subTurnCounter atomic.Int64 - turnSeq atomic.Uint64 - activeRequests sync.WaitGroup + turnSeq atomic.Uint64 + + // activeReqMu/activeReqCond/activeReqCount replace sync.WaitGroup to + // avoid the "WaitGroup is reused before previous Wait has returned" panic + // that occurs when Add(1) races with a goroutine-launched Wait(). + activeReqMu sync.Mutex + activeReqCond *sync.Cond + activeReqCount int reloadFunc func() error @@ -208,7 +214,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { // Session claimed — spawn a worker goroutine that acquires a semaphore // slot. The goroutine is spawned immediately so the main loop keeps // draining the inbound channel. The goroutine blocks on the semaphore. - go func(m bus.InboundMessage) { + go func(m bus.InboundMessage, ph *turnState) { var releaseSession bool // Acquire semaphore slot (blocks if at capacity) select { @@ -227,7 +233,14 @@ func (al *AgentLoop) Run(ctx context.Context) error { // this becomes a no-op (the key is already gone). defer func() { if releaseSession { - al.releaseSessionTurnState(sessionKey, nil) + // Conditional delete: only remove the entry if it still points + // to our placeholder. A new message may have claimed the slot + // between the panic and this defer. + if actual, ok := al.activeTurnStates.Load(sessionKey); ok { + if ts, ok := actual.(*turnState); ok && ts == ph { + al.releaseSessionTurnState(sessionKey, ts) + } + } return } if actual, ok := al.activeTurnStates.Load(sessionKey); ok { @@ -276,7 +289,7 @@ func (al *AgentLoop) Run(ctx context.Context) error { } al.runTurnWithSteering(ctx, m) - }(msg) + }(msg, placeholder) // TODO: Re-enable media cleanup after inbound media is properly consumed by the agent. // Currently disabled because files are deleted before the LLM can access their content. diff --git a/pkg/agent/agent_init.go b/pkg/agent/agent_init.go index 50f0227a..b8a7c7b0 100644 --- a/pkg/agent/agent_init.go +++ b/pkg/agent/agent_init.go @@ -5,6 +5,7 @@ package agent import ( "context" "fmt" + "sync" "time" "github.com/sipeed/picoclaw/pkg/agent/interfaces" @@ -91,6 +92,7 @@ func NewAgentLoop( }) } } + al.activeReqCond = sync.NewCond(&al.activeReqMu) al.refreshRuntimeEventLogger(cfg) al.providerFactory = providers.CreateProviderFromConfig al.hooks = NewHookManager(al.runtimeEvents.Channel()) diff --git a/pkg/agent/agent_utils.go b/pkg/agent/agent_utils.go index e1da3bf6..11fbfb40 100644 --- a/pkg/agent/agent_utils.go +++ b/pkg/agent/agent_utils.go @@ -585,33 +585,55 @@ func closeProviderIfStateful(provider providers.LLMProvider) { } } +// activeRequestsInc atomically increments the active request count. +func (al *AgentLoop) activeRequestsInc() { + al.activeReqMu.Lock() + al.activeReqCount++ + al.activeReqMu.Unlock() +} + +// activeRequestsDec atomically decrements the active request count +// and wakes any goroutine blocked in waitForActiveRequests when the +// count reaches zero. +func (al *AgentLoop) activeRequestsDec() { + al.activeReqMu.Lock() + al.activeReqCount-- + if al.activeReqCount == 0 { + al.activeReqCond.Broadcast() + } + al.activeReqMu.Unlock() +} + func (al *AgentLoop) waitForActiveRequests(ctx context.Context, timeout time.Duration) bool { - done := make(chan struct{}) + al.activeReqMu.Lock() + if al.activeReqCount == 0 { + al.activeReqMu.Unlock() + return true + } + + // Wake blocked Wait() callers on timeout or context cancellation. + var timedOut bool + if timeout > 0 { + time.AfterFunc(timeout, func() { + al.activeReqMu.Lock() + timedOut = true + al.activeReqCond.Broadcast() + al.activeReqMu.Unlock() + }) + } go func() { - al.activeRequests.Wait() - close(done) + <-ctx.Done() + al.activeReqMu.Lock() + al.activeReqCond.Broadcast() + al.activeReqMu.Unlock() }() - if timeout <= 0 { - select { - case <-done: - return true - case <-ctx.Done(): - return false - } - } - - timer := time.NewTimer(timeout) - defer timer.Stop() - - select { - case <-done: - return true - case <-timer.C: - return false - case <-ctx.Done(): - return false + for al.activeReqCount > 0 && !timedOut && ctx.Err() == nil { + al.activeReqCond.Wait() } + result := al.activeReqCount == 0 + al.activeReqMu.Unlock() + return result } func (al *AgentLoop) closeReloadedProvider(ctx context.Context, provider providers.StatefulProvider) { diff --git a/pkg/agent/context_legacy.go b/pkg/agent/context_legacy.go index 94ef5367..908c3d85 100644 --- a/pkg/agent/context_legacy.go +++ b/pkg/agent/context_legacy.go @@ -294,9 +294,9 @@ func (m *legacyContextManager) retryLLMCall( var err error for attempt := 0; attempt < maxRetries; attempt++ { - m.al.activeRequests.Add(1) + m.al.activeRequestsInc() resp, err = func() (*providers.LLMResponse, error) { - defer m.al.activeRequests.Done() + defer m.al.activeRequestsDec() return agent.Provider.Chat( ctx, []providers.Message{{Role: "user", Content: prompt}}, diff --git a/pkg/agent/pipeline_llm.go b/pkg/agent/pipeline_llm.go index aaae765e..5950c4bf 100644 --- a/pkg/agent/pipeline_llm.go +++ b/pkg/agent/pipeline_llm.go @@ -158,8 +158,8 @@ func (p *Pipeline) CallLLM( ts.clearProviderCancel(providerCancel) }() - al.activeRequests.Add(1) - defer al.activeRequests.Done() + al.activeRequestsInc() + defer al.activeRequestsDec() if response, handled, streamErr := p.tryConfiguredStreamingLLM( providerCtx, diff --git a/pkg/agent/runtime_event_logger_test.go b/pkg/agent/runtime_event_logger_test.go index 1da9803c..479257e3 100644 --- a/pkg/agent/runtime_event_logger_test.go +++ b/pkg/agent/runtime_event_logger_test.go @@ -3,6 +3,7 @@ package agent import ( "context" "errors" + "sync" "sync/atomic" "testing" "time" @@ -307,8 +308,9 @@ func TestReloadProviderAndConfigWaitsForInFlightRequestsBeforeClosingOldProvider func TestWaitForActiveRequestsHonorsContextCancellation(t *testing.T) { al := &AgentLoop{} - al.activeRequests.Add(1) - defer al.activeRequests.Done() + al.activeReqCond = sync.NewCond(&al.activeReqMu) + al.activeRequestsInc() + defer al.activeRequestsDec() ctx, cancel := context.WithCancel(context.Background()) cancel()