diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index f20a56b9c..70cf756bd 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -227,6 +227,13 @@ func registerSharedTools( if cfg.Tools.IsToolEnabled("subagent") { subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace) subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) + // Set model resolver so spawn can use target agent's model + subagentManager.SetModelResolver(func(targetAgentID string) string { + if targetAgent, ok := registry.GetAgent(targetAgentID); ok { + return targetAgent.Model + } + return "" + }) spawnTool := tools.NewSpawnTool(subagentManager) currentAgentID := agentID spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { diff --git a/pkg/tools/spawn_test.go b/pkg/tools/spawn_test.go index 43223b8db..e2939f9ba 100644 --- a/pkg/tools/spawn_test.go +++ b/pkg/tools/spawn_test.go @@ -77,3 +77,38 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) { t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM) } } + +func TestSubagentManager_ModelResolver(t *testing.T) { + provider := &MockLLMProvider{} + manager := NewSubagentManager(provider, "default-model", "/tmp/test") + + // Set up model resolver + resolvedAgentID := "" + manager.SetModelResolver(func(agentID string) string { + resolvedAgentID = agentID + if agentID == "premium-agent" { + return "gpt-4" + } + return "" + }) + + // Verify resolver is set + if manager.modelResolver == nil { + t.Fatal("Model resolver should be set") + } + + // Test resolver is called with correct agent ID + result := manager.modelResolver("premium-agent") + if resolvedAgentID != "premium-agent" { + t.Errorf("Expected resolver to be called with 'premium-agent', got '%s'", resolvedAgentID) + } + if result != "gpt-4" { + t.Errorf("Expected 'gpt-4', got '%s'", result) + } + + // Test fallback for unknown agent + result = manager.modelResolver("unknown-agent") + if result != "" { + t.Errorf("Expected empty string for unknown agent, got '%s'", result) + } +} diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index e51cbaafa..78ee3afa7 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -34,6 +34,9 @@ type SubagentManager struct { hasMaxTokens bool hasTemperature bool nextID int + // modelResolver resolves agentID to model name. + // Returns empty string if agent not found (falls back to defaultModel). + modelResolver func(agentID string) string } func NewSubagentManager( @@ -61,6 +64,16 @@ func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) { sm.hasTemperature = true } +// SetModelResolver sets a function to resolve agentID to model name. +// When spawn is called with agent_id, this resolver is used to get the +// target agent's configured model. If the resolver returns empty string +// or is not set, falls back to the defaultModel. +func (sm *SubagentManager) SetModelResolver(resolver func(agentID string) string) { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.modelResolver = resolver +} + // SetTools sets the tool registry for subagent execution. // If not set, subagent will have access to the provided tools. func (sm *SubagentManager) SetTools(tools *ToolRegistry) { @@ -147,8 +160,18 @@ After completing the task, provide a clear summary of what was done.` temperature := sm.temperature hasMaxTokens := sm.hasMaxTokens hasTemperature := sm.hasTemperature + modelResolver := sm.modelResolver + defaultModel := sm.defaultModel sm.mu.RUnlock() + // Resolve target agent model if agentID is specified + model := defaultModel + if task.AgentID != "" && modelResolver != nil { + if resolvedModel := modelResolver(task.AgentID); resolvedModel != "" { + model = resolvedModel + } + } + var llmOptions map[string]any if hasMaxTokens || hasTemperature { llmOptions = map[string]any{} @@ -162,7 +185,7 @@ After completing the task, provide a clear summary of what was done.` loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ Provider: sm.provider, - Model: sm.defaultModel, + Model: model, Tools: tools, MaxIterations: maxIter, LLMOptions: llmOptions,