fix(subturn): preserve spawn/subagent tool snapshots in child turns

This commit is contained in:
Alix-007 2026-03-26 03:45:17 +08:00
parent ed618e14aa
commit 19979e731f
6 changed files with 219 additions and 5 deletions

View file

@ -197,6 +197,16 @@ func (al *AgentLoop) generateSubTurnID() string {
return fmt.Sprintf("subturn-%d", al.subTurnCounter.Add(1)) return fmt.Sprintf("subturn-%d", al.subTurnCounter.Add(1))
} }
func toolRegistryFromSlice(toolSlice []tools.Tool) *tools.ToolRegistry {
registry := tools.NewToolRegistry()
for _, tool := range toolSlice {
if tool != nil {
registry.Register(tool)
}
}
return registry
}
// ====================== Core Function: spawnSubTurn ====================== // ====================== Core Function: spawnSubTurn ======================
// AgentLoopSpawner implements tools.SubTurnSpawner interface. // AgentLoopSpawner implements tools.SubTurnSpawner interface.
@ -344,10 +354,14 @@ func spawnSubTurn(
ephemeralStore := newEphemeralSession(nil) ephemeralStore := newEphemeralSession(nil)
agent := *baseAgent // shallow copy agent := *baseAgent // shallow copy
agent.Sessions = ephemeralStore agent.Sessions = ephemeralStore
// Clone the tool registry so child turn's tool registrations if cfg.Tools != nil {
// don't pollute the parent's registry. // Explicit tool slice from caller has highest priority for sub-turns.
if baseAgent.Tools != nil { agent.Tools = toolRegistryFromSlice(cfg.Tools)
} else if baseAgent.Tools != nil {
// Otherwise inherit the parent's tool snapshot.
agent.Tools = baseAgent.Tools.Clone() agent.Tools = baseAgent.Tools.Clone()
} else {
agent.Tools = tools.NewToolRegistry()
} }
// Create processOptions for the child turn // Create processOptions for the child turn

View file

@ -933,6 +933,89 @@ func (m *simpleMockProviderAPI) GetDefaultModel() string {
return "gpt-4o-mini" return "gpt-4o-mini"
} }
type toolCaptureProvider struct {
lastToolNames []string
}
func (p *toolCaptureProvider) Chat(
ctx context.Context,
messages []providers.Message,
toolDefs []providers.ToolDefinition,
model string,
options map[string]any,
) (*providers.LLMResponse, error) {
p.lastToolNames = p.lastToolNames[:0]
for _, td := range toolDefs {
p.lastToolNames = append(p.lastToolNames, td.Function.Name)
}
return &providers.LLMResponse{Content: "ok"}, nil
}
func (p *toolCaptureProvider) GetDefaultModel() string {
return "test-model"
}
type subturnProbeTool struct{}
func (t *subturnProbeTool) Name() string { return "subturn_probe_tool" }
func (t *subturnProbeTool) Description() string { return "subturn probe tool" }
func (t *subturnProbeTool) Parameters() map[string]any {
return map[string]any{
"type": "object",
}
}
func (t *subturnProbeTool) Execute(ctx context.Context, args map[string]any) *tools.ToolResult {
return tools.SilentResult("ok")
}
func TestSpawnSubTurn_UsesExplicitConfigTools(t *testing.T) {
provider := &toolCaptureProvider{}
cfg := &config.Config{
Agents: config.AgentsConfig{
Defaults: config.AgentDefaults{
Workspace: t.TempDir(),
ModelName: "test-model",
MaxTokens: 4096,
MaxToolIterations: 10,
},
},
}
al := NewAgentLoop(cfg, bus.NewMessageBus(), provider)
parentAgent := al.registry.GetDefaultAgent()
if parentAgent == nil {
t.Fatal("expected default agent")
}
// Keep parent tools empty so child tools can only come from cfg.Tools.
parentAgent.Tools = tools.NewToolRegistry()
parent := &turnState{
ctx: context.Background(),
turnID: "parent-explicit-tools",
depth: 0,
pendingResults: make(chan *tools.ToolResult, 1),
concurrencySem: make(chan struct{}, testMaxConcurrentSubTurns),
session: &ephemeralSessionStore{},
agent: parentAgent,
}
_, err := spawnSubTurn(context.Background(), al, parent, SubTurnConfig{
Model: "test-model",
SystemPrompt: "run task",
Tools: []tools.Tool{&subturnProbeTool{}},
})
if err != nil {
t.Fatalf("spawnSubTurn returned error: %v", err)
}
if len(provider.lastToolNames) != 1 || provider.lastToolNames[0] != "subturn_probe_tool" {
t.Fatalf("expected explicit cfg tool to be sent, got %v", provider.lastToolNames)
}
}
// TestGetActiveTurn verifies that GetActiveTurn returns correct turn information // TestGetActiveTurn verifies that GetActiveTurn returns correct turn information
func TestGetActiveTurn(t *testing.T) { func TestGetActiveTurn(t *testing.T) {
cfg := &config.Config{ cfg := &config.Config{

View file

@ -11,6 +11,7 @@ type SpawnTool struct {
defaultModel string defaultModel string
maxTokens int maxTokens int
temperature float64 temperature float64
tools *ToolRegistry
allowlistCheck func(targetAgentID string) bool allowlistCheck func(targetAgentID string) bool
} }
@ -25,6 +26,7 @@ func NewSpawnTool(manager *SubagentManager) *SpawnTool {
defaultModel: manager.defaultModel, defaultModel: manager.defaultModel,
maxTokens: manager.maxTokens, maxTokens: manager.maxTokens,
temperature: manager.temperature, temperature: manager.temperature,
tools: manager.tools,
} }
} }
@ -124,7 +126,7 @@ Task: %s`,
go func() { go func() {
result, err := t.spawner.SpawnSubTurn(ctx, SubTurnConfig{ result, err := t.spawner.SpawnSubTurn(ctx, SubTurnConfig{
Model: t.defaultModel, Model: t.defaultModel,
Tools: nil, // Will inherit from parent via context Tools: snapshotTools(t.tools),
SystemPrompt: systemPrompt, SystemPrompt: systemPrompt,
MaxTokens: t.maxTokens, MaxTokens: t.maxTokens,
Temperature: t.temperature, Temperature: t.temperature,

View file

@ -4,6 +4,7 @@ import (
"context" "context"
"strings" "strings"
"testing" "testing"
"time"
) )
// mockSpawner implements SubTurnSpawner for testing // mockSpawner implements SubTurnSpawner for testing
@ -24,6 +25,48 @@ func (m *mockSpawner) SpawnSubTurn(ctx context.Context, cfg SubTurnConfig) (*Too
}, nil }, nil
} }
type managerSnapshotTool struct {
name string
}
func (t *managerSnapshotTool) Name() string {
return t.name
}
func (t *managerSnapshotTool) Description() string {
return "test tool"
}
func (t *managerSnapshotTool) Parameters() map[string]any {
return map[string]any{
"type": "object",
}
}
func (t *managerSnapshotTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
return SilentResult("ok")
}
type recordingSpawner struct {
toolNames []string
done chan struct{}
}
func (s *recordingSpawner) SpawnSubTurn(ctx context.Context, cfg SubTurnConfig) (*ToolResult, error) {
for _, tool := range cfg.Tools {
if tool != nil {
s.toolNames = append(s.toolNames, tool.Name())
}
}
if s.done != nil {
close(s.done)
}
return &ToolResult{
ForLLM: "Task completed",
ForUser: "Task completed",
}, nil
}
func TestSpawnTool_Execute_EmptyTask(t *testing.T) { func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
provider := &MockLLMProvider{} provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test") manager := NewSubagentManager(provider, "test-model", "/tmp/test")
@ -96,3 +139,34 @@ 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_Execute_PassesManagerToolsToSubTurn(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager.RegisterTool(&managerSnapshotTool{name: "snapshot_tool"})
tool := NewSpawnTool(manager)
spawner := &recordingSpawner{done: make(chan struct{})}
tool.SetSpawner(spawner)
result := tool.Execute(context.Background(), map[string]any{"task": "inspect tools"})
if result == nil {
t.Fatal("Result should not be nil")
}
if result.IsError {
t.Fatalf("Expected success, got error: %s", result.ForLLM)
}
if !result.Async {
t.Fatal("spawn result should be async")
}
select {
case <-spawner.done:
case <-time.After(2 * time.Second):
t.Fatal("timed out waiting for async spawn execution")
}
if len(spawner.toolNames) != 1 || spawner.toolNames[0] != "snapshot_tool" {
t.Fatalf("expected tool snapshot [snapshot_tool], got %v", spawner.toolNames)
}
}

View file

@ -115,6 +115,21 @@ func (sm *SubagentManager) RegisterTool(tool Tool) {
sm.tools.Register(tool) sm.tools.Register(tool)
} }
func snapshotTools(toolsRegistry *ToolRegistry) []Tool {
if toolsRegistry == nil {
return nil
}
names := toolsRegistry.List()
snapshot := make([]Tool, 0, len(names))
for _, name := range names {
if tool, ok := toolsRegistry.Get(name); ok {
snapshot = append(snapshot, tool)
}
}
return snapshot
}
func (sm *SubagentManager) Spawn( func (sm *SubagentManager) Spawn(
ctx context.Context, ctx context.Context,
task, label, agentID, originChannel, originChatID string, task, label, agentID, originChannel, originChatID string,
@ -322,6 +337,7 @@ type SubagentTool struct {
defaultModel string defaultModel string
maxTokens int maxTokens int
temperature float64 temperature float64
tools *ToolRegistry
} }
func NewSubagentTool(manager *SubagentManager) *SubagentTool { func NewSubagentTool(manager *SubagentManager) *SubagentTool {
@ -332,6 +348,7 @@ func NewSubagentTool(manager *SubagentManager) *SubagentTool {
defaultModel: manager.defaultModel, defaultModel: manager.defaultModel,
maxTokens: manager.maxTokens, maxTokens: manager.maxTokens,
temperature: manager.temperature, temperature: manager.temperature,
tools: manager.tools,
} }
} }
@ -395,7 +412,7 @@ Task: %s`,
if t.spawner != nil { if t.spawner != nil {
result, err := t.spawner.SpawnSubTurn(ctx, SubTurnConfig{ result, err := t.spawner.SpawnSubTurn(ctx, SubTurnConfig{
Model: t.defaultModel, Model: t.defaultModel,
Tools: nil, // Will inherit from parent via context Tools: snapshotTools(t.tools),
SystemPrompt: systemPrompt, SystemPrompt: systemPrompt,
MaxTokens: t.maxTokens, MaxTokens: t.maxTokens,
Temperature: t.temperature, Temperature: t.temperature,

View file

@ -324,3 +324,27 @@ func TestSubagentTool_ForUserTruncation(t *testing.T) {
t.Error("ForLLM should contain reference to original task") t.Error("ForLLM should contain reference to original task")
} }
} }
func TestSubagentTool_Execute_PassesManagerToolsToSubTurn(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager.RegisterTool(&managerSnapshotTool{name: "snapshot_tool"})
tool := NewSubagentTool(manager)
spawner := &recordingSpawner{}
tool.SetSpawner(spawner)
result := tool.Execute(context.Background(), map[string]any{
"task": "inspect tools",
})
if result == nil {
t.Fatal("Result should not be nil")
}
if result.IsError {
t.Fatalf("Expected success, got error: %s", result.ForLLM)
}
if len(spawner.toolNames) != 1 || spawner.toolNames[0] != "snapshot_tool" {
t.Fatalf("expected tool snapshot [snapshot_tool], got %v", spawner.toolNames)
}
}