Prevent infinite recursion when agents delegate to each other as tools

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

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

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), &params); 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{

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

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