diff --git a/pkg/agent/loop_execute.go b/pkg/agent/loop_execute.go index 1b2ee47..61f054f 100644 --- a/pkg/agent/loop_execute.go +++ b/pkg/agent/loop_execute.go @@ -26,6 +26,24 @@ func newToolCallID() string { // short-circuit with a synthesized tool-error result. A fatal error // recorded by any goroutine (e.g. anti-loop detector) breaks out of // the wave loop. +// defaultMaxParallelToolCalls bounds concurrent tool execution within one +// dependency wave. Wide fan-outs are the pathological case this exists for: +// a model asking for 40 calls at once would otherwise open 40 sockets or +// spawn 40 subprocesses. Waves are usually smaller than this, so the cap is +// invisible in normal operation. +const defaultMaxParallelToolCalls = 8 + +// toolCallSemaphore returns a buffered channel bounding wave concurrency, or +// nil when no cap applies — an unlimited setting, or a wave already at or +// below the cap, both skip the channel and its send/receive pair entirely. +func (al *AgentLoop) toolCallSemaphore(waveSize int) chan struct{} { + n := al.MaxParallelToolCalls + if n <= 0 || waveSize <= n { + return nil + } + return make(chan struct{}, n) +} + func (al *AgentLoop) executeToolWaves(ctx context.Context, st *iterationState, scheduled []PendingToolCall) *waveState { waves, schedErr := scheduleToolCalls(scheduled) if schedErr != nil { @@ -43,11 +61,16 @@ func (al *AgentLoop) executeToolWaves(ctx context.Context, st *iterationState, s })) }) + sem := al.toolCallSemaphore(len(substitutedWave)) var wg sync.WaitGroup for _, tc := range substitutedWave { wg.Add(1) go func(tCall PendingToolCall) { defer wg.Done() + if sem != nil { + sem <- struct{}{} + defer func() { <-sem }() + } al.executeToolCall(ctx, st, ws, tCall) }(tc) } diff --git a/pkg/agent/loop_parallel_cap_test.go b/pkg/agent/loop_parallel_cap_test.go new file mode 100644 index 0000000..a86c83c --- /dev/null +++ b/pkg/agent/loop_parallel_cap_test.go @@ -0,0 +1,123 @@ +package agent + +import ( + "context" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/hung12ct/gopheragent/pkg/tools" +) + +// concurrencyProbe records the high-water mark of overlapping Execute calls +// so a test can assert the wave semaphore actually bounds them. +type concurrencyProbe struct { + live atomic.Int32 + peak atomic.Int32 + runs atomic.Int32 +} + +func (c *concurrencyProbe) Descriptor() tools.ToolDescriptor { + return tools.ToolDescriptor{ + Name: "probe", + Description: "records concurrency", + Display: tools.DefaultDisplay("probe", "records concurrency"), + } +} + +func (c *concurrencyProbe) Execute(_ context.Context, args string) (tools.Result, error) { + live := c.live.Add(1) + for { + peak := c.peak.Load() + if live <= peak || c.peak.CompareAndSwap(peak, live) { + break + } + } + // Hold the slot long enough that an unbounded fan-out genuinely overlaps. + time.Sleep(5 * time.Millisecond) + c.live.Add(-1) + c.runs.Add(1) + return tools.Text("probe:" + args), nil +} + +// distinctCalls builds n calls to "probe" with unique arguments, so neither +// the anti-loop detector nor the result cache collapses them. +func distinctCalls(n int) []PendingToolCall { + out := make([]PendingToolCall, n) + for i := range n { + out[i] = PendingToolCall{ + ID: fmt.Sprintf("p%d", i), + Name: "probe", + ArgsJSON: fmt.Sprintf(`{"i":%d}`, i), + } + } + return out +} + +func runProbeWave(t *testing.T, calls int, opts ...Option) *concurrencyProbe { + t.Helper() + probe := &concurrencyProbe{} + provider := &scriptProvider{turns: []LLMResult{ + {ToolCalls: distinctCalls(calls)}, + {Content: "final"}, + }} + loop, _ := setup(provider, probe) + for _, opt := range opts { + opt(loop) + } + + if _, err := loop.RunIteration(context.Background(), "s1", "go"); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got := probe.runs.Load(); got != int32(calls) { + t.Fatalf("executed %d calls, want %d — the cap must delay calls, never drop them", got, calls) + } + return probe +} + +func TestMaxParallelToolCalls_BoundsWaveConcurrency(t *testing.T) { + probe := runProbeWave(t, 20, WithMaxParallelToolCalls(4)) + + if peak := probe.peak.Load(); peak > 4 { + t.Fatalf("peak concurrency = %d, want <= 4", peak) + } +} + +func TestMaxParallelToolCalls_DefaultApplies(t *testing.T) { + probe := runProbeWave(t, 20) + + if peak := probe.peak.Load(); peak > defaultMaxParallelToolCalls { + t.Fatalf("peak concurrency = %d, want <= %d (default)", peak, defaultMaxParallelToolCalls) + } +} + +func TestMaxParallelToolCalls_ZeroIsUnlimited(t *testing.T) { + probe := runProbeWave(t, 20, WithMaxParallelToolCalls(0)) + + // Not asserting an exact peak: the scheduler is free to interleave. The + // point is that no cap is applied, which only shows up as a peak above + // the default the constructor would otherwise have installed. + if peak := probe.peak.Load(); peak <= defaultMaxParallelToolCalls { + t.Skipf("peak %d did not exceed the default cap; scheduling-dependent, not a failure", peak) + } +} + +// toolCallSemaphore must skip the channel entirely when it cannot bind, so +// the common small-wave path pays nothing. +func TestToolCallSemaphore_NilWhenNoCapBinds(t *testing.T) { + al := &AgentLoop{MaxParallelToolCalls: 8} + for _, size := range []int{0, 1, 8} { + if sem := al.toolCallSemaphore(size); sem != nil { + t.Fatalf("waveSize %d: got a semaphore, want nil", size) + } + } + if sem := al.toolCallSemaphore(9); sem == nil || cap(sem) != 8 { + t.Fatalf("waveSize 9: want a semaphore of cap 8, got %v", sem) + } + + unlimited := &AgentLoop{MaxParallelToolCalls: 0} + if sem := unlimited.toolCallSemaphore(100); sem != nil { + t.Fatal("unlimited: got a semaphore, want nil") + } +} diff --git a/pkg/agent/loop_stream.go b/pkg/agent/loop_stream.go index d033bef..ac09b6e 100644 --- a/pkg/agent/loop_stream.go +++ b/pkg/agent/loop_stream.go @@ -426,6 +426,14 @@ type AgentLoop struct { // they did not execute. A "thought" event announces the truncation. MaxToolCallsPerTurn int + // MaxParallelToolCalls caps how many tool calls execute concurrently + // within one dependency wave. Defaults to defaultMaxParallelToolCalls; + // 0 means unlimited. Unlike MaxToolCallsPerTurn this drops nothing — a + // call over the cap waits for a slot, so the wave's result set is + // identical either way. It bounds live resource use (sockets, + // subprocesses, provider rate limits) when a model emits a wide fan-out. + MaxParallelToolCalls int + // MaxToolCallsPerSession caps the cumulative number of tool calls // scheduled across all iterations of a single Run. 0 (default) means // unlimited. Distinct from MaxToolCallsPerTurn, which only bounds the @@ -518,13 +526,14 @@ type AgentLoop struct { // Optional hooks run before each iteration for security/policy enforcement. func NewAgentLoop(sessions SessionManager, registry *tools.Registry, llm LLMProvider, hooks ...Hook) *AgentLoop { return &AgentLoop{ - Sessions: sessions, - Tools: registry, - LLM: llm, - MaxIters: 15, - EmitThoughts: true, - BeforeHooks: hooks, - AutoCacheSystem: true, + Sessions: sessions, + Tools: registry, + LLM: llm, + MaxIters: 15, + EmitThoughts: true, + BeforeHooks: hooks, + AutoCacheSystem: true, + MaxParallelToolCalls: defaultMaxParallelToolCalls, } } diff --git a/pkg/agent/options.go b/pkg/agent/options.go index 78b629c..6affbb3 100644 --- a/pkg/agent/options.go +++ b/pkg/agent/options.go @@ -245,6 +245,16 @@ func WithMaxToolCallsPerTurn(n int) Option { return func(al *AgentLoop) { al.MaxToolCallsPerTurn = n } } +// WithMaxParallelToolCalls caps how many tool calls run concurrently inside +// one dependency wave. Defaults to 8; pass 0 for unlimited. Nothing is +// dropped — calls over the cap wait for a slot, so results are unchanged and +// only peak resource use differs. Raise it for latency-bound tools (HTTP +// fan-out), lower it for tools that each hold a scarce resource (database +// connections, subprocesses). +func WithMaxParallelToolCalls(n int) Option { + return func(al *AgentLoop) { al.MaxParallelToolCalls = n } +} + // WithMaxToolCallsPerSession caps cumulative tool calls across all // iterations of a single Run. 0 (default) means unlimited. When the cap // trips, the loop emits LimitExhaustedEvent and saves history.