From 375e64bba55a7c1e99270bfe7b433982185dcef9 Mon Sep 17 00:00:00 2001 From: Administrator <1280842908@qq.com> Date: Tue, 24 Mar 2026 10:21:58 +0800 Subject: [PATCH] feat(team): add DAG execution tests and improve mock provider for concurrency --- pkg/tools/team.go | 7 +- pkg/tools/team_test.go | 384 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 390 insertions(+), 1 deletion(-) diff --git a/pkg/tools/team.go b/pkg/tools/team.go index eb7b8873c..13e97c9fe 100644 --- a/pkg/tools/team.go +++ b/pkg/tools/team.go @@ -899,6 +899,7 @@ func (t *TeamTool) executeDAG(ctx context.Context, cancel context.CancelFunc, ba // 3. Initialize queue with nodes having 0 in-degree nodesToProcess := len(members) + completedNodes := 0 for id, deg := range inDegree { if deg == 0 { readyChan <- id @@ -914,7 +915,8 @@ func (t *TeamTool) executeDAG(ctx context.Context, cancel context.CancelFunc, ba var finalResultsMu sync.Mutex // 4. DAG Execution Loop - for i := 0; i < nodesToProcess; i++ { + // Continue until all nodes have completed + for completedNodes < nodesToProcess { select { case <-ctx.Done(): return ErrorResult("DAG execution timed out or cancelled") @@ -987,6 +989,9 @@ func (t *TeamTool) executeDAG(ctx context.Context, cancel context.CancelFunc, ba return ErrorResult(res.err.Error()) } + // Mark this node as completed + completedNodes++ + // Update dependents for _, dependentID := range graph[res.id] { contextMu.Lock() diff --git a/pkg/tools/team_test.go b/pkg/tools/team_test.go index 1fa3224fa..232eb3903 100644 --- a/pkg/tools/team_test.go +++ b/pkg/tools/team_test.go @@ -2,6 +2,7 @@ package tools import ( "context" + "sync" "testing" "github.com/sipeed/picoclaw/pkg/config" @@ -192,6 +193,34 @@ func (m *mockProvider) GetDefaultModel() string { return "mock-model" } +// mockProviderWithID returns responses based on the role in the system message +// This is needed for DAG tests where execution order is non-deterministic +type mockProviderWithID struct { + responses map[string]string + mu sync.Mutex + callCount int +} + +func (m *mockProviderWithID) Chat(ctx context.Context, messages []providers.Message, tools []providers.ToolDefinition, model string, options map[string]any) (*providers.LLMResponse, error) { + m.mu.Lock() + defer m.mu.Unlock() + m.callCount++ + + // Extract role from system message to determine which response to return + for _, msg := range messages { + if msg.Role == "system" { + if resp, ok := m.responses[msg.Content]; ok { + return &providers.LLMResponse{Content: resp}, nil + } + } + } + return &providers.LLMResponse{Content: "Default response"}, nil +} + +func (m *mockProviderWithID) GetDefaultModel() string { + return "mock-model" +} + func TestExecuteSequential(t *testing.T) { // 1. Setup mock provider to return specific outputs for each agent mock := &mockProvider{ @@ -225,3 +254,358 @@ func TestExecuteSequential(t *testing.T) { assert.Contains(t, result.ForLLM, "Derived result from Agent B") assert.Equal(t, 2, mock.callCount, "Should have called mock provider exactly twice") } + +func TestExecuteDAG(t *testing.T) { + t.Run("Simple DAG with dependencies", func(t *testing.T) { + // Setup: A -> C, B -> C (C depends on both A and B) + mock := &mockProviderWithID{ + responses: map[string]string{ + "Data Collector A": "Data from A", + "Data Collector B": "Data from B", + "Aggregator": "Combined result from C using A and B", + }, + } + + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{}, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: mock, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), // Empty registry for test + } + + members := []TeamMember{ + {ID: "A", Role: "Data Collector A", Task: "Collect data A"}, + {ID: "B", Role: "Data Collector B", Task: "Collect data B"}, + {ID: "C", Role: "Aggregator", Task: "Combine data", DependsOn: []string{"A", "B"}}, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + result := tool.executeDAG(ctx, cancel, baseConfig, members, 1000) + + assert.False(t, result.IsError, "Should not return error") + assert.Contains(t, result.ForLLM, "Data from A") + assert.Contains(t, result.ForLLM, "Data from B") + assert.Contains(t, result.ForLLM, "Combined result from C") + assert.Equal(t, 3, mock.callCount, "Should have called mock provider 3 times") + }) + + t.Run("Linear chain DAG", func(t *testing.T) { + // Setup: A -> B -> C (linear dependency chain) + mock := &mockProviderWithID{ + responses: map[string]string{ + "Step 1": "Step 1 output", + "Step 2": "Step 2 output", + "Step 3": "Step 3 output", + }, + } + + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{}, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: mock, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), + } + + members := []TeamMember{ + {ID: "A", Role: "Step 1", Task: "Do step 1"}, + {ID: "B", Role: "Step 2", Task: "Do step 2", DependsOn: []string{"A"}}, + {ID: "C", Role: "Step 3", Task: "Do step 3", DependsOn: []string{"B"}}, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + result := tool.executeDAG(ctx, cancel, baseConfig, members, 1000) + + assert.False(t, result.IsError) + assert.Contains(t, result.ForLLM, "Step 1 output") + assert.Contains(t, result.ForLLM, "Step 2 output") + assert.Contains(t, result.ForLLM, "Step 3 output") + }) + + t.Run("Detect circular dependency", func(t *testing.T) { + // Setup: A -> B -> C -> A (circular) + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{}, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: &mockProvider{}, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), + } + + members := []TeamMember{ + {ID: "A", Role: "Worker A", Task: "Task A", DependsOn: []string{"C"}}, + {ID: "B", Role: "Worker B", Task: "Task B", DependsOn: []string{"A"}}, + {ID: "C", Role: "Worker C", Task: "Task C", DependsOn: []string{"B"}}, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + result := tool.executeDAG(ctx, cancel, baseConfig, members, 1000) + + assert.True(t, result.IsError, "Should detect circular dependency") + assert.Contains(t, result.ForLLM, "cycle") + }) + + t.Run("Detect undefined dependency", func(t *testing.T) { + // Setup: A depends on non-existent "X" + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{}, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: &mockProvider{}, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), + } + + members := []TeamMember{ + {ID: "A", Role: "Worker A", Task: "Task A", DependsOn: []string{"X"}}, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + result := tool.executeDAG(ctx, cancel, baseConfig, members, 1000) + + assert.True(t, result.IsError, "Should detect undefined dependency") + assert.Contains(t, result.ForLLM, "undefined member") + }) + + t.Run("Complex DAG with multiple roots", func(t *testing.T) { + // Setup: A, B (roots) -> C, D -> E + mock := &mockProviderWithID{ + responses: map[string]string{ + "Root A": "Root A output", + "Root B": "Root B output", + "Worker C": "C output", + "Worker D": "D output", + "Final": "Final E output", + }, + } + + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{}, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: mock, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), + } + + members := []TeamMember{ + {ID: "A", Role: "Root A", Task: "Root task A"}, + {ID: "B", Role: "Root B", Task: "Root task B"}, + {ID: "C", Role: "Worker C", Task: "Task C", DependsOn: []string{"A"}}, + {ID: "D", Role: "Worker D", Task: "Task D", DependsOn: []string{"B"}}, + {ID: "E", Role: "Final", Task: "Final task", DependsOn: []string{"C", "D"}}, + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + result := tool.executeDAG(ctx, cancel, baseConfig, members, 1000) + + assert.False(t, result.IsError) + assert.Contains(t, result.ForLLM, "Root A output") + assert.Contains(t, result.ForLLM, "Root B output") + assert.Contains(t, result.ForLLM, "Final E output") + assert.Equal(t, 5, mock.callCount) + }) +} + +func TestExecuteEvaluatorOptimizer(t *testing.T) { + t.Run("Pass on first attempt", func(t *testing.T) { + mock := &mockProvider{ + responses: []string{ + "Perfect code implementation", + "[PASS] The code is correct", + }, + } + + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{ + MaxEvaluatorLoops: 3, + }, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: mock, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), + } + + members := []TeamMember{ + {ID: "worker", Role: "Coder", Task: "Write a function"}, + {ID: "evaluator", Role: "Code Reviewer", Task: "Review the code"}, + } + + result := tool.executeEvaluatorOptimizer(context.Background(), baseConfig, members, 1000) + + assert.False(t, result.IsError) + assert.Contains(t, result.ForLLM, "Perfect code implementation") + assert.Contains(t, result.ForLLM, "[PASS]") + assert.Contains(t, result.ForUser, "passed on attempt 1") + assert.Equal(t, 2, mock.callCount, "Should call worker once and evaluator once") + }) + + t.Run("Pass on second attempt after feedback", func(t *testing.T) { + mock := &mockProvider{ + responses: []string{ + "Initial code with bug", + "Missing error handling", + "Fixed code with error handling", + "[PASS] Now it's correct", + }, + } + + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{ + MaxEvaluatorLoops: 3, + }, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: mock, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), + } + + members := []TeamMember{ + {ID: "worker", Role: "Coder", Task: "Write a function"}, + {ID: "evaluator", Role: "Code Reviewer", Task: "Review the code"}, + } + + result := tool.executeEvaluatorOptimizer(context.Background(), baseConfig, members, 1000) + + assert.False(t, result.IsError) + assert.Contains(t, result.ForLLM, "Initial code with bug") + assert.Contains(t, result.ForLLM, "Missing error handling") + assert.Contains(t, result.ForLLM, "Fixed code with error handling") + assert.Contains(t, result.ForLLM, "[PASS]") + assert.Contains(t, result.ForUser, "passed on attempt 2") + assert.Equal(t, 4, mock.callCount, "Should call worker twice and evaluator twice") + }) + + t.Run("Exhaust max loops without pass", func(t *testing.T) { + mock := &mockProvider{ + responses: []string{ + "Attempt 1", + "Still has issues", + "Attempt 2", + "Still not good", + "Attempt 3", + "Still failing", + }, + } + + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{ + MaxEvaluatorLoops: 3, + }, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: mock, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), + } + + members := []TeamMember{ + {ID: "worker", Role: "Coder", Task: "Write a function"}, + {ID: "evaluator", Role: "Code Reviewer", Task: "Review the code"}, + } + + result := tool.executeEvaluatorOptimizer(context.Background(), baseConfig, members, 1000) + + assert.False(t, result.IsError, "Should not error, just report exhaustion") + assert.Contains(t, result.ForLLM, "Maximum evaluation loops reached") + assert.Contains(t, result.ForUser, "exhausted 3 attempts") + assert.Equal(t, 6, mock.callCount, "Should call worker 3 times and evaluator 3 times") + }) + + t.Run("Require exactly two members", func(t *testing.T) { + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{}, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: &mockProvider{}, + Model: "mock-model", + MaxIterations: 1, + } + + // Test with 1 member + members := []TeamMember{ + {ID: "worker", Role: "Coder", Task: "Write a function"}, + } + + result := tool.executeEvaluatorOptimizer(context.Background(), baseConfig, members, 1000) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "exactly two members") + + // Test with 3 members + members = []TeamMember{ + {ID: "worker", Role: "Coder", Task: "Write a function"}, + {ID: "evaluator", Role: "Reviewer", Task: "Review"}, + {ID: "extra", Role: "Extra", Task: "Extra task"}, + } + + result = tool.executeEvaluatorOptimizer(context.Background(), baseConfig, members, 1000) + assert.True(t, result.IsError) + assert.Contains(t, result.ForLLM, "exactly two members") + }) + + t.Run("Stateful worker memory across iterations", func(t *testing.T) { + // This test verifies that the worker's message history is preserved + // across iterations, allowing it to "remember" previous feedback + mockWithMemory := &mockProvider{ + responses: []string{ + "First attempt", + "Needs improvement: add validation", + "Second attempt with validation", + "[PASS] Good now", + }, + } + + manager := NewSubagentManager(nil, "mock-model", nil, "", config.TeamToolsConfig{ + MaxEvaluatorLoops: 3, + }, nil) + tool := NewTeamTool(manager, &config.Config{}) + + baseConfig := ToolLoopConfig{ + Provider: mockWithMemory, + Model: "mock-model", + MaxIterations: 1, + Tools: NewToolRegistry(), + } + + members := []TeamMember{ + {ID: "worker", Role: "Coder", Task: "Write validation logic"}, + {ID: "evaluator", Role: "Reviewer", Task: "Check validation"}, + } + + result := tool.executeEvaluatorOptimizer(context.Background(), baseConfig, members, 1000) + + assert.False(t, result.IsError) + assert.Contains(t, result.ForLLM, "First attempt") + assert.Contains(t, result.ForLLM, "Needs improvement") + assert.Contains(t, result.ForLLM, "Second attempt with validation") + assert.Contains(t, result.ForLLM, "[PASS]") + + // Verify the worker was called twice (once initially, once after feedback) + // and evaluator was called twice + assert.Equal(t, 4, mockWithMemory.callCount) + }) +}