fix(subturn): preserve spawn/subagent tool snapshots in child turns
This commit is contained in:
parent
ed618e14aa
commit
19979e731f
6 changed files with 219 additions and 5 deletions
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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{
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue