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.
This commit is contained in:
SiYue-ZO
2026-06-15 01:32:13 +08:00
parent a7208983d0
commit 8290c72d44
6 changed files with 72 additions and 33 deletions
+18 -5
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
@@ -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.
+2
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())
+44 -22
View File
@@ -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) {
+2 -2
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}},
+2 -2
View File
@@ -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,
+4 -2
View File
@@ -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()