diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 91ebff636..893deab8a 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -36,6 +36,8 @@ type SubagentManager struct { hasMaxTokens bool hasTemperature bool nextID int + maxConcurrent int + activeCount int } func NewSubagentManager( @@ -52,6 +54,7 @@ func NewSubagentManager( tools: NewToolRegistry(), maxIterations: 10, nextID: 1, + maxConcurrent: 3, } } @@ -88,8 +91,13 @@ func (sm *SubagentManager) Spawn( sm.mu.Lock() defer sm.mu.Unlock() + if sm.maxConcurrent > 0 && sm.activeCount >= sm.maxConcurrent { + return "", fmt.Errorf("too many concurrent subagents (max %d)", sm.maxConcurrent) + } + taskID := fmt.Sprintf("subagent-%d", sm.nextID) sm.nextID++ + sm.activeCount++ subagentTask := &SubagentTask{ ID: taskID, @@ -113,6 +121,12 @@ func (sm *SubagentManager) Spawn( } func (sm *SubagentManager) runTask(ctx context.Context, task *SubagentTask, callback AsyncCallback) { + defer func() { + sm.mu.Lock() + sm.activeCount-- + sm.mu.Unlock() + }() + task.Status = "running" task.Created = time.Now().UnixMilli() diff --git a/pkg/tools/subagent_tool_test.go b/pkg/tools/subagent_tool_test.go index 59bfdffae..69ea561b2 100644 --- a/pkg/tools/subagent_tool_test.go +++ b/pkg/tools/subagent_tool_test.go @@ -2,8 +2,11 @@ package tools import ( "context" + "fmt" "strings" + "sync" "testing" + "time" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/providers" @@ -348,3 +351,83 @@ func TestSubagentTool_ForUserTruncation(t *testing.T) { t.Error("ForLLM should contain reference to original task") } } + +// SlowMockLLMProvider delays responses to allow testing concurrency. +type SlowMockLLMProvider struct { + delay time.Duration +} + +func (m *SlowMockLLMProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + select { + case <-time.After(m.delay): + case <-ctx.Done(): + return nil, ctx.Err() + } + return &providers.LLMResponse{Content: "done"}, nil +} + +func (m *SlowMockLLMProvider) GetDefaultModel() string { return "test-model" } +func (m *SlowMockLLMProvider) SupportsTools() bool { return false } +func (m *SlowMockLLMProvider) GetContextWindow() int { return 4096 } + +// TestSubagentManager_ConcurrencyLimit verifies that spawning beyond maxConcurrent is rejected. +func TestSubagentManager_ConcurrencyLimit(t *testing.T) { + provider := &SlowMockLLMProvider{delay: 2 * time.Second} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + // maxConcurrent defaults to 3 + + ctx := context.Background() + + // Spawn 3 tasks — should all succeed + for i := 0; i < 3; i++ { + _, err := manager.Spawn(ctx, fmt.Sprintf("task-%d", i), fmt.Sprintf("label-%d", i), "agent", "cli", "direct", nil) + if err != nil { + t.Fatalf("Spawn %d should succeed, got: %v", i, err) + } + } + + // 4th spawn should fail + _, err := manager.Spawn(ctx, "task-overflow", "overflow", "agent", "cli", "direct", nil) + if err == nil { + t.Fatal("Expected error for exceeding maxConcurrent, got nil") + } + if !strings.Contains(err.Error(), "too many concurrent subagents") { + t.Errorf("Expected 'too many concurrent subagents' error, got: %v", err) + } +} + +// TestSubagentManager_ConcurrencyReleasesAfterCompletion verifies slots are freed after task completion. +func TestSubagentManager_ConcurrencyReleasesAfterCompletion(t *testing.T) { + provider := &SlowMockLLMProvider{delay: 50 * time.Millisecond} + manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + + ctx := context.Background() + + // Fill all 3 slots + var wg sync.WaitGroup + for i := 0; i < 3; i++ { + wg.Add(1) + cb := func(_ context.Context, _ *ToolResult) { wg.Done() } + _, err := manager.Spawn(ctx, fmt.Sprintf("task-%d", i), fmt.Sprintf("label-%d", i), "agent", "cli", "direct", cb) + if err != nil { + t.Fatalf("Spawn %d should succeed, got: %v", i, err) + } + } + + // Wait for all tasks to complete + wg.Wait() + // Small delay to ensure activeCount is decremented + time.Sleep(20 * time.Millisecond) + + // Now spawning should succeed again + _, err := manager.Spawn(ctx, "task-after", "after", "agent", "cli", "direct", nil) + if err != nil { + t.Fatalf("Spawn after completion should succeed, got: %v", err) + } +}