From e7de26ee1be658e8ad21f44de4e11219edde9e3f Mon Sep 17 00:00:00 2001 From: Administrator <1280842908@qq.com> Date: Sat, 28 Feb 2026 13:45:47 +0800 Subject: [PATCH] feat(teams): implement Agent Teams architecture - Added 'team' and 'spawn_sub_agent' tools to support Coordinator-Worker patterns. - Implemented execution strategies: Sequential, Parallel, DAG, and Evaluator-Optimizer. - Added global token budget tracking via 'RemainingTokenBudget' for team cost control. - Designed 'ConcurrencyUpgradeable' interface and 'ConcurrentFS' wrapper for opt-in, thread-safe file operations during parallel/DAG runs. - Tested file locking mechanisms preventing race conditions during high concurrency. --- pkg/agent/loop.go | 6 + pkg/tools/edit.go | 44 ++- pkg/tools/filesystem.go | 135 +++++++++ pkg/tools/filesystem_test.go | 58 ++++ pkg/tools/registry.go | 11 + pkg/tools/spawn_sub_agent.go | 105 +++++++ pkg/tools/subagent.go | 25 ++ pkg/tools/team.go | 571 +++++++++++++++++++++++++++++++++++ pkg/tools/team_test.go | 62 ++++ pkg/tools/toolloop.go | 26 +- 10 files changed, 1015 insertions(+), 28 deletions(-) create mode 100644 pkg/tools/spawn_sub_agent.go create mode 100644 pkg/tools/team.go create mode 100644 pkg/tools/team_test.go diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 29827d0b2..1b8654eb6 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -155,6 +155,12 @@ func registerSharedTools( return registry.CanSpawnSubagent(currentAgentID, targetAgentID) }) agent.Tools.Register(spawnTool) + + teamTool := tools.NewTeamTool(subagentManager) + agent.Tools.Register(teamTool) + + spawnSubAgentTool := tools.NewSpawnSubAgentTool(subagentManager) + agent.Tools.Register(spawnSubAgentTool) } } diff --git a/pkg/tools/edit.go b/pkg/tools/edit.go index d3ab267bf..e1d6d6821 100644 --- a/pkg/tools/edit.go +++ b/pkg/tools/edit.go @@ -2,9 +2,7 @@ package tools import ( "context" - "errors" "fmt" - "io/fs" "strings" ) @@ -76,6 +74,12 @@ func (t *EditFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe return SilentResult(fmt.Sprintf("File edited: %s", path)) } +func (t *EditFileTool) UpgradeToConcurrent() Tool { + return &EditFileTool{ + fs: &ConcurrentFS{baseFS: t.fs}, + } +} + type AppendFileTool struct { fs fileSystem } @@ -132,31 +136,25 @@ func (t *AppendFileTool) Execute(ctx context.Context, args map[string]any) *Tool return SilentResult(fmt.Sprintf("Appended to %s", path)) } -// editFile reads the file via sysFs, performs the replacement, and writes back. -// It uses a fileSystem interface, allowing the same logic for both restricted and unrestricted modes. -func editFile(sysFs fileSystem, path, oldText, newText string) error { - content, err := sysFs.ReadFile(path) - if err != nil { - return err +func (t *AppendFileTool) UpgradeToConcurrent() Tool { + return &AppendFileTool{ + fs: &ConcurrentFS{baseFS: t.fs}, } - - newContent, err := replaceEditContent(content, oldText, newText) - if err != nil { - return err - } - - return sysFs.WriteFile(path, newContent) } -// appendFile reads the existing content (if any) via sysFs, appends new content, and writes back. -func appendFile(sysFs fileSystem, path, appendContent string) error { - content, err := sysFs.ReadFile(path) - if err != nil && !errors.Is(err, fs.ErrNotExist) { - return err - } +// editFile reads the file via sysFs, performs the replacement, and writes back atomically. +func editFile(sysFs fileSystem, path, oldText, newText string) error { + return sysFs.EditFile(path, func(content []byte) ([]byte, error) { + return replaceEditContent(content, oldText, newText) + }) +} - newContent := append(content, []byte(appendContent)...) - return sysFs.WriteFile(path, newContent) +// appendFile reads the existing content (if any) via sysFs, appends new content, and writes back atomically. +func appendFile(sysFs fileSystem, path, appendContent string) error { + return sysFs.EditFile(path, func(content []byte) ([]byte, error) { + newContent := append(content, []byte(appendContent)...) + return newContent, nil + }) } // replaceEditContent handles the core logic of finding and replacing a single occurrence of oldText. diff --git a/pkg/tools/filesystem.go b/pkg/tools/filesystem.go index 03d461dcc..c9c99cfe5 100644 --- a/pkg/tools/filesystem.go +++ b/pkg/tools/filesystem.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "strings" + "sync" "time" "github.com/sipeed/picoclaw/pkg/fileutil" @@ -131,6 +132,12 @@ func (t *ReadFileTool) Execute(ctx context.Context, args map[string]any) *ToolRe return NewToolResult(string(content)) } +func (t *ReadFileTool) UpgradeToConcurrent() Tool { + return &ReadFileTool{ + fs: &ConcurrentFS{baseFS: t.fs}, + } +} + type WriteFileTool struct { fs fileSystem } @@ -188,6 +195,12 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]any) *ToolR return SilentResult(fmt.Sprintf("File written: %s", path)) } +func (t *WriteFileTool) UpgradeToConcurrent() Tool { + return &WriteFileTool{ + fs: &ConcurrentFS{baseFS: t.fs}, + } +} + type ListDirTool struct { fs fileSystem } @@ -253,6 +266,7 @@ func formatDirEntries(entries []os.DirEntry) *ToolResult { type fileSystem interface { ReadFile(path string) ([]byte, error) WriteFile(path string, data []byte) error + EditFile(path string, editFn func([]byte) ([]byte, error)) error ReadDir(path string) ([]os.DirEntry, error) } @@ -273,6 +287,23 @@ func (h *hostFs) ReadFile(path string) ([]byte, error) { return content, nil } +func (h *hostFs) EditFile(path string, editFn func([]byte) ([]byte, error)) error { + + content, err := os.ReadFile(path) + if err != nil && !os.IsNotExist(err) { + return fmt.Errorf("failed to read file for editing: %w", err) + } + // If it doesn't exist, we pass an empty byte slice to the editFn. + // This is important for "append" operations which might create new files. + + newContent, err := editFn(content) + if err != nil { + return err + } + + return fileutil.WriteFileAtomic(path, newContent, 0o600) +} + func (h *hostFs) ReadDir(path string) ([]os.DirEntry, error) { return os.ReadDir(path) } @@ -381,6 +412,63 @@ func (r *sandboxFs) WriteFile(path string, data []byte) error { }) } +func (r *sandboxFs) EditFile(path string, editFn func([]byte) ([]byte, error)) error { + return r.execute(path, func(root *os.Root, relPath string) error { + // 1. Read + content, err := root.ReadFile(relPath) + if err != nil && !os.IsNotExist(err) { + return fmt.Errorf("failed to read file for editing: %w", err) + } + + // 2. Modify + newContent, err := editFn(content) + if err != nil { + return err + } + + // 3. Write (reusing the atomic write logic) + dir := filepath.Dir(relPath) + if dir != "." && dir != "/" { + if err := root.MkdirAll(dir, 0o755); err != nil { + return fmt.Errorf("failed to create parent directories: %w", err) + } + } + + tmpRelPath := fmt.Sprintf(".tmp-%d-%d", os.Getpid(), time.Now().UnixNano()) + tmpFile, err := root.OpenFile(tmpRelPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + root.Remove(tmpRelPath) + return fmt.Errorf("failed to open temp file: %w", err) + } + + if _, err := tmpFile.Write(newContent); err != nil { + tmpFile.Close() + root.Remove(tmpRelPath) + return fmt.Errorf("failed to write temp file: %w", err) + } + if err := tmpFile.Sync(); err != nil { + tmpFile.Close() + root.Remove(tmpRelPath) + return fmt.Errorf("failed to sync temp file: %w", err) + } + if err := tmpFile.Close(); err != nil { + root.Remove(tmpRelPath) + return fmt.Errorf("failed to close temp file: %w", err) + } + + if err := root.Rename(tmpRelPath, relPath); err != nil { + root.Remove(tmpRelPath) + return fmt.Errorf("failed to rename temp file over target: %w", err) + } + if dirFile, err := root.Open("."); err == nil { + _ = dirFile.Sync() + dirFile.Close() + } + + return nil + }) +} + func (r *sandboxFs) ReadDir(path string) ([]os.DirEntry, error) { var entries []os.DirEntry err := r.execute(path, func(root *os.Root, relPath string) error { @@ -415,3 +503,50 @@ func getSafeRelPath(workspace, path string) (string, error) { return rel, nil } + +// ConcurrencyUpgradeable indicates a Tool operates on files and can be upgraded +// to use a thread-safe locking proxy backend (`ConcurrentFS`) for Parallel or DAG agent teams. +type ConcurrencyUpgradeable interface { + UpgradeToConcurrent() Tool +} + +// Global file locks explicitly for concurrent agent strategies +var globalFileLocks sync.Map // map[string]*sync.RWMutex + +func getPathLock(path string) *sync.RWMutex { + cleanPath := filepath.Clean(path) + actual, _ := globalFileLocks.LoadOrStore(cleanPath, &sync.RWMutex{}) + return actual.(*sync.RWMutex) +} + +// ConcurrentFS is a lightweight proxy wrapper around any `fileSystem`. +// It guarantees thread-safe, race-condition-free access by locking the absolute file path globally. +type ConcurrentFS struct { + baseFS fileSystem +} + +func (c *ConcurrentFS) ReadFile(path string) ([]byte, error) { + lock := getPathLock(path) + lock.RLock() + defer lock.RUnlock() + return c.baseFS.ReadFile(path) +} + +func (c *ConcurrentFS) WriteFile(path string, data []byte) error { + lock := getPathLock(path) + lock.Lock() + defer lock.Unlock() + return c.baseFS.WriteFile(path, data) +} + +func (c *ConcurrentFS) EditFile(path string, editFn func([]byte) ([]byte, error)) error { + lock := getPathLock(path) + lock.Lock() + defer lock.Unlock() + return c.baseFS.EditFile(path, editFn) +} + +func (c *ConcurrentFS) ReadDir(path string) ([]os.DirEntry, error) { + // Directories rarely suffer from single-file corruption, but we delegate anyway. + return c.baseFS.ReadDir(path) +} diff --git a/pkg/tools/filesystem_test.go b/pkg/tools/filesystem_test.go index 6f896e22d..6ccfec3be 100644 --- a/pkg/tools/filesystem_test.go +++ b/pkg/tools/filesystem_test.go @@ -486,3 +486,61 @@ func TestRootRW_Write(t *testing.T) { assert.NoError(t, err) assert.Equal(t, newData, content) } + +// TestConcurrentFS_RaceCondition simulates a high-concurrency environment +// to prove that the ConcurrentFS proxy correctly prevents file corruption. +func TestConcurrentFS_RaceCondition(t *testing.T) { + tmpDir := t.TempDir() + testFile := filepath.Join(tmpDir, "race_test.txt") + + // Pre-fill file with "0" + err := os.WriteFile(testFile, []byte("0"), 0o600) + assert.NoError(t, err) + + // Create a base FS (hostFs) and wrap it in ConcurrentFS + baseFS := &hostFs{} + concurrentFs := &ConcurrentFS{baseFS: baseFS} + + numGoroutines := 100 + done := make(chan bool) + + // Simulate 100 goroutines trying to append/edit simultaneously + for i := 0; i < numGoroutines; i++ { + go func() { + _ = concurrentFs.EditFile(testFile, func(content []byte) ([]byte, error) { + // Artificial parsing of a number to increment + newContent := append(content, []byte("-x")...) + return newContent, nil + }) + done <- true + }() + } + + // Wait for all to finish + for i := 0; i < numGoroutines; i++ { + <-done + } + + // Verify the file isn't corrupted and has exactly 100 "-x" additions + finalData, err := os.ReadFile(testFile) + assert.NoError(t, err) + + finalStr := string(finalData) + xCount := strings.Count(finalStr, "-x") + assert.Equal(t, numGoroutines, xCount, "Race condition detected! The file was corrupted or missed writes.") +} + +func TestConcurrencyUpgradeable(t *testing.T) { + // Verify that ReadFileTool implements the interface and upgrades correctly + readTool := NewReadFileTool("", false) + upgradable, ok := interface{}(readTool).(ConcurrencyUpgradeable) + assert.True(t, ok, "ReadFileTool should implement ConcurrencyUpgradeable") + + upgradedTool := upgradable.UpgradeToConcurrent() + upgradedReadTool, ok := upgradedTool.(*ReadFileTool) + assert.True(t, ok, "Upgraded tool should still be a *ReadFileTool") + + // Ensure the internal fs is now a *ConcurrentFS + _, isConcurrent := upgradedReadTool.fs.(*ConcurrentFS) + assert.True(t, isConcurrent, "Internal fileSystem should be upgraded to *ConcurrentFS") +} diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index d37a093a8..894f74834 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -35,6 +35,17 @@ func (r *ToolRegistry) Get(name string) (Tool, bool) { return tool, ok } +// ListTools returns a slice of all registered tool names. +func (r *ToolRegistry) ListTools() []string { + r.mu.RLock() + defer r.mu.RUnlock() + names := make([]string, 0, len(r.tools)) + for name := range r.tools { + names = append(names, name) + } + return names +} + func (r *ToolRegistry) Execute(ctx context.Context, name string, args map[string]any) *ToolResult { return r.ExecuteWithContext(ctx, name, args, "", "", nil) } diff --git a/pkg/tools/spawn_sub_agent.go b/pkg/tools/spawn_sub_agent.go new file mode 100644 index 000000000..cc740f16f --- /dev/null +++ b/pkg/tools/spawn_sub_agent.go @@ -0,0 +1,105 @@ +package tools + +import ( + "context" + "fmt" + "strings" + + "github.com/sipeed/picoclaw/pkg/providers" +) + +// SpawnSubAgentTool executes a customized subagent task synchronously using Anthology-style single worker delegation. +type SpawnSubAgentTool struct { + manager *SubagentManager + originChannel string + originChatID string +} + +func NewSpawnSubAgentTool(manager *SubagentManager) *SpawnSubAgentTool { + return &SpawnSubAgentTool{ + manager: manager, + originChannel: "cli", + originChatID: "direct", + } +} + +func (t *SpawnSubAgentTool) Name() string { + return "spawn_sub_agent" +} + +func (t *SpawnSubAgentTool) Description() string { + return "Directly delegate a specific task to a new, isolated sub-agent. You (the main agent) should autonomously determine the appropriate expert role and specific task based on the user's high-level request. It will execute independently and return the final result." +} + +func (t *SpawnSubAgentTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "task": map[string]any{ + "type": "string", + "description": "The specific task the sub-agent needs to accomplish.", + }, + "role": map[string]any{ + "type": "string", + "description": "The system prompt/role assignment for the sub-agent (e.g., 'You are an expert code reviewer').", + }, + }, + "required": []string{"task", "role"}, + } +} + +func (t *SpawnSubAgentTool) SetContext(channel, chatID string) { + t.originChannel = channel + t.originChatID = chatID +} + +func (t *SpawnSubAgentTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + task, ok := args["task"].(string) + if !ok || strings.TrimSpace(task) == "" { + return ErrorResult("task is required").WithError(fmt.Errorf("task parameter is required")) + } + + role, ok := args["role"].(string) + if !ok || strings.TrimSpace(role) == "" { + return ErrorResult("role is required").WithError(fmt.Errorf("role parameter is required")) + } + + if t.manager == nil { + return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil")) + } + + // 1. Isolation: Each SubAgent gets a completely fresh message set + messages := []providers.Message{ + { + Role: "system", + Content: role, + }, + { + Role: "user", + Content: task, + }, + } + + // 2. Base Configuration (Timeout & LLM constraints) + config := t.manager.BuildBaseWorkerConfig(ctx) + + // Note: For MVP, we pass the current ToolRegistry unmodified. + // To enforce strict sandboxing later, we can construct a new ToolRegistry here based on args['allowed_tools']. + + loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID) + if err != nil { + return ErrorResult(fmt.Sprintf("Subagent execution failed: %v", err)).WithError(err) + } + + // Return full details to LLM + llmContent := fmt.Sprintf("Subagent (Role: %s) task completed:\nIterations: %d\nResult: %s", + role, loopResult.Iterations, loopResult.Content) + + return &ToolResult{ + ForLLM: llmContent, + ForUser: "Sub-agent finished task.", + Silent: false, + IsError: false, + Async: false, + } +} diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 69f1a49a2..3e3551d19 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -248,6 +248,31 @@ func (sm *SubagentManager) ListTasks() []*SubagentTask { return tasks } +// BuildBaseWorkerConfig returns a base ToolLoopConfig that can be customized for isolated workers. +func (sm *SubagentManager) BuildBaseWorkerConfig(ctx context.Context) ToolLoopConfig { + sm.mu.RLock() + defer sm.mu.RUnlock() + + var llmOptions map[string]any + if sm.hasMaxTokens || sm.hasTemperature { + llmOptions = map[string]any{} + if sm.hasMaxTokens { + llmOptions["max_tokens"] = sm.maxTokens + } + if sm.hasTemperature { + llmOptions["temperature"] = sm.temperature + } + } + + return ToolLoopConfig{ + Provider: sm.provider, + Model: sm.defaultModel, + Tools: sm.tools, // Note: Caller should replace this for isolated registry + MaxIterations: sm.maxIterations, + LLMOptions: llmOptions, + } +} + // SubagentTool executes a subagent task synchronously and returns the result. // Unlike SpawnTool which runs tasks asynchronously, SubagentTool waits for completion // and returns the result directly in the ToolResult. diff --git a/pkg/tools/team.go b/pkg/tools/team.go new file mode 100644 index 000000000..8fb500b70 --- /dev/null +++ b/pkg/tools/team.go @@ -0,0 +1,571 @@ +package tools + +import ( + "context" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/sipeed/picoclaw/pkg/providers" +) + +type TeamTool struct { + manager *SubagentManager + originChannel string + originChatID string +} + +type TeamMember struct { + ID string + Role string + Task string + DependsOn []string // List of member IDs this member depends on +} + +func NewTeamTool(manager *SubagentManager) *TeamTool { + return &TeamTool{ + manager: manager, + originChannel: "cli", + originChatID: "direct", + } +} + +func (t *TeamTool) Name() string { + return "team" +} + +func (t *TeamTool) Description() string { + return "Compose and execute a team of distinct sub-agents. You (the main agent) should autonomously analyze the user's request, determine the necessary specialized roles, break down the work into sub-tasks, and assign them. Execute sequentially (passing output from one to the next) or concurrently in parallel." +} + +func (t *TeamTool) Parameters() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "strategy": map[string]any{ + "type": "string", + "enum": []string{"sequential", "parallel", "dag", "evaluator_optimizer"}, + "description": "How to run the team members. 'sequential': one after another. 'parallel': all at once. 'dag': execute based on declared dependencies. 'evaluator_optimizer': EXACTLY two members (worker & evaluator). The evaluator will check the worker's output; if it fails, the worker is revived with its FULL stateful memory intact and asked to fix it. Use this for complex generation tasks (like coding) requiring deep reasoning.", + }, + "max_team_tokens": map[string]any{ + "type": "integer", + "description": "The maximum combined LLM tokens (prompt + completion) this entire team is allowed to consume. Once exceeded, the team is instantly killed.", + }, + "members": map[string]any{ + "type": "array", + "description": "The list of sub-agents in the team.", + "items": map[string]any{ + "type": "object", + "properties": map[string]any{ + "id": map[string]any{ + "type": "string", + "description": "Unique identifier for this member, used for dependencies in 'dag' strategy.", + }, + "role": map[string]any{ + "type": "string", + "description": "The system prompt/role assignment for the member.", + }, + "task": map[string]any{ + "type": "string", + "description": "The specific task this member needs to accomplish.", + }, + "depends_on": map[string]any{ + "type": "array", + "description": "List of 'id' strings this member depends on. Only applicable for 'dag' strategy.", + "items": map[string]any{"type": "string"}, + }, + }, + "required": []string{"role", "task"}, + }, + }, + }, + "required": []string{"strategy", "members"}, + } +} + +func (t *TeamTool) SetContext(channel, chatID string) { + t.originChannel = channel + t.originChatID = chatID +} + +func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult { + strategy, ok := args["strategy"].(string) + if !ok || (strategy != "sequential" && strategy != "parallel" && strategy != "dag" && strategy != "evaluator_optimizer") { + return ErrorResult("strategy must be 'sequential', 'parallel', 'dag', or 'evaluator_optimizer'") + } + + membersRaw, ok := args["members"].([]any) + if !ok || len(membersRaw) == 0 { + return ErrorResult("members map array is required and must not be empty") + } + + maxTokensFloat, ok := args["max_team_tokens"].(float64) + var budget *atomic.Int64 + if ok && maxTokensFloat > 0 { + budget = &atomic.Int64{} + budget.Store(int64(maxTokensFloat)) + } + + if t.manager == nil { + return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil")) + } + + var members []TeamMember + for i, mRaw := range membersRaw { + mMap, ok := mRaw.(map[string]any) + if !ok { + return ErrorResult(fmt.Sprintf("member at index %d is invalid", i)) + } + + id, iOk := mMap["id"].(string) + role, rOk := mMap["role"].(string) + task, tOk := mMap["task"].(string) + + if !rOk || !tOk || strings.TrimSpace(role) == "" || strings.TrimSpace(task) == "" { + return ErrorResult(fmt.Sprintf("member at index %d is missing required 'role' or 'task'", i)) + } + + // ID is highly recommended, generate one if missing for backwards compatibility + if !iOk || strings.TrimSpace(id) == "" { + id = fmt.Sprintf("member_%d", i) + } + + var dependsOn []string + if depRaw, dOk := mMap["depends_on"].([]any); dOk { + for _, d := range depRaw { + if dStr, dsOk := d.(string); dsOk { + dependsOn = append(dependsOn, dStr) + } + } + } + + members = append(members, TeamMember{ + ID: id, + Role: role, + Task: task, + DependsOn: dependsOn, + }) + } + + // Base struct setup + baseConfig := t.manager.BuildBaseWorkerConfig(ctx) + if budget != nil { + baseConfig.RemainingTokenBudget = budget + } + + // Create a new master context for team bounding + // In the future this could be overridden by an argument + teamCtx, cancel := context.WithTimeout(ctx, 15*time.Minute) + defer cancel() + + // If strategy is parallel or dag, we must upgrade the file tools to be concurrent-safe (locking) + if strategy == "parallel" || strategy == "dag" { + baseConfig.Tools = upgradeRegistryForConcurrency(baseConfig.Tools) + } + + switch strategy { + case "sequential": + return t.executeSequential(teamCtx, baseConfig, members) + case "dag": + return t.executeDAG(teamCtx, baseConfig, members) + case "evaluator_optimizer": + return t.executeEvaluatorOptimizer(teamCtx, baseConfig, members) + } + return t.executeParallel(teamCtx, baseConfig, members) +} + +// upgradeRegistryForConcurrency takes an existing ToolRegistry, clones it, +// and upgrades any tools that implement ConcurrencyUpgradeable to their locking counterparts. +func upgradeRegistryForConcurrency(original *ToolRegistry) *ToolRegistry { + if original == nil { + return nil + } + + upgraded := NewToolRegistry() + for _, name := range original.ListTools() { + tool, ok := original.Get(name) + if !ok { + continue + } + + if upgradeable, isUpgradeable := tool.(ConcurrencyUpgradeable); isUpgradeable { + upgraded.Register(upgradeable.UpgradeToConcurrent()) + } else { + upgraded.Register(tool) + } + } + return upgraded +} + +func (t *TeamTool) executeSequential(ctx context.Context, config ToolLoopConfig, members []TeamMember) *ToolResult { + var finalOutput strings.Builder + finalOutput.WriteString("Team Execution Summary (Sequential):\n\n") + + var previousResult string + + for i, m := range members { + // If there is a previous result, we append it to the task so the new agent sees it. + actualTask := m.Task + if i > 0 && previousResult != "" { + actualTask = fmt.Sprintf("%s\n\n--- Context from previous phase ---\n%s", m.Task, truncateContext(previousResult)) + } + + messages := []providers.Message{ + {Role: "system", Content: m.Role}, + {Role: "user", Content: actualTask}, + } + + loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID) + if err != nil { + errStr := fmt.Sprintf("Phase %d (Role: %s) failed: %v", i+1, m.Role, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) // Fail fast + } + + previousResult = loopResult.Content + + finalOutput.WriteString(fmt.Sprintf("### Phase %d completed by Role: [%s]\n%s\n\n", i+1, m.Role, previousResult)) + } + + return &ToolResult{ + ForLLM: finalOutput.String(), + ForUser: "Team completed sequential execution successfully.", + } +} + +func (t *TeamTool) executeParallel(ctx context.Context, config ToolLoopConfig, members []TeamMember) *ToolResult { + var wg sync.WaitGroup + type workResult struct { + index int + role string + res string + err error + } + + resultsChan := make(chan workResult, len(members)) + + for i, m := range members { + wg.Add(1) + go func(index int, role, task string) { + defer wg.Done() + + messages := []providers.Message{ + {Role: "system", Content: role}, + {Role: "user", Content: task}, + } + + loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID) + + if err != nil { + resultsChan <- workResult{index: index, role: role, err: err} + return + } + resultsChan <- workResult{index: index, role: role, res: loopResult.Content} + }(i, m.Role, m.Task) + } + + // Wait for all goroutines to finish + wg.Wait() + close(resultsChan) + + var finalOutput strings.Builder + finalOutput.WriteString("Team Execution Summary (Parallel):\n\n") + + // Pre-allocate to maintain order since channels don't guarantee arrival order + orderedResults := make([]workResult, len(members)) + for res := range resultsChan { + orderedResults[res.index] = res + } + + hasError := false + for _, res := range orderedResults { + if res.err != nil { + hasError = true + finalOutput.WriteString(fmt.Sprintf("### Worker [%s] FAILED:\n%v\n\n", res.role, res.err)) + } else { + finalOutput.WriteString(fmt.Sprintf("### Worker [%s] Output:\n%s\n\n", res.role, res.res)) + } + } + + if hasError { + return ErrorResult("One or more parallel workers failed.\n" + finalOutput.String()) + } + + return &ToolResult{ + ForLLM: finalOutput.String(), + ForUser: "Team completed parallel execution.", + } +} + +func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, config ToolLoopConfig, members []TeamMember) *ToolResult { + if len(members) != 2 { + return ErrorResult("The evaluator_optimizer strategy requires exactly two members: [0] Worker, [1] Evaluator.") + } + + worker := members[0] + evaluator := members[1] + + var finalOutput strings.Builder + finalOutput.WriteString("Team Execution Summary (Evaluator-Optimizer):\n\n") + + // 1. Initialize the stateful memory for the worker + workerMessages := []providers.Message{ + {Role: "system", Content: worker.Role}, + {Role: "user", Content: worker.Task}, + } + + maxLoops := 5 + for attempt := 1; attempt <= maxLoops; attempt++ { + finalOutput.WriteString(fmt.Sprintf("## Attempt %d\n", attempt)) + + // 2. Trigger Worker (resumes from its exact previous state!) + workerResult, err := RunToolLoop(ctx, config, workerMessages, t.originChannel, t.originChatID) + if err != nil { + errStr := fmt.Sprintf("Worker failed on attempt %d: %v", attempt, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) + } + + // Save the worker's cognitive state so it remembers its thought process for the next loop + workerMessages = workerResult.Messages + + finalOutput.WriteString(fmt.Sprintf("### Worker Output:\n%s\n\n", workerResult.Content)) + + // 3. Trigger Evaluator (Ephemeral, stateless evaluation) + evalContext := fmt.Sprintf("%s\n\n--- Worker's Output to Evaluate ---\n%s\n\nIf the output is completely correct and fulfills the task, you MUST reply starting with strictly '[PASS]'. Otherwise, explain the issues in detail.", evaluator.Task, truncateContext(workerResult.Content)) + + evalMessages := []providers.Message{ + {Role: "system", Content: evaluator.Role}, + {Role: "user", Content: evalContext}, + } + + evalResult, err := RunToolLoop(ctx, config, evalMessages, t.originChannel, t.originChatID) + if err != nil { + errStr := fmt.Sprintf("Evaluator failed on attempt %d: %v", attempt, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) + } + + finalOutput.WriteString(fmt.Sprintf("### Evaluator Feedback:\n%s\n\n", evalResult.Content)) + + // 4. Check for PASS condition + if strings.HasPrefix(strings.TrimSpace(evalResult.Content), "[PASS]") { + finalOutput.WriteString("✅ Evaluation Passed! Loop finished successfully.\n") + return &ToolResult{ + ForLLM: finalOutput.String(), + ForUser: "Evaluator-Optimizer loop completed successfully.", + } + } + + // 5. If not passed, and not the last attempt, inject feedback into Worker's stateful memory + if attempt < maxLoops { + injection := fmt.Sprintf("The evaluator rejected your previous attempt. Please fix the issues based on this feedback:\n\n%s", evalResult.Content) + workerMessages = append(workerMessages, providers.Message{ + Role: "user", + Content: injection, + }) + } + } + + finalOutput.WriteString("❌ Maximum evaluation loops reached without a [PASS]. Returning current state.\n") + return &ToolResult{ + ForLLM: finalOutput.String(), + ForUser: "Evaluator-Optimizer loop exhausted maximum attempts.", + } +} + +func (t *TeamTool) executeDAG(ctx context.Context, config ToolLoopConfig, members []TeamMember) *ToolResult { + // 1. Build and VALIDATE dependency graph + memberMap := make(map[string]TeamMember) + inDegree := make(map[string]int) + graph := make(map[string][]string) // node -> nodes that depend on it + + // Register all valid members first + for _, m := range members { + memberMap[m.ID] = m + inDegree[m.ID] = 0 + graph[m.ID] = []string{} + } + + // Build edges and check for ghost nodes + for _, m := range members { + for _, dep := range m.DependsOn { + if _, exists := memberMap[dep]; !exists { + return ErrorResult(fmt.Sprintf("DAG Validation Error: Member [%s] depends on undefined member [%s]", m.ID, dep)) + } + graph[dep] = append(graph[dep], m.ID) + inDegree[m.ID]++ + } + } + + // 1.5. Cycle Detection using Kahn's Algorithm + var kahnQueue []string + kahnInDegree := make(map[string]int) + for k, v := range inDegree { + kahnInDegree[k] = v + if v == 0 { + kahnQueue = append(kahnQueue, k) + } + } + + processedCount := 0 + for len(kahnQueue) > 0 { + curr := kahnQueue[0] + kahnQueue = kahnQueue[1:] + processedCount++ + + for _, dependent := range graph[curr] { + kahnInDegree[dependent]-- + if kahnInDegree[dependent] == 0 { + kahnQueue = append(kahnQueue, dependent) + } + } + } + + if processedCount != len(members) { + return ErrorResult("DAG Validation Error: Circular dependency (cycle) detected in the team layout. Please fix your 'depends_on' definitions.") + } + + // 2. Channels for coordination + type nodeResult struct { + id string + res string + err error + } + readyChan := make(chan string, len(members)) + resultChan := make(chan nodeResult, len(members)) + + // Channels specifically for passing context from dependencies to dependants + contextMap := make(map[string]*strings.Builder) + var contextMu sync.Mutex + + // 3. Initialize queue with nodes having 0 in-degree + nodesToProcess := len(members) + for id, deg := range inDegree { + if deg == 0 { + readyChan <- id + } + } + + var wg sync.WaitGroup + var masterErr error + var masterErrMu sync.Mutex + + // Shared results store for the final output + finalResults := make(map[string]string) + var finalResultsMu sync.Mutex + + // 4. DAG Execution Loop + for i := 0; i < nodesToProcess; i++ { + select { + case <-ctx.Done(): + return ErrorResult("DAG execution timed out or cancelled") + + case memberID := <-readyChan: + wg.Add(1) + go func(id string) { + defer wg.Done() + + m := memberMap[id] + + // Construct the task with context from all dependencies + actualTask := m.Task + contextMu.Lock() + b := contextMap[id] + depsContext := "" + if b != nil { + depsContext = b.String() + } + contextMu.Unlock() + + if depsContext != "" { + actualTask = fmt.Sprintf("%s\n\n--- Context from dependencies ---\n%s", m.Task, truncateContext(depsContext)) + } + + messages := []providers.Message{ + {Role: "system", Content: m.Role}, + {Role: "user", Content: actualTask}, + } + + loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID) + + if err != nil { + masterErrMu.Lock() + if masterErr == nil { + masterErr = fmt.Errorf("worker [%s] failed: %v", m.ID, err) + } + masterErrMu.Unlock() + resultChan <- nodeResult{id: id, err: err} + return + } + + // Store result for final output + finalResultsMu.Lock() + finalResults[id] = loopResult.Content + finalResultsMu.Unlock() + + // Pass result to dependents + resultChan <- nodeResult{id: id, res: loopResult.Content} + }(memberID) + + case res := <-resultChan: + if res.err != nil { + // Fast fail on first error + return ErrorResult(res.err.Error()) + } + + // Update dependents + for _, dependentID := range graph[res.id] { + contextMu.Lock() + b := contextMap[dependentID] + if b == nil { + b = &strings.Builder{} + } + b.WriteString(fmt.Sprintf("--- Result from [%s] ---\n%s\n\n", res.id, res.res)) + contextMap[dependentID] = b + contextMu.Unlock() + + inDegree[dependentID]-- + if inDegree[dependentID] == 0 { + readyChan <- dependentID + } + } + } + } + + // Wait for any remaining goroutines (though the select loop handles the exact count) + wg.Wait() + + if masterErr != nil { + return ErrorResult(masterErr.Error()) + } + + // 5. Format final output + var finalOutput strings.Builder + finalOutput.WriteString("Team Execution Summary (DAG):\n\n") + + // Preserve original member order for final output readability + for _, m := range members { + if res, ok := finalResults[m.ID]; ok { + finalOutput.WriteString(fmt.Sprintf("### Worker [%s] (Role: %s) Output:\n%s\n\n", m.ID, m.Role, res)) + } + } + + return &ToolResult{ + ForLLM: finalOutput.String(), + ForUser: "Team completed DAG execution.", + } +} + +// truncateContext prevents Context Window Explosion (Token Bombs) +// by limiting the size of upstream results injected into downstream prompts. +func truncateContext(ctx string) string { + maxRunes := 8000 + runes := []rune(ctx) + if len(runes) > maxRunes { + return string(runes[:maxRunes]) + "\n...[Context truncated due to length]..." + } + return ctx +} diff --git a/pkg/tools/team_test.go b/pkg/tools/team_test.go new file mode 100644 index 000000000..9d76bdb2e --- /dev/null +++ b/pkg/tools/team_test.go @@ -0,0 +1,62 @@ +package tools + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestUpgradeRegistryForConcurrency(t *testing.T) { + // Create a standard tool registry + original := NewToolRegistry() + + // Register a mix of tools: some upgradeable, some not + readTool := NewReadFileTool("", false) + listTool := NewListDirTool("", false) // Not upgradeable + writeTool := NewWriteFileTool("", false) + + original.Register(readTool) + original.Register(listTool) + original.Register(writeTool) + + // Perform the upgrade + upgraded := upgradeRegistryForConcurrency(original) + + // Verify count matches + assert.Equal(t, len(original.ListTools()), len(upgraded.ListTools()), "Upgraded registry should have same number of tools") + + // Verify ReadFileTool got upgraded + actualReadTool, ok := upgraded.Get("read_file") + assert.True(t, ok) + if upgradedRead, isUpgraded := actualReadTool.(*ReadFileTool); isUpgraded { + _, isConcurrent := upgradedRead.fs.(*ConcurrentFS) + assert.True(t, isConcurrent, "read_file should have been upgraded to ConcurrentFS") + } + + // Verify WriteFileTool got upgraded + actualWriteTool, ok := upgraded.Get("write_file") + assert.True(t, ok) + if upgradedWrite, isUpgraded := actualWriteTool.(*WriteFileTool); isUpgraded { + _, isConcurrent := upgradedWrite.fs.(*ConcurrentFS) + assert.True(t, isConcurrent, "write_file should have been upgraded to ConcurrentFS") + } + + // Verify ListDirTool remained the same + actualListTool, ok := upgraded.Get("list_dir") + assert.True(t, ok) + _, isListDir := actualListTool.(*ListDirTool) + assert.True(t, isListDir, "list_dir should still be ListDirTool") + + // Double check list_dir doesn't randomly have ConcurrentFS injected + if listImpl, ok := actualListTool.(*ListDirTool); ok { + _, isConcurrent := listImpl.fs.(*ConcurrentFS) + assert.False(t, isConcurrent, "list_dir should NOT have ConcurrentFS because it's not upgradeable") + } + + // Double check original registry was entirely unmodified + origReadTool, _ := original.Get("read_file") + if origRead, _ := origReadTool.(*ReadFileTool); origRead != nil { + _, isConcurrent := origRead.fs.(*ConcurrentFS) + assert.False(t, isConcurrent, "Original registry components MUST REMAIN completely lock-free") + } +} diff --git a/pkg/tools/toolloop.go b/pkg/tools/toolloop.go index cdfe0d6ce..abdd0d762 100644 --- a/pkg/tools/toolloop.go +++ b/pkg/tools/toolloop.go @@ -10,6 +10,7 @@ import ( "context" "encoding/json" "fmt" + "sync/atomic" "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" @@ -18,17 +19,19 @@ import ( // ToolLoopConfig configures the tool execution loop. type ToolLoopConfig struct { - Provider providers.LLMProvider - Model string - Tools *ToolRegistry - MaxIterations int - LLMOptions map[string]any + Provider providers.LLMProvider + Model string + Tools *ToolRegistry + MaxIterations int + LLMOptions map[string]any + RemainingTokenBudget *atomic.Int64 } // ToolLoopResult contains the result of running the tool loop. type ToolLoopResult struct { Content string Iterations int + Messages []providers.Message // Allows caller to retain stateful context across executions } // RunToolLoop executes the LLM + tool call iteration loop. @@ -73,6 +76,18 @@ func RunToolLoop( return nil, fmt.Errorf("LLM call failed: %w", err) } + // 3.5 Token Budget Enforcement + if response.Usage != nil && config.RemainingTokenBudget != nil { + newBudget := config.RemainingTokenBudget.Add(-int64(response.Usage.TotalTokens)) + if newBudget < 0 { + logger.ErrorCF("toolloop", "Token budget exceeded", map[string]any{ + "used_iteration": response.Usage.TotalTokens, + "deficit": newBudget, + }) + return nil, fmt.Errorf("Token Budget Exceeded: Team budget completely consumed") + } + } + // 4. If no tool calls, we're done if len(response.ToolCalls) == 0 { finalContent = response.Content @@ -158,5 +173,6 @@ func RunToolLoop( return &ToolLoopResult{ Content: finalContent, Iterations: iteration, + Messages: messages, }, nil }