From a710a7472dfc7cca1d5d1d04e5ab2e3abfa56de5 Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Fri, 13 Mar 2026 18:15:45 +0100 Subject: [PATCH] Prevent infinite recursion when agents delegate to each other as tools Signed-off-by: Bryan Frimin --- pkg/agent/agent.go | 11 ++++ pkg/agent/agent_tool.go | 18 ++++++ pkg/agent/agent_tool_test.go | 109 +++++++++++++++++++++++++++++++++++ pkg/agent/errors.go | 8 +++ 4 files changed, 146 insertions(+) diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go index 2fe3c9750..3d9a6ccad 100644 --- a/pkg/agent/agent.go +++ b/pkg/agent/agent.go @@ -39,6 +39,7 @@ type ( handoffs []*Handoff mcpServers []*MCPServer maxTurns int + maxToolDepth int client *llm.Client logger *log.Logger hooks []RunHooks @@ -60,6 +61,7 @@ func New(name string, client *llm.Client, opts ...Option) *Agent { name: name, client: client, maxTurns: DefaultMaxTurns, + maxToolDepth: DefaultMaxToolDepth, toolUseBehavior: RunLLMAgain(), resetToolChoice: true, logger: log.NewLogger(log.WithOutput(io.Discard)), @@ -196,6 +198,15 @@ func WithMaxTurns(n int) Option { } } +func WithMaxToolDepth(n int) Option { + return func(a *Agent) { + if n < 1 { + n = 1 + } + a.maxToolDepth = n + } +} + func WithTemperature(t float64) Option { return func(a *Agent) { a.modelSettings.Temperature = &t diff --git a/pkg/agent/agent_tool.go b/pkg/agent/agent_tool.go index 54b32adb1..0a8fda8f3 100644 --- a/pkg/agent/agent_tool.go +++ b/pkg/agent/agent_tool.go @@ -22,6 +22,8 @@ import ( "go.probo.inc/probo/pkg/llm" ) +const DefaultMaxToolDepth = 16 + type ( agentTool struct { agent *Agent @@ -33,12 +35,21 @@ type ( agentToolParams struct { Input string `json:"input" jsonschema:"The input to send to the agent"` } + + agentToolDepthKey struct{} ) var ( agentToolParamsSchema = jsonSchemaFor[agentToolParams]() ) +func agentToolDepth(ctx context.Context) int { + if v, ok := ctx.Value(agentToolDepthKey{}).(int); ok { + return v + } + return 0 +} + func newAgentTool(agent *Agent, name, description string) *agentTool { return &agentTool{ agent: agent, @@ -59,6 +70,11 @@ func (t *agentTool) Definition() llm.Tool { } func (t *agentTool) Execute(ctx context.Context, arguments string) (ToolResult, error) { + depth := agentToolDepth(ctx) + if depth >= t.agent.maxToolDepth { + return ToolResult{}, &MaxToolDepthExceededError{MaxDepth: t.agent.maxToolDepth} + } + var params agentToolParams if err := json.Unmarshal([]byte(arguments), ¶ms); err != nil { @@ -68,6 +84,8 @@ func (t *agentTool) Execute(ctx context.Context, arguments string) (ToolResult, }, nil } + ctx = context.WithValue(ctx, agentToolDepthKey{}, depth+1) + result, err := t.agent.Run( ctx, []llm.Message{ diff --git a/pkg/agent/agent_tool_test.go b/pkg/agent/agent_tool_test.go index a95447b21..73ca93028 100644 --- a/pkg/agent/agent_tool_test.go +++ b/pkg/agent/agent_tool_test.go @@ -789,6 +789,115 @@ func TestAgentTool_Execute_NestedApproval(t *testing.T) { ) } +func TestAgentTool_Execute_DepthLimit(t *testing.T) { + t.Parallel() + + t.Run( + "deep agent-tool chain stops at depth limit", + func(t *testing.T) { + t.Parallel() + + innerProvider := &mockProvider{} + + innerAgent := agent.New( + "inner", + newTestClient(innerProvider), + agent.WithModel("test-model"), + agent.WithMaxToolDepth(1), + ) + + middleProvider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_mid", + Function: llm.FunctionCall{Name: "call_inner", Arguments: `{"input":"ping"}`}, + }), + stopResponse("inner was unreachable"), + }, + } + + middleAgent := agent.New( + "middle", + newTestClient(middleProvider), + agent.WithModel("test-model"), + agent.WithTools(innerAgent.AsTool("call_inner", "Call inner")), + ) + + outerProvider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_out", + Function: llm.FunctionCall{Name: "call_middle", Arguments: `{"input":"start"}`}, + }), + stopResponse("outer done"), + }, + } + + outerAgent := agent.New( + "outer", + newTestClient(outerProvider), + agent.WithModel("test-model"), + agent.WithTools(middleAgent.AsTool("call_middle", "Call middle")), + ) + + result, err := outerAgent.Run( + context.Background(), + []llm.Message{userMessage("go")}, + ) + + require.NoError(t, err) + assert.Equal(t, "outer done", result.FinalMessage().Text()) + assert.Equal(t, 0, innerProvider.calls, "inner agent should never be called") + }, + ) + + t.Run( + "delegation within depth limit succeeds", + func(t *testing.T) { + t.Parallel() + + innerProvider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + stopResponse("inner result"), + }, + } + + innerAgent := agent.New( + "inner", + newTestClient(innerProvider), + agent.WithModel("test-model"), + agent.WithMaxToolDepth(2), + ) + + outerProvider := &mockProvider{ + responses: []*llm.ChatCompletionResponse{ + toolCallResponse(llm.ToolCall{ + ID: "tc_1", + Function: llm.FunctionCall{Name: "call_inner", Arguments: `{"input":"hello"}`}, + }), + stopResponse("outer result"), + }, + } + + outerAgent := agent.New( + "outer", + newTestClient(outerProvider), + agent.WithModel("test-model"), + agent.WithTools(innerAgent.AsTool("call_inner", "Call inner")), + ) + + result, err := outerAgent.Run( + context.Background(), + []llm.Message{userMessage("go")}, + ) + + require.NoError(t, err) + assert.Equal(t, "outer result", result.FinalMessage().Text()) + assert.Equal(t, 1, innerProvider.calls, "inner agent should be called once") + }, + ) +} + func TestAgentTool_InterfaceSatisfaction(t *testing.T) { t.Parallel() diff --git a/pkg/agent/errors.go b/pkg/agent/errors.go index 595597f3f..3f44fec58 100644 --- a/pkg/agent/errors.go +++ b/pkg/agent/errors.go @@ -25,6 +25,10 @@ type ( MaxTurns int } + MaxToolDepthExceededError struct { + MaxDepth int + } + InputGuardrailTrippedError struct { Guardrail string Message string @@ -79,6 +83,10 @@ func (e *MaxTurnsExceededError) Error() string { return fmt.Sprintf("agent exceeded maximum number of turns (%d)", e.MaxTurns) } +func (e *MaxToolDepthExceededError) Error() string { + return fmt.Sprintf("agent-tool delegation exceeded maximum depth (%d)", e.MaxDepth) +} + func (e *InputGuardrailTrippedError) Error() string { return fmt.Sprintf("input guardrail %q tripped: %s", e.Guardrail, e.Message) }