feat(security): add spawn concurrency limit to prevent fork-bombing

Add maxConcurrent (default 3) and activeCount fields to SubagentManager.
Spawn rejects new subagents when the limit is reached. Active count is
decremented when tasks complete, freeing slots for new spawns.
This commit is contained in:
Paul De Velder 2026-02-23 15:04:14 +01:00
parent ca390846a5
commit b6638a5067
2 changed files with 97 additions and 0 deletions

View file

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

View file

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