From b6638a506701ff26bca44154d26dbcea1c30c532 Mon Sep 17 00:00:00 2001
From: Paul De Velder <paul.develder@attentia.be>
Date: Mon, 23 Feb 2026 15:04:14 +0100
Subject: [PATCH] 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.
---
 pkg/tools/subagent.go           | 14 ++++++
 pkg/tools/subagent_tool_test.go | 83 +++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

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