fix(agent): resolve critical race conditions and resource leaks in SubTurn
- Fix turnState hierarchy corruption when SubTurns recursively call runAgentLoop by checking context for existing turnState before creating new root - Fix deadlock risk in deliverSubTurnResult by separating lock and channel operations - Fix session rollback race in HardAbort by calling Finish() before rollback - Fix resource leak by closing pendingResults channel in Finish() with panic recovery - Add thread-safety documentation for childTurnIDs and isFinished fields - Move globalTurnCounter to AgentLoop.subTurnCounter to prevent ID conflicts - Improve semaphore acquisition to ensure release even on early validation failures - Document design choice: ephemeral sessions start empty for complete isolation - Add 5 new tests: hierarchy, deadlock, order, channel close, and semaphore
This commit is contained in:
parent
9d761b7f5b
commit
6b5d7e3fd7
5 changed files with 347 additions and 67 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -56,3 +56,5 @@ dist/
|
||||||
!web/backend/dist/
|
!web/backend/dist/
|
||||||
web/backend/dist/*
|
web/backend/dist/*
|
||||||
!web/backend/dist/.gitkeep
|
!web/backend/dist/.gitkeep
|
||||||
|
|
||||||
|
.claude/
|
||||||
|
|
@ -36,21 +36,22 @@ import (
|
||||||
)
|
)
|
||||||
|
|
||||||
type AgentLoop struct {
|
type AgentLoop struct {
|
||||||
bus *bus.MessageBus
|
bus *bus.MessageBus
|
||||||
cfg *config.Config
|
cfg *config.Config
|
||||||
registry *AgentRegistry
|
registry *AgentRegistry
|
||||||
state *state.Manager
|
state *state.Manager
|
||||||
running atomic.Bool
|
running atomic.Bool
|
||||||
summarizing sync.Map
|
summarizing sync.Map
|
||||||
fallback *providers.FallbackChain
|
fallback *providers.FallbackChain
|
||||||
channelManager *channels.Manager
|
channelManager *channels.Manager
|
||||||
mediaStore media.MediaStore
|
mediaStore media.MediaStore
|
||||||
transcriber voice.Transcriber
|
transcriber voice.Transcriber
|
||||||
cmdRegistry *commands.Registry
|
cmdRegistry *commands.Registry
|
||||||
mcp mcpRuntime
|
mcp mcpRuntime
|
||||||
steering *steeringQueue
|
steering *steeringQueue
|
||||||
subTurnResults sync.Map // key: sessionKey (string), value: chan *tools.ToolResult
|
subTurnResults sync.Map // key: sessionKey (string), value: chan *tools.ToolResult
|
||||||
activeTurnStates sync.Map // key: sessionKey (string), value: *turnState
|
activeTurnStates sync.Map // key: sessionKey (string), value: *turnState
|
||||||
|
subTurnCounter atomic.Int64 // Counter for generating unique SubTurn IDs
|
||||||
mu sync.RWMutex
|
mu sync.RWMutex
|
||||||
// Track active requests for safe provider cleanup
|
// Track active requests for safe provider cleanup
|
||||||
activeRequests sync.WaitGroup
|
activeRequests sync.WaitGroup
|
||||||
|
|
@ -964,25 +965,39 @@ func (al *AgentLoop) runAgentLoop(
|
||||||
agent *AgentInstance,
|
agent *AgentInstance,
|
||||||
opts processOptions,
|
opts processOptions,
|
||||||
) (string, error) {
|
) (string, error) {
|
||||||
// Initialize a root TurnState for this iteration, allowing sub-turns to be spawned.
|
// Check if we're already inside a SubTurn (context already has a turnState).
|
||||||
rootTS := &turnState{
|
// If so, reuse it instead of creating a new root turnState.
|
||||||
ctx: ctx,
|
// This prevents turnState hierarchy corruption when SubTurns recursively call runAgentLoop.
|
||||||
turnID: opts.SessionKey, // Associate this turn graph with the current session key
|
existingTS := turnStateFromContext(ctx)
|
||||||
depth: 0,
|
var rootTS *turnState
|
||||||
session: agent.Sessions,
|
var isRootTurn bool
|
||||||
initialHistoryLength: len(agent.Sessions.GetHistory("")), // Snapshot for rollback on hard abort
|
|
||||||
pendingResults: make(chan *tools.ToolResult, 16),
|
if existingTS != nil {
|
||||||
concurrencySem: make(chan struct{}, 5), // maxConcurrentSubTurns
|
// We're inside a SubTurn — reuse the existing turnState
|
||||||
|
rootTS = existingTS
|
||||||
|
isRootTurn = false
|
||||||
|
} else {
|
||||||
|
// This is a top-level turn — initialize a new root TurnState
|
||||||
|
rootTS = &turnState{
|
||||||
|
ctx: ctx,
|
||||||
|
turnID: opts.SessionKey, // Associate this turn graph with the current session key
|
||||||
|
depth: 0,
|
||||||
|
session: agent.Sessions,
|
||||||
|
initialHistoryLength: len(agent.Sessions.GetHistory("")), // Snapshot for rollback on hard abort
|
||||||
|
pendingResults: make(chan *tools.ToolResult, 16),
|
||||||
|
concurrencySem: make(chan struct{}, 5), // maxConcurrentSubTurns
|
||||||
|
}
|
||||||
|
ctx = withTurnState(ctx, rootTS)
|
||||||
|
isRootTurn = true
|
||||||
|
|
||||||
|
// Register this root turn state so HardAbort can find it
|
||||||
|
al.activeTurnStates.Store(opts.SessionKey, rootTS)
|
||||||
|
defer al.activeTurnStates.Delete(opts.SessionKey)
|
||||||
|
|
||||||
|
// Ensure the parent's pending results channel is cleaned up when this root turn finishes
|
||||||
|
defer al.unregisterSubTurnResultChannel(rootTS.turnID)
|
||||||
|
al.registerSubTurnResultChannel(rootTS.turnID, rootTS.pendingResults)
|
||||||
}
|
}
|
||||||
ctx = withTurnState(ctx, rootTS)
|
|
||||||
|
|
||||||
// Register this root turn state so HardAbort can find it
|
|
||||||
al.activeTurnStates.Store(opts.SessionKey, rootTS)
|
|
||||||
defer al.activeTurnStates.Delete(opts.SessionKey)
|
|
||||||
|
|
||||||
// Ensure the parent's pending results channel is cleaned up when this root turn finishes
|
|
||||||
defer al.unregisterSubTurnResultChannel(rootTS.turnID)
|
|
||||||
al.registerSubTurnResultChannel(rootTS.turnID, rootTS.pendingResults)
|
|
||||||
|
|
||||||
// 0. Record last channel for heartbeat notifications (skip internal channels and cli)
|
// 0. Record last channel for heartbeat notifications (skip internal channels and cli)
|
||||||
if opts.Channel != "" && opts.ChatID != "" {
|
if opts.Channel != "" && opts.ChatID != "" {
|
||||||
|
|
@ -1028,8 +1043,11 @@ func (al *AgentLoop) runAgentLoop(
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
|
|
||||||
// Signal completion to rootTS so it knows it is finished, terminating any active sub-turns
|
// Signal completion to rootTS so it knows it is finished, terminating any active sub-turns.
|
||||||
rootTS.Finish()
|
// Only call Finish() if this is a root turn (not a SubTurn recursively calling runAgentLoop).
|
||||||
|
if isRootTurn {
|
||||||
|
rootTS.Finish()
|
||||||
|
}
|
||||||
|
|
||||||
// If last tool had ForUser content and we already sent it, we might not need to send final response
|
// If last tool had ForUser content and we already sent it, we might not need to send final response
|
||||||
// This is controlled by the tool's Silent flag and ForUser content
|
// This is controlled by the tool's Silent flag and ForUser content
|
||||||
|
|
|
||||||
|
|
@ -255,7 +255,13 @@ func (al *AgentLoop) HardAbort(sessionKey string) error {
|
||||||
"initial_history_length": ts.initialHistoryLength,
|
"initial_history_length": ts.initialHistoryLength,
|
||||||
})
|
})
|
||||||
|
|
||||||
// Rollback session history to the state before this turn started
|
// IMPORTANT: Trigger cascading cancellation FIRST to stop all child SubTurns
|
||||||
|
// from adding more messages to the session. This prevents race conditions
|
||||||
|
// where rollback happens while children are still writing.
|
||||||
|
ts.Finish()
|
||||||
|
|
||||||
|
// Rollback session history to the state before this turn started.
|
||||||
|
// This must happen AFTER Finish() to ensure no child turns are still writing.
|
||||||
if ts.session != nil {
|
if ts.session != nil {
|
||||||
currentHistory := ts.session.GetHistory("")
|
currentHistory := ts.session.GetHistory("")
|
||||||
if len(currentHistory) > ts.initialHistoryLength {
|
if len(currentHistory) > ts.initialHistoryLength {
|
||||||
|
|
@ -268,8 +274,5 @@ func (al *AgentLoop) HardAbort(sessionKey string) error {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Trigger cascading cancellation to all child SubTurns
|
|
||||||
ts.Finish()
|
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -5,7 +5,6 @@ import (
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"sync"
|
"sync"
|
||||||
"sync/atomic"
|
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
"github.com/sipeed/picoclaw/pkg/session"
|
"github.com/sipeed/picoclaw/pkg/session"
|
||||||
|
|
@ -14,8 +13,8 @@ import (
|
||||||
|
|
||||||
// ====================== Config & Constants ======================
|
// ====================== Config & Constants ======================
|
||||||
const (
|
const (
|
||||||
maxSubTurnDepth = 3
|
maxSubTurnDepth = 3
|
||||||
maxConcurrentSubTurns = 5
|
maxConcurrentSubTurns = 5
|
||||||
)
|
)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
|
|
@ -78,20 +77,19 @@ type turnState struct {
|
||||||
turnID string
|
turnID string
|
||||||
parentTurnID string
|
parentTurnID string
|
||||||
depth int
|
depth int
|
||||||
childTurnIDs []string
|
childTurnIDs []string // MUST be accessed under mu lock or maybe add a getter method
|
||||||
pendingResults chan *tools.ToolResult
|
pendingResults chan *tools.ToolResult
|
||||||
session session.SessionStore
|
session session.SessionStore
|
||||||
initialHistoryLength int // Snapshot of session history length at turn start, for rollback on hard abort
|
initialHistoryLength int // Snapshot of session history length at turn start, for rollback on hard abort
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
isFinished bool // Marks if the parent Turn has ended
|
isFinished bool // MUST be accessed under mu lock
|
||||||
concurrencySem chan struct{} // Limits concurrent child sub-turns
|
concurrencySem chan struct{} // Limits concurrent child sub-turns
|
||||||
}
|
}
|
||||||
|
|
||||||
// ====================== Helper Functions ======================
|
// ====================== Helper Functions ======================
|
||||||
var globalTurnCounter int64
|
|
||||||
|
|
||||||
func generateTurnID() string {
|
func (al *AgentLoop) generateSubTurnID() string {
|
||||||
return fmt.Sprintf("subturn-%d", atomic.AddInt64(&globalTurnCounter, 1))
|
return fmt.Sprintf("subturn-%d", al.subTurnCounter.Add(1))
|
||||||
}
|
}
|
||||||
|
|
||||||
func newTurnState(ctx context.Context, id string, parent *turnState) *turnState {
|
func newTurnState(ctx context.Context, id string, parent *turnState) *turnState {
|
||||||
|
|
@ -113,13 +111,27 @@ func newTurnState(ctx context.Context, id string, parent *turnState) *turnState
|
||||||
}
|
}
|
||||||
|
|
||||||
// Finish marks the turn as finished and cancels its context, aborting any running sub-turns.
|
// Finish marks the turn as finished and cancels its context, aborting any running sub-turns.
|
||||||
|
// It also closes the pendingResults channel to signal that no more results will be delivered.
|
||||||
func (ts *turnState) Finish() {
|
func (ts *turnState) Finish() {
|
||||||
ts.mu.Lock()
|
ts.mu.Lock()
|
||||||
defer ts.mu.Unlock()
|
defer ts.mu.Unlock()
|
||||||
|
|
||||||
|
if ts.isFinished {
|
||||||
|
// Already finished - avoid double close of channel
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
ts.isFinished = true
|
ts.isFinished = true
|
||||||
|
|
||||||
if ts.cancelFunc != nil {
|
if ts.cancelFunc != nil {
|
||||||
ts.cancelFunc()
|
ts.cancelFunc()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Close the pendingResults channel to signal no more results will arrive.
|
||||||
|
// This prevents goroutine leaks from readers waiting on the channel.
|
||||||
|
if ts.pendingResults != nil {
|
||||||
|
close(ts.pendingResults)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// ephemeralSessionStore is a pure in-memory SessionStore for SubTurns.
|
// ephemeralSessionStore is a pure in-memory SessionStore for SubTurns.
|
||||||
|
|
@ -186,6 +198,24 @@ func newEphemeralSession(_ session.SessionStore) session.SessionStore {
|
||||||
|
|
||||||
// ====================== Core Function: spawnSubTurn ======================
|
// ====================== Core Function: spawnSubTurn ======================
|
||||||
func spawnSubTurn(ctx context.Context, al *AgentLoop, parentTS *turnState, cfg SubTurnConfig) (result *tools.ToolResult, err error) {
|
func spawnSubTurn(ctx context.Context, al *AgentLoop, parentTS *turnState, cfg SubTurnConfig) (result *tools.ToolResult, err error) {
|
||||||
|
// 0. Acquire concurrency semaphore FIRST to ensure it's released even if early validation fails.
|
||||||
|
// Blocks if parent already has maxConcurrentSubTurns running.
|
||||||
|
// Also respects context cancellation so we don't block forever if parent is aborted.
|
||||||
|
var semAcquired bool
|
||||||
|
if parentTS.concurrencySem != nil {
|
||||||
|
select {
|
||||||
|
case parentTS.concurrencySem <- struct{}{}:
|
||||||
|
semAcquired = true
|
||||||
|
defer func() {
|
||||||
|
if semAcquired {
|
||||||
|
<-parentTS.concurrencySem
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
case <-ctx.Done():
|
||||||
|
return nil, ctx.Err()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// 1. Depth limit check
|
// 1. Depth limit check
|
||||||
if parentTS.depth >= maxSubTurnDepth {
|
if parentTS.depth >= maxSubTurnDepth {
|
||||||
return nil, ErrDepthLimitExceeded
|
return nil, ErrDepthLimitExceeded
|
||||||
|
|
@ -196,42 +226,31 @@ func spawnSubTurn(ctx context.Context, al *AgentLoop, parentTS *turnState, cfg S
|
||||||
return nil, ErrInvalidSubTurnConfig
|
return nil, ErrInvalidSubTurnConfig
|
||||||
}
|
}
|
||||||
|
|
||||||
// 3. Acquire concurrency semaphore — blocks if parent already has maxConcurrentSubTurns running.
|
|
||||||
// Also respects context cancellation so we don't block forever if parent is aborted.
|
|
||||||
if parentTS.concurrencySem != nil {
|
|
||||||
select {
|
|
||||||
case parentTS.concurrencySem <- struct{}{}:
|
|
||||||
defer func() { <-parentTS.concurrencySem }()
|
|
||||||
case <-ctx.Done():
|
|
||||||
return nil, ctx.Err()
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Create a sub-context for the child turn to support cancellation
|
// Create a sub-context for the child turn to support cancellation
|
||||||
childCtx, cancel := context.WithCancel(ctx)
|
childCtx, cancel := context.WithCancel(ctx)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
// 4. Create child Turn state
|
// 3. Create child Turn state
|
||||||
childID := generateTurnID()
|
childID := al.generateSubTurnID()
|
||||||
childTS := newTurnState(childCtx, childID, parentTS)
|
childTS := newTurnState(childCtx, childID, parentTS)
|
||||||
|
|
||||||
// 5. Establish parent-child relationship (thread-safe)
|
// 4. Establish parent-child relationship (thread-safe)
|
||||||
parentTS.mu.Lock()
|
parentTS.mu.Lock()
|
||||||
parentTS.childTurnIDs = append(parentTS.childTurnIDs, childID)
|
parentTS.childTurnIDs = append(parentTS.childTurnIDs, childID)
|
||||||
parentTS.mu.Unlock()
|
parentTS.mu.Unlock()
|
||||||
|
|
||||||
// 6. Register the parent's pendingResults channel so the parent loop can poll it
|
// 5. Register the parent's pendingResults channel so the parent loop can poll it
|
||||||
al.registerSubTurnResultChannel(parentTS.turnID, parentTS.pendingResults)
|
al.registerSubTurnResultChannel(parentTS.turnID, parentTS.pendingResults)
|
||||||
defer al.unregisterSubTurnResultChannel(parentTS.turnID)
|
defer al.unregisterSubTurnResultChannel(parentTS.turnID)
|
||||||
|
|
||||||
// 7. Emit Spawn event (currently using Mock, will be replaced by real EventBus)
|
// 6. Emit Spawn event (currently using Mock, will be replaced by real EventBus)
|
||||||
MockEventBus.Emit(SubTurnSpawnEvent{
|
MockEventBus.Emit(SubTurnSpawnEvent{
|
||||||
ParentID: parentTS.turnID,
|
ParentID: parentTS.turnID,
|
||||||
ChildID: childID,
|
ChildID: childID,
|
||||||
Config: cfg,
|
Config: cfg,
|
||||||
})
|
})
|
||||||
|
|
||||||
// 8. Defer emitting End event, and recover from panics to ensure it's always fired
|
// 7. Defer emitting End event, and recover from panics to ensure it's always fired
|
||||||
defer func() {
|
defer func() {
|
||||||
if r := recover(); r != nil {
|
if r := recover(); r != nil {
|
||||||
err = fmt.Errorf("subturn panicked: %v", r)
|
err = fmt.Errorf("subturn panicked: %v", r)
|
||||||
|
|
@ -244,11 +263,11 @@ func spawnSubTurn(ctx context.Context, al *AgentLoop, parentTS *turnState, cfg S
|
||||||
})
|
})
|
||||||
}()
|
}()
|
||||||
|
|
||||||
// 9. Execute sub-turn via the real agent loop.
|
// 8. Execute sub-turn via the real agent loop.
|
||||||
// Build a child AgentInstance from SubTurnConfig, inheriting defaults from the parent agent.
|
// Build a child AgentInstance from SubTurnConfig, inheriting defaults from the parent agent.
|
||||||
result, err = runTurn(childCtx, al, childTS, cfg)
|
result, err = runTurn(childCtx, al, childTS, cfg)
|
||||||
|
|
||||||
// 10. Deliver result back to parent Turn
|
// 9. Deliver result back to parent Turn
|
||||||
deliverSubTurnResult(parentTS, childID, result)
|
deliverSubTurnResult(parentTS, childID, result)
|
||||||
|
|
||||||
return result, err
|
return result, err
|
||||||
|
|
@ -256,8 +275,11 @@ func spawnSubTurn(ctx context.Context, al *AgentLoop, parentTS *turnState, cfg S
|
||||||
|
|
||||||
// ====================== Result Delivery ======================
|
// ====================== Result Delivery ======================
|
||||||
func deliverSubTurnResult(parentTS *turnState, childID string, result *tools.ToolResult) {
|
func deliverSubTurnResult(parentTS *turnState, childID string, result *tools.ToolResult) {
|
||||||
|
// Check parent state under lock, but don't hold lock while sending to channel
|
||||||
parentTS.mu.Lock()
|
parentTS.mu.Lock()
|
||||||
defer parentTS.mu.Unlock()
|
isFinished := parentTS.isFinished
|
||||||
|
resultChan := parentTS.pendingResults
|
||||||
|
parentTS.mu.Unlock()
|
||||||
|
|
||||||
// Emit ResultDelivered event
|
// Emit ResultDelivered event
|
||||||
MockEventBus.Emit(SubTurnResultDeliveredEvent{
|
MockEventBus.Emit(SubTurnResultDeliveredEvent{
|
||||||
|
|
@ -266,10 +288,24 @@ func deliverSubTurnResult(parentTS *turnState, childID string, result *tools.Too
|
||||||
Result: result,
|
Result: result,
|
||||||
})
|
})
|
||||||
|
|
||||||
if !parentTS.isFinished {
|
if !isFinished && resultChan != nil {
|
||||||
// Parent Turn is still running → Place in pending queue (handled automatically by parent loop in next round)
|
// Parent Turn is still running → Place in pending queue (handled automatically by parent loop in next round)
|
||||||
|
// Use defer/recover to handle the case where the channel is closed between our check and the send.
|
||||||
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
// Channel was closed - treat as orphan result
|
||||||
|
if result != nil {
|
||||||
|
MockEventBus.Emit(SubTurnOrphanResultEvent{
|
||||||
|
ParentID: parentTS.turnID,
|
||||||
|
ChildID: childID,
|
||||||
|
Result: result,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
select {
|
select {
|
||||||
case parentTS.pendingResults <- result:
|
case resultChan <- result:
|
||||||
default:
|
default:
|
||||||
fmt.Println("[SubTurn] warning: pendingResults channel full")
|
fmt.Println("[SubTurn] warning: pendingResults channel full")
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -2,8 +2,11 @@ package agent
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"fmt"
|
||||||
"reflect"
|
"reflect"
|
||||||
|
"sync"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/providers"
|
"github.com/sipeed/picoclaw/pkg/providers"
|
||||||
"github.com/sipeed/picoclaw/pkg/tools"
|
"github.com/sipeed/picoclaw/pkg/tools"
|
||||||
|
|
@ -500,3 +503,221 @@ func TestHardAbortSessionRollback(t *testing.T) {
|
||||||
t.Error("history content does not match initial state after rollback")
|
t.Error("history content does not match initial state after rollback")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestNestedSubTurnHierarchy verifies that nested SubTurns maintain correct
|
||||||
|
// parent-child relationships and depth tracking when recursively calling runAgentLoop.
|
||||||
|
func TestNestedSubTurnHierarchy(t *testing.T) {
|
||||||
|
al, _, _, _, cleanup := newTestAgentLoop(t)
|
||||||
|
defer cleanup()
|
||||||
|
|
||||||
|
// Track spawned turns and their depths
|
||||||
|
type turnInfo struct {
|
||||||
|
parentID string
|
||||||
|
childID string
|
||||||
|
depth int
|
||||||
|
}
|
||||||
|
var spawnedTurns []turnInfo
|
||||||
|
var mu sync.Mutex
|
||||||
|
|
||||||
|
// Override MockEventBus to capture spawn events
|
||||||
|
originalEmit := MockEventBus.Emit
|
||||||
|
defer func() { MockEventBus.Emit = originalEmit }()
|
||||||
|
|
||||||
|
MockEventBus.Emit = func(event any) {
|
||||||
|
if spawnEvent, ok := event.(SubTurnSpawnEvent); ok {
|
||||||
|
mu.Lock()
|
||||||
|
// Extract depth from context (we'll verify this matches expected depth)
|
||||||
|
spawnedTurns = append(spawnedTurns, turnInfo{
|
||||||
|
parentID: spawnEvent.ParentID,
|
||||||
|
childID: spawnEvent.ChildID,
|
||||||
|
})
|
||||||
|
mu.Unlock()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create a root turn
|
||||||
|
rootSession := &ephemeralSessionStore{}
|
||||||
|
rootTS := &turnState{
|
||||||
|
ctx: context.Background(),
|
||||||
|
turnID: "root-turn",
|
||||||
|
depth: 0,
|
||||||
|
session: rootSession,
|
||||||
|
pendingResults: make(chan *tools.ToolResult, 16),
|
||||||
|
concurrencySem: make(chan struct{}, 5),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Spawn a child (depth 1)
|
||||||
|
childCfg := SubTurnConfig{Model: "gpt-4o-mini"}
|
||||||
|
_, err := spawnSubTurn(context.Background(), al, rootTS, childCfg)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("failed to spawn child: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify we captured the spawn event
|
||||||
|
mu.Lock()
|
||||||
|
if len(spawnedTurns) != 1 {
|
||||||
|
t.Fatalf("expected 1 spawn event, got %d", len(spawnedTurns))
|
||||||
|
}
|
||||||
|
if spawnedTurns[0].parentID != "root-turn" {
|
||||||
|
t.Errorf("expected parent ID 'root-turn', got %s", spawnedTurns[0].parentID)
|
||||||
|
}
|
||||||
|
mu.Unlock()
|
||||||
|
|
||||||
|
// Verify root turn has the child in its childTurnIDs
|
||||||
|
rootTS.mu.Lock()
|
||||||
|
if len(rootTS.childTurnIDs) != 1 {
|
||||||
|
t.Errorf("expected root to have 1 child, got %d", len(rootTS.childTurnIDs))
|
||||||
|
}
|
||||||
|
rootTS.mu.Unlock()
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestDeliverSubTurnResultNoDeadlock verifies that deliverSubTurnResult doesn't
|
||||||
|
// deadlock when multiple goroutines are accessing the parent turnState concurrently.
|
||||||
|
func TestDeliverSubTurnResultNoDeadlock(t *testing.T) {
|
||||||
|
parent := &turnState{
|
||||||
|
ctx: context.Background(),
|
||||||
|
turnID: "parent-deadlock-test",
|
||||||
|
depth: 0,
|
||||||
|
pendingResults: make(chan *tools.ToolResult, 2), // Small buffer to test blocking
|
||||||
|
isFinished: false,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Simulate multiple child turns delivering results concurrently
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
numChildren := 10
|
||||||
|
|
||||||
|
for i := 0; i < numChildren; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func(id int) {
|
||||||
|
defer wg.Done()
|
||||||
|
result := &tools.ToolResult{ForLLM: fmt.Sprintf("result-%d", id)}
|
||||||
|
deliverSubTurnResult(parent, fmt.Sprintf("child-%d", id), result)
|
||||||
|
}(i)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Concurrently read from the channel to prevent blocking
|
||||||
|
go func() {
|
||||||
|
for i := 0; i < numChildren; i++ {
|
||||||
|
select {
|
||||||
|
case <-parent.pendingResults:
|
||||||
|
case <-time.After(2 * time.Second):
|
||||||
|
t.Error("timeout waiting for result")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
// Wait for all deliveries to complete (with timeout)
|
||||||
|
done := make(chan struct{})
|
||||||
|
go func() {
|
||||||
|
wg.Wait()
|
||||||
|
close(done)
|
||||||
|
}()
|
||||||
|
|
||||||
|
select {
|
||||||
|
case <-done:
|
||||||
|
// Success - no deadlock
|
||||||
|
case <-time.After(3 * time.Second):
|
||||||
|
t.Fatal("deadlock detected: deliverSubTurnResult blocked")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestHardAbortOrderOfOperations verifies that HardAbort calls Finish() before
|
||||||
|
// rolling back session history, minimizing the race window where new messages
|
||||||
|
// could be added after rollback.
|
||||||
|
func TestHardAbortOrderOfOperations(t *testing.T) {
|
||||||
|
al, _, _, _, cleanup := newTestAgentLoop(t)
|
||||||
|
defer cleanup()
|
||||||
|
|
||||||
|
sess := &ephemeralSessionStore{
|
||||||
|
history: []providers.Message{
|
||||||
|
{Role: "user", Content: "initial message"},
|
||||||
|
{Role: "assistant", Content: "response 1"},
|
||||||
|
{Role: "user", Content: "follow-up"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
rootTS := &turnState{
|
||||||
|
ctx: ctx,
|
||||||
|
cancelFunc: cancel,
|
||||||
|
turnID: "test-session-order",
|
||||||
|
depth: 0,
|
||||||
|
session: sess,
|
||||||
|
initialHistoryLength: 1, // Snapshot: 1 message
|
||||||
|
pendingResults: make(chan *tools.ToolResult, 16),
|
||||||
|
concurrencySem: make(chan struct{}, 5),
|
||||||
|
}
|
||||||
|
|
||||||
|
al.activeTurnStates.Store("test-session-order", rootTS)
|
||||||
|
|
||||||
|
// Trigger HardAbort
|
||||||
|
err := al.HardAbort("test-session-order")
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("HardAbort failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify context was cancelled (Finish() was called)
|
||||||
|
select {
|
||||||
|
case <-rootTS.ctx.Done():
|
||||||
|
// Good - context was cancelled
|
||||||
|
default:
|
||||||
|
t.Error("expected context to be cancelled after HardAbort")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify history was rolled back
|
||||||
|
finalHistory := sess.GetHistory("")
|
||||||
|
if len(finalHistory) != 1 {
|
||||||
|
t.Errorf("expected history to rollback to 1 message, got %d", len(finalHistory))
|
||||||
|
}
|
||||||
|
|
||||||
|
if finalHistory[0].Content != "initial message" {
|
||||||
|
t.Error("history content does not match initial state after rollback")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestFinishClosesChannel verifies that Finish() closes the pendingResults channel
|
||||||
|
// and that deliverSubTurnResult handles closed channels gracefully.
|
||||||
|
func TestFinishClosesChannel(t *testing.T) {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
ts := &turnState{
|
||||||
|
ctx: ctx,
|
||||||
|
cancelFunc: cancel,
|
||||||
|
turnID: "test-finish-channel",
|
||||||
|
depth: 0,
|
||||||
|
pendingResults: make(chan *tools.ToolResult, 2),
|
||||||
|
isFinished: false,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify channel is open initially
|
||||||
|
select {
|
||||||
|
case ts.pendingResults <- &tools.ToolResult{ForLLM: "test"}:
|
||||||
|
// Good - channel is open
|
||||||
|
// Drain the message we just sent
|
||||||
|
<-ts.pendingResults
|
||||||
|
default:
|
||||||
|
t.Fatal("channel should be open initially")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Call Finish()
|
||||||
|
ts.Finish()
|
||||||
|
|
||||||
|
// Verify channel is closed
|
||||||
|
_, ok := <-ts.pendingResults
|
||||||
|
if ok {
|
||||||
|
t.Error("expected channel to be closed after Finish()")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify Finish() is idempotent (can be called multiple times)
|
||||||
|
ts.Finish() // Should not panic
|
||||||
|
|
||||||
|
// Verify deliverSubTurnResult doesn't panic when sending to closed channel
|
||||||
|
result := &tools.ToolResult{ForLLM: "late result"}
|
||||||
|
|
||||||
|
// This should not panic - it should recover and emit OrphanResultEvent
|
||||||
|
deliverSubTurnResult(ts, "child-1", result)
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue