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

Merge pull request #2904 from SiYue-ZO/feature/fix-agent-loop-stability

Browse source 
Fix agent loop reload and panic cleanup stability
This commit is contained in:
Mauro 2026-06-14 21:22:53 +02:00 committed by GitHub
commit 13a38bd1c4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 804 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

72
pkg/agent/agent.go
View file

@ -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
@ -118,6 +124,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,7 +215,8 @@ 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 {
case al.workerSem <- struct{}{}:
@ -216,7 +224,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
}
@ -225,16 +233,28 @@ 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 {
// 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 {
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{
@ -252,7 +272,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,
@ -270,7 +290,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.
@ -366,37 +386,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
@ -472,17 +478,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)
}
}

2
pkg/agent/agent_init.go
View file

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

106
pkg/agent/agent_test.go
View file

@ -14,6 +14,7 @@ import (
"slices"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
@ -120,6 +121,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,
@ -7492,3 +7518,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)
}
}

72
pkg/agent/agent_utils.go
View file

@ -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"
@ -591,6 +592,77 @@ 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 {
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() {
<-ctx.Done()
al.activeReqMu.Lock()
al.activeReqCond.Broadcast()
al.activeReqMu.Unlock()
}()
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) {
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)
}

4
pkg/agent/context_legacy.go
View file

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

4
pkg/agent/pipeline_llm.go
View file

@ -160,8 +160,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,

161
pkg/agent/runtime_event_logger_test.go
View file

@ -2,6 +2,8 @@ package agent
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
@ -9,6 +11,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 +194,164 @@ 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.activeReqCond = sync.NewCond(&al.activeReqMu)
al.activeRequestsInc()
defer al.activeRequestsDec()
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() {

12
pkg/agent/turn_state.go
View file

@ -290,7 +290,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 {

225
pkg/bus/bus.go
View file

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

147
pkg/bus/bus_test.go
View file

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

51
pkg/bus/events.go
View file

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

8
pkg/channels/discord/voice.go
View file

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

4
pkg/events/kind.go
View file

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