fix(spawn): use target agent model

This commit is contained in:
Alix-007 2026-03-17 21:42:04 +08:00
parent 8a188cf7fc
commit 4d6ac5fc16
4 changed files with 104 additions and 1 deletions

View file

@ -244,6 +244,13 @@ func registerSharedTools(
// spawn_status which are added below — preventing recursive // spawn_status which are added below — preventing recursive
// subagent spawning. // subagent spawning.
subagentManager.SetTools(agent.Tools.Clone()) subagentManager.SetTools(agent.Tools.Clone())
subagentManager.SetAgentModelResolver(func(targetAgentID string) (string, bool) {
targetAgent, ok := registry.GetAgent(targetAgentID)
if !ok || targetAgent == nil || targetAgent.Model == "" {
return "", false
}
return targetAgent.Model, true
})
if spawnEnabled { if spawnEnabled {
spawnTool := tools.NewSpawnTool(subagentManager) spawnTool := tools.NewSpawnTool(subagentManager)
currentAgentID := agentID currentAgentID := agentID

View file

@ -4,6 +4,7 @@ import (
"context" "context"
"strings" "strings"
"testing" "testing"
"time"
) )
func TestSpawnTool_Execute_EmptyTask(t *testing.T) { func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
@ -77,3 +78,39 @@ func TestSpawnTool_Execute_NilManager(t *testing.T) {
t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM) t.Errorf("Error message should mention manager not configured, got: %s", result.ForLLM)
} }
} }
func TestSpawnTool_ExecuteAsync_UsesTargetAgentModel(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "caller-model", "/tmp/test")
manager.SetAgentModelResolver(func(agentID string) (string, bool) {
if agentID == "analyst" {
return "premium-model", true
}
return "", false
})
tool := NewSpawnTool(manager)
done := make(chan *ToolResult, 1)
result := tool.ExecuteAsync(context.Background(), map[string]any{
"task": "Investigate the issue",
"agent_id": "analyst",
}, func(_ context.Context, result *ToolResult) {
done <- result
})
if result == nil || result.IsError {
t.Fatalf("expected async spawn success, got: %+v", result)
}
select {
case callbackResult := <-done:
if callbackResult == nil || callbackResult.IsError {
t.Fatalf("expected successful callback result, got: %+v", callbackResult)
}
case <-time.After(time.Second):
t.Fatal("timed out waiting for async spawn callback")
}
if provider.lastModel != "premium-model" {
t.Fatalf("model = %q, want %q", provider.lastModel, "premium-model")
}
}

View file

@ -26,6 +26,7 @@ type SubagentManager struct {
mu sync.RWMutex mu sync.RWMutex
provider providers.LLMProvider provider providers.LLMProvider
defaultModel string defaultModel string
resolveModel func(agentID string) (string, bool)
workspace string workspace string
tools *ToolRegistry tools *ToolRegistry
maxIterations int maxIterations int
@ -61,6 +62,13 @@ func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) {
sm.hasTemperature = true sm.hasTemperature = true
} }
// SetAgentModelResolver resolves the effective model for a target agent ID.
func (sm *SubagentManager) SetAgentModelResolver(resolve func(agentID string) (string, bool)) {
sm.mu.Lock()
defer sm.mu.Unlock()
sm.resolveModel = resolve
}
// SetTools sets the tool registry for subagent execution. // SetTools sets the tool registry for subagent execution.
// If not set, subagent will have access to the provided tools. // If not set, subagent will have access to the provided tools.
func (sm *SubagentManager) SetTools(tools *ToolRegistry) { func (sm *SubagentManager) SetTools(tools *ToolRegistry) {
@ -144,8 +152,16 @@ After completing the task, provide a clear summary of what was done.`
temperature := sm.temperature temperature := sm.temperature
hasMaxTokens := sm.hasMaxTokens hasMaxTokens := sm.hasMaxTokens
hasTemperature := sm.hasTemperature hasTemperature := sm.hasTemperature
resolveModel := sm.resolveModel
sm.mu.RUnlock() sm.mu.RUnlock()
model := sm.defaultModel
if task.AgentID != "" && resolveModel != nil {
if resolvedModel, ok := resolveModel(task.AgentID); ok && resolvedModel != "" {
model = resolvedModel
}
}
var llmOptions map[string]any var llmOptions map[string]any
if hasMaxTokens || hasTemperature { if hasMaxTokens || hasTemperature {
llmOptions = map[string]any{} llmOptions = map[string]any{}
@ -159,7 +175,7 @@ After completing the task, provide a clear summary of what was done.`
loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ loopResult, err := RunToolLoop(ctx, ToolLoopConfig{
Provider: sm.provider, Provider: sm.provider,
Model: sm.defaultModel, Model: model,
Tools: tools, Tools: tools,
MaxIterations: maxIter, MaxIterations: maxIter,
LLMOptions: llmOptions, LLMOptions: llmOptions,

View file

@ -4,6 +4,7 @@ import (
"context" "context"
"strings" "strings"
"testing" "testing"
"time"
"github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/providers"
) )
@ -11,6 +12,7 @@ import (
// MockLLMProvider is a test implementation of LLMProvider // MockLLMProvider is a test implementation of LLMProvider
type MockLLMProvider struct { type MockLLMProvider struct {
lastOptions map[string]any lastOptions map[string]any
lastModel string
} }
func (m *MockLLMProvider) Chat( func (m *MockLLMProvider) Chat(
@ -21,6 +23,7 @@ func (m *MockLLMProvider) Chat(
options map[string]any, options map[string]any,
) (*providers.LLMResponse, error) { ) (*providers.LLMResponse, error) {
m.lastOptions = options m.lastOptions = options
m.lastModel = model
// Find the last user message to generate a response // Find the last user message to generate a response
for i := len(messages) - 1; i >= 0; i-- { for i := len(messages) - 1; i >= 0; i-- {
if messages[i].Role == "user" { if messages[i].Role == "user" {
@ -69,6 +72,46 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
} }
} }
func TestSubagentManager_RunTask_UsesResolvedTargetAgentModel(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "caller-model", "/tmp/test")
manager.SetAgentModelResolver(func(agentID string) (string, bool) {
if agentID == "analyst" {
return "premium-model", true
}
return "", false
})
done := make(chan *ToolResult, 1)
_, err := manager.Spawn(
context.Background(),
"Investigate the issue",
"analysis",
"analyst",
"cli",
"direct",
func(_ context.Context, result *ToolResult) {
done <- result
},
)
if err != nil {
t.Fatalf("Spawn() error: %v", err)
}
select {
case result := <-done:
if result == nil || result.IsError {
t.Fatalf("expected successful async result, got: %+v", result)
}
case <-time.After(time.Second):
t.Fatal("timed out waiting for subagent completion")
}
if provider.lastModel != "premium-model" {
t.Fatalf("model = %q, want %q", provider.lastModel, "premium-model")
}
}
// TestSubagentTool_Name verifies tool name // TestSubagentTool_Name verifies tool name
func TestSubagentTool_Name(t *testing.T) { func TestSubagentTool_Name(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}