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