diff --git a/pkg/agent/agent_init.go b/pkg/agent/agent_init.go index 50f0227a1..a6a8e84ad 100644 --- a/pkg/agent/agent_init.go +++ b/pkg/agent/agent_init.go @@ -345,6 +345,18 @@ func registerSharedTools( // Also register the synchronous subagent tool subagentTool := tools.NewSubagentTool(subagentManager) subagentTool.SetSpawner(NewSubTurnSpawner(al)) + subagentTool.SetAllowlistChecker(func(targetAgentID string) bool { + return registry.CanSpawnSubagent(currentAgentID, targetAgentID) + }) + subagentTool.SetTargetModelResolver(func(targetAgentID string) string { + if targetAgentID == "" { + return agent.Model + } + if targetAgent, ok := al.GetRegistry().GetAgent(targetAgentID); ok { + return targetAgent.Model + } + return agent.Model + }) agent.Tools.Register(subagentTool) } if spawnStatusEnabled { diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index feeabe536..3587293bd 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -338,10 +338,12 @@ func (sm *SubagentManager) ListTaskCopies() []SubagentTask { // SubagentTool executes a subagent task synchronously and returns the result. // It directly calls SubTurnSpawner with Async=false for synchronous execution. type SubagentTool struct { - spawner SubTurnSpawner - defaultModel string - maxTokens int - temperature float64 + spawner SubTurnSpawner + defaultModel string + maxTokens int + temperature float64 + allowlistCheck func(targetAgentID string) bool + targetModelResolver func(targetAgentID string) string } func NewSubagentTool(manager *SubagentManager) *SubagentTool { @@ -360,6 +362,14 @@ func (t *SubagentTool) SetSpawner(spawner SubTurnSpawner) { t.spawner = spawner } +func (t *SubagentTool) SetAllowlistChecker(check func(targetAgentID string) bool) { + t.allowlistCheck = check +} + +func (t *SubagentTool) SetTargetModelResolver(resolver func(targetAgentID string) string) { + t.targetModelResolver = resolver +} + func (t *SubagentTool) Name() string { return "subagent" } @@ -380,6 +390,10 @@ func (t *SubagentTool) Parameters() map[string]any { "type": "string", "description": "Optional short label for the task (for display)", }, + "agent_id": map[string]any{ + "type": "string", + "description": "Optional target agent ID to delegate the task to", + }, }, "required": []string{"task"}, } @@ -392,6 +406,17 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe } label, _ := args["label"].(string) + agentID, _ := args["agent_id"].(string) + if agentID != "" && t.allowlistCheck != nil && !t.allowlistCheck(agentID) { + return ErrorResult(fmt.Sprintf("Not allowed to target agent '%s'", agentID)) + } + + modelToUse := t.defaultModel + if agentID != "" && t.targetModelResolver != nil { + if resolved := t.targetModelResolver(agentID); resolved != "" { + modelToUse = resolved + } + } // Build system prompt for subagent systemPrompt := fmt.Sprintf( @@ -414,7 +439,7 @@ Task: %s`, // Use spawner if available (direct SpawnSubTurn call) if t.spawner != nil { result, err := t.spawner.SpawnSubTurn(ctx, SubTurnConfig{ - Model: t.defaultModel, + Model: modelToUse, Tools: nil, // Will inherit from parent via context SystemPrompt: systemPrompt, MaxTokens: t.maxTokens, diff --git a/pkg/tools/subagent_targeting_test.go b/pkg/tools/subagent_targeting_test.go new file mode 100644 index 000000000..cd3c68783 --- /dev/null +++ b/pkg/tools/subagent_targeting_test.go @@ -0,0 +1,75 @@ +package tools + +import ( + "context" + "testing" +) + +type capturingSubTurnSpawner struct { + lastCfg SubTurnConfig +} + +func (s *capturingSubTurnSpawner) SpawnSubTurn(_ context.Context, cfg SubTurnConfig) (*ToolResult, error) { + s.lastCfg = cfg + return &ToolResult{ForLLM: "ok", ForUser: "ok"}, nil +} + +func TestSubagentToolRejectsDisallowedAgentID(t *testing.T) { + tool := &SubagentTool{defaultModel: "main-model"} + tool.SetAllowlistChecker(func(targetAgentID string) bool { + return targetAgentID != "code" + }) + + result := tool.Execute(context.Background(), map[string]any{ + "task": "do something", + "agent_id": "code", + }) + + if result == nil || !result.IsError { + t.Fatalf("expected error result, got %#v", result) + } +} + +func TestSubagentToolUsesResolvedTargetModelForAgentID(t *testing.T) { + spawner := &capturingSubTurnSpawner{} + tool := &SubagentTool{defaultModel: "main-model"} + tool.SetSpawner(spawner) + tool.SetAllowlistChecker(func(targetAgentID string) bool { + return targetAgentID == "code" + }) + tool.SetTargetModelResolver(func(targetAgentID string) string { + if targetAgentID == "code" { + return "code-model" + } + return "" + }) + + result := tool.Execute(context.Background(), map[string]any{ + "task": "do something", + "agent_id": "code", + }) + + if result == nil || result.IsError { + t.Fatalf("expected success result, got %#v", result) + } + if spawner.lastCfg.Model != "code-model" { + t.Fatalf("expected model code-model, got %q", spawner.lastCfg.Model) + } +} + +func TestSubagentToolFallsBackToDefaultModelWithoutAgentID(t *testing.T) { + spawner := &capturingSubTurnSpawner{} + tool := &SubagentTool{defaultModel: "main-model"} + tool.SetSpawner(spawner) + + result := tool.Execute(context.Background(), map[string]any{ + "task": "do something", + }) + + if result == nil || result.IsError { + t.Fatalf("expected success result, got %#v", result) + } + if spawner.lastCfg.Model != "main-model" { + t.Fatalf("expected model main-model, got %q", spawner.lastCfg.Model) + } +}