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 1/5] 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 } From 429ce66a7d0cf6bbbd7ca53f1dd7e97813f109a8 Mon Sep 17 00:00:00 2001 From: Administrator <1280842908@qq.com> Date: Sat, 28 Feb 2026 14:16:56 +0800 Subject: [PATCH 2/5] feat(teams): implement heterogeneous agent model routing - Added 'model' property to 'team' and 'spawn_sub_agent' JSON schemas. - Modified 'buildWorkerConfig' to override 'baseConfig.Model' when a specific LLM model is requested by the coordinator. - Allows teams to dynamically mix and match specialized vision, coding, and logical models within the same execution loop. --- pkg/tools/spawn_sub_agent.go | 9 ++++++ pkg/tools/team.go | 56 ++++++++++++++++++++++++++---------- 2 files changed, 50 insertions(+), 15 deletions(-) diff --git a/pkg/tools/spawn_sub_agent.go b/pkg/tools/spawn_sub_agent.go index cc740f16f..7617f8a12 100644 --- a/pkg/tools/spawn_sub_agent.go +++ b/pkg/tools/spawn_sub_agent.go @@ -43,6 +43,10 @@ func (t *SpawnSubAgentTool) Parameters() map[string]any { "type": "string", "description": "The system prompt/role assignment for the sub-agent (e.g., 'You are an expert code reviewer').", }, + "model": map[string]any{ + "type": "string", + "description": "Optional specific LLM model ID to route this task to (e.g., 'gpt-4o' for vision, 'claude-3-5-sonnet' for logic). If omitted, inherits the parent's model.", + }, }, "required": []string{"task", "role"}, } @@ -83,6 +87,11 @@ func (t *SpawnSubAgentTool) Execute(ctx context.Context, args map[string]any) *T // 2. Base Configuration (Timeout & LLM constraints) config := t.manager.BuildBaseWorkerConfig(ctx) + // 2.1 Model Override (Heterogeneous Agents) + if modelParam, ok := args["model"].(string); ok && strings.TrimSpace(modelParam) != "" { + config.Model = strings.TrimSpace(modelParam) + } + // 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']. diff --git a/pkg/tools/team.go b/pkg/tools/team.go index 8fb500b70..4290487b0 100644 --- a/pkg/tools/team.go +++ b/pkg/tools/team.go @@ -21,6 +21,7 @@ type TeamMember struct { ID string Role string Task string + Model string // Heterogeneous Agents: Optional specific model for this task DependsOn []string // List of member IDs this member depends on } @@ -71,6 +72,10 @@ func (t *TeamTool) Parameters() map[string]any { "type": "string", "description": "The specific task this member needs to accomplish.", }, + "model": map[string]any{ + "type": "string", + "description": "Optional specific LLM model ID to route this task to (e.g., 'gpt-4o' for vision, 'claude-3-5-sonnet' for logic). If omitted, inherits the parent's model.", + }, "depends_on": map[string]any{ "type": "array", "description": "List of 'id' strings this member depends on. Only applicable for 'dag' strategy.", @@ -132,6 +137,9 @@ func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult id = fmt.Sprintf("member_%d", i) } + modelStr, _ := mMap["model"].(string) + modelStr = strings.TrimSpace(modelStr) + var dependsOn []string if depRaw, dOk := mMap["depends_on"].([]any); dOk { for _, d := range depRaw { @@ -145,6 +153,7 @@ func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult ID: id, Role: role, Task: task, + Model: modelStr, DependsOn: dependsOn, }) } @@ -199,7 +208,19 @@ func upgradeRegistryForConcurrency(original *ToolRegistry) *ToolRegistry { return upgraded } -func (t *TeamTool) executeSequential(ctx context.Context, config ToolLoopConfig, members []TeamMember) *ToolResult { +// buildWorkerConfig creates a ToolLoopConfig for a specific team member, +// potentially overriding the model based on the member's definition. +func buildWorkerConfig(baseConfig ToolLoopConfig, registry *ToolRegistry, m TeamMember) ToolLoopConfig { + cfg := baseConfig + cfg.Tools = registry + // Heterogeneous Agents: Override model if this team member requested a specific one + if m.Model != "" { + cfg.Model = m.Model + } + return cfg +} + +func (t *TeamTool) executeSequential(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult { var finalOutput strings.Builder finalOutput.WriteString("Team Execution Summary (Sequential):\n\n") @@ -217,7 +238,8 @@ func (t *TeamTool) executeSequential(ctx context.Context, config ToolLoopConfig, {Role: "user", Content: actualTask}, } - loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID) + workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, m) + loopResult, err := RunToolLoop(ctx, workerConfig, 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") @@ -235,7 +257,7 @@ func (t *TeamTool) executeSequential(ctx context.Context, config ToolLoopConfig, } } -func (t *TeamTool) executeParallel(ctx context.Context, config ToolLoopConfig, members []TeamMember) *ToolResult { +func (t *TeamTool) executeParallel(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult { var wg sync.WaitGroup type workResult struct { index int @@ -248,22 +270,23 @@ func (t *TeamTool) executeParallel(ctx context.Context, config ToolLoopConfig, m for i, m := range members { wg.Add(1) - go func(index int, role, task string) { + go func(index int, member TeamMember) { defer wg.Done() messages := []providers.Message{ - {Role: "system", Content: role}, - {Role: "user", Content: task}, + {Role: "system", Content: member.Role}, + {Role: "user", Content: member.Task}, } - loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID) + workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, member) + loopResult, err := RunToolLoop(ctx, workerConfig, messages, t.originChannel, t.originChatID) if err != nil { - resultsChan <- workResult{index: index, role: role, err: err} + resultsChan <- workResult{index: index, role: member.Role, err: err} return } - resultsChan <- workResult{index: index, role: role, res: loopResult.Content} - }(i, m.Role, m.Task) + resultsChan <- workResult{index: index, role: member.Role, res: loopResult.Content} + }(i, m) } // Wait for all goroutines to finish @@ -299,7 +322,7 @@ func (t *TeamTool) executeParallel(ctx context.Context, config ToolLoopConfig, m } } -func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, config ToolLoopConfig, members []TeamMember) *ToolResult { +func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult { if len(members) != 2 { return ErrorResult("The evaluator_optimizer strategy requires exactly two members: [0] Worker, [1] Evaluator.") } @@ -321,7 +344,8 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, config ToolLoo 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) + workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, worker) + workerResult, err := RunToolLoop(ctx, workerConfig, workerMessages, t.originChannel, t.originChatID) if err != nil { errStr := fmt.Sprintf("Worker failed on attempt %d: %v", attempt, err) finalOutput.WriteString(errStr + "\n") @@ -341,7 +365,8 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, config ToolLoo {Role: "user", Content: evalContext}, } - evalResult, err := RunToolLoop(ctx, config, evalMessages, t.originChannel, t.originChatID) + evalConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, evaluator) + evalResult, err := RunToolLoop(ctx, evalConfig, evalMessages, t.originChannel, t.originChatID) if err != nil { errStr := fmt.Sprintf("Evaluator failed on attempt %d: %v", attempt, err) finalOutput.WriteString(errStr + "\n") @@ -376,7 +401,7 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, config ToolLoo } } -func (t *TeamTool) executeDAG(ctx context.Context, config ToolLoopConfig, members []TeamMember) *ToolResult { +func (t *TeamTool) executeDAG(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult { // 1. Build and VALIDATE dependency graph memberMap := make(map[string]TeamMember) inDegree := make(map[string]int) @@ -489,7 +514,8 @@ func (t *TeamTool) executeDAG(ctx context.Context, config ToolLoopConfig, member {Role: "user", Content: actualTask}, } - loopResult, err := RunToolLoop(ctx, config, messages, t.originChannel, t.originChatID) + workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, m) + loopResult, err := RunToolLoop(ctx, workerConfig, messages, t.originChannel, t.originChatID) if err != nil { masterErrMu.Lock() From efb9dc27ad657b414337491bda6623a38ec3f943 Mon Sep 17 00:00:00 2001 From: Administrator <1280842908@qq.com> Date: Sat, 28 Feb 2026 14:44:46 +0800 Subject: [PATCH 3/5] feat(teams): add model tags system for capability-aware agent routing - Added ModelTag* constants (vision, code, fast, long-context, reasoning) to subagent.go - Added Tags []string to config.ModelConfig (json:"tags,omitempty") - Piped tags from config through FallbackCandidate and into SubagentManager.allowedModels - Changed ResolveCandidatesWithLookup lookup signature to return (string, []string, bool) to carry tags - Added ModelCapabilityHint() that generates rich per-model routing guidance for the LLM - Dynamically injected capability hints into TeamTool and SpawnSubAgentTool descriptions - Fixed fallback_test.go and subagent test files to match updated signatures --- pkg/agent/instance.go | 12 +++--- pkg/agent/loop.go | 2 +- pkg/config/config.go | 5 ++- pkg/providers/fallback.go | 9 +++- pkg/providers/fallback_test.go | 18 ++++---- pkg/tools/spawn_sub_agent.go | 14 ++++++- pkg/tools/spawn_test.go | 4 +- pkg/tools/subagent.go | 73 ++++++++++++++++++++++++++++++++- pkg/tools/subagent_tool_test.go | 20 ++++----- pkg/tools/team.go | 58 ++++++++++++++++++++++---- 10 files changed, 172 insertions(+), 43 deletions(-) diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 65a1fe04d..097cb4efb 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -92,7 +92,7 @@ func NewAgentInstance( Primary: model, Fallbacks: fallbacks, } - resolveFromModelList := func(raw string) (string, bool) { + resolveFromModelList := func(raw string) (string, []string, bool) { ensureProtocol := func(model string) string { model = strings.TrimSpace(model) if model == "" { @@ -106,12 +106,12 @@ func NewAgentInstance( raw = strings.TrimSpace(raw) if raw == "" { - return "", false + return "", nil, false } if cfg != nil { if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" { - return ensureProtocol(mc.Model), true + return ensureProtocol(mc.Model), mc.Tags, true } for i := range cfg.ModelList { @@ -120,16 +120,16 @@ func NewAgentInstance( continue } if fullModel == raw { - return ensureProtocol(fullModel), true + return ensureProtocol(fullModel), cfg.ModelList[i].Tags, true } _, modelID := providers.ExtractProtocol(fullModel) if modelID == raw { - return ensureProtocol(fullModel), true + return ensureProtocol(fullModel), cfg.ModelList[i].Tags, true } } } - return "", false + return "", nil, false } candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 1b8654eb6..6e9d9d4c2 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -147,7 +147,7 @@ func registerSharedTools( agent.Tools.Register(tools.NewInstallSkillTool(registryMgr, agent.Workspace)) // Spawn tool with allowlist checker - subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace, msgBus) + subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, msgBus) subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) spawnTool := tools.NewSpawnTool(subagentManager) currentAgentID := agentID diff --git a/pkg/config/config.go b/pkg/config/config.go index d84772d2b..80cd8253c 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -457,8 +457,9 @@ type OpenAIProviderConfig struct { // Default protocol is "openai" if no prefix is specified. type ModelConfig struct { // Required fields - ModelName string `json:"model_name"` // User-facing alias for the model - Model string `json:"model"` // Protocol/model-identifier (e.g., "openai/gpt-4o", "anthropic/claude-sonnet-4.6") + ModelName string `json:"model_name"` // User-facing alias for the model + Model string `json:"model"` // Protocol/model-identifier (e.g., "openai/gpt-4o", "anthropic/claude-sonnet-4.6") + Tags []string `json:"tags,omitempty"` // Model capability labels like 'vision' // HTTP-based providers APIBase string `json:"api_base,omitempty"` // API endpoint URL diff --git a/pkg/providers/fallback.go b/pkg/providers/fallback.go index 7ba563b66..2c0f52f21 100644 --- a/pkg/providers/fallback.go +++ b/pkg/providers/fallback.go @@ -16,6 +16,7 @@ type FallbackChain struct { type FallbackCandidate struct { Provider string Model string + Tags []string } // FallbackResult contains the successful response and metadata about all attempts. @@ -49,16 +50,19 @@ func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandid func ResolveCandidatesWithLookup( cfg ModelConfig, defaultProvider string, - lookup func(raw string) (resolved string, ok bool), + lookup func(raw string) (resolved string, tags []string, ok bool), ) []FallbackCandidate { seen := make(map[string]bool) var candidates []FallbackCandidate addCandidate := func(raw string) { candidateRaw := strings.TrimSpace(raw) + var modelTags []string + if lookup != nil { - if resolved, ok := lookup(candidateRaw); ok { + if resolved, tags, ok := lookup(candidateRaw); ok { candidateRaw = resolved + modelTags = tags } } @@ -74,6 +78,7 @@ func ResolveCandidatesWithLookup( candidates = append(candidates, FallbackCandidate{ Provider: ref.Provider, Model: ref.Model, + Tags: modelTags, }) } diff --git a/pkg/providers/fallback_test.go b/pkg/providers/fallback_test.go index 1783ebcb5..f9014db2d 100644 --- a/pkg/providers/fallback_test.go +++ b/pkg/providers/fallback_test.go @@ -459,11 +459,11 @@ func TestResolveCandidatesWithLookup_AliasResolvesToNestedModel(t *testing.T) { Fallbacks: nil, } - lookup := func(raw string) (string, bool) { + lookup := func(raw string) (string, []string, bool) { if raw == "step-3.5-flash" { - return "openrouter/stepfun/step-3.5-flash:free", true + return "openrouter/stepfun/step-3.5-flash:free", nil, true } - return "", false + return "", nil, false } candidates := ResolveCandidatesWithLookup(cfg, "", lookup) @@ -484,11 +484,11 @@ func TestResolveCandidatesWithLookup_DeduplicateAfterLookup(t *testing.T) { Fallbacks: []string{"openrouter/stepfun/step-3.5-flash:free"}, } - lookup := func(raw string) (string, bool) { + lookup := func(raw string) (string, []string, bool) { if raw == "step-3.5-flash" { - return "openrouter/stepfun/step-3.5-flash:free", true + return "openrouter/stepfun/step-3.5-flash:free", nil, true } - return "", false + return "", nil, false } candidates := ResolveCandidatesWithLookup(cfg, "", lookup) @@ -503,11 +503,11 @@ func TestResolveCandidatesWithLookup_AliasWithoutProtocolUsesDefaultProvider(t * Fallbacks: nil, } - lookup := func(raw string) (string, bool) { + lookup := func(raw string) (string, []string, bool) { if raw == "glm-5" { - return "glm-5", true + return "glm-5", nil, true } - return "", false + return "", nil, false } candidates := ResolveCandidatesWithLookup(cfg, "openai", lookup) diff --git a/pkg/tools/spawn_sub_agent.go b/pkg/tools/spawn_sub_agent.go index 7617f8a12..ecc71b69f 100644 --- a/pkg/tools/spawn_sub_agent.go +++ b/pkg/tools/spawn_sub_agent.go @@ -28,7 +28,13 @@ func (t *SpawnSubAgentTool) Name() string { } 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." + base := "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." + if t.manager != nil { + if hint := t.manager.ModelCapabilityHint(); hint != "" { + return base + "\n\n" + hint + } + } + return base } func (t *SpawnSubAgentTool) Parameters() map[string]any { @@ -89,7 +95,11 @@ func (t *SpawnSubAgentTool) Execute(ctx context.Context, args map[string]any) *T // 2.1 Model Override (Heterogeneous Agents) if modelParam, ok := args["model"].(string); ok && strings.TrimSpace(modelParam) != "" { - config.Model = strings.TrimSpace(modelParam) + requestedModel := strings.TrimSpace(modelParam) + if !t.manager.IsModelAllowed(requestedModel) { + return ErrorResult(fmt.Sprintf("requested model '%s' is not in the allowed fallback candidates list for this agent workspace", requestedModel)).WithError(fmt.Errorf("model %s not allowed", requestedModel)) + } + config.Model = requestedModel } // Note: For MVP, we pass the current ToolRegistry unmodified. diff --git a/pkg/tools/spawn_test.go b/pkg/tools/spawn_test.go index 0646c82a9..b65a5e928 100644 --- a/pkg/tools/spawn_test.go +++ b/pkg/tools/spawn_test.go @@ -8,7 +8,7 @@ import ( func TestSpawnTool_Execute_EmptyTask(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSpawnTool(manager) ctx := context.Background() @@ -42,7 +42,7 @@ func TestSpawnTool_Execute_EmptyTask(t *testing.T) { func TestSpawnTool_Execute_ValidTask(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSpawnTool(manager) ctx := context.Background() diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 3e3551d19..03eeb8743 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -3,6 +3,7 @@ package tools import ( "context" "fmt" + "strings" "sync" "time" @@ -10,6 +11,26 @@ import ( "github.com/sipeed/picoclaw/pkg/providers" ) +// ModelTag constants define the recognized capability labels for models in config.json. +// These are set via `"tags": ["vision", "code"]` under each model in the model list. +const ( + ModelTagVision = "vision" // Supports image/screenshot input (multimodal) + ModelTagCode = "code" // Specialized for code generation and analysis + ModelTagFast = "fast" // Low-latency model, suited for lightweight tasks + ModelTagLongContext = "long-context" // Supports very long context windows (>100k tokens) + ModelTagReasoning = "reasoning" // Strong logical/math reasoning (e.g., o1, deepseek-r1) +) + +// modelTagDescriptions provides LLM-readable explanations of each known tag, +// injected at runtime into the tool description to guide model selection. +var modelTagDescriptions = map[string]string{ + ModelTagVision: "can analyze images and screenshots", + ModelTagCode: "specialized in code generation and debugging", + ModelTagFast: "fast and lightweight, ideal for simple or high-frequency tasks", + ModelTagLongContext: "handles very long inputs (>100k tokens)", + ModelTagReasoning: "excels at logical reasoning, math, and multi-step planning", +} + type SubagentTask struct { ID string Task string @@ -27,6 +48,7 @@ type SubagentManager struct { mu sync.RWMutex provider providers.LLMProvider defaultModel string + allowedModels []providers.FallbackCandidate bus *bus.MessageBus workspace string tools *ToolRegistry @@ -40,13 +62,16 @@ type SubagentManager struct { func NewSubagentManager( provider providers.LLMProvider, - defaultModel, workspace string, + defaultModel string, + candidates []providers.FallbackCandidate, + workspace string, bus *bus.MessageBus, ) *SubagentManager { return &SubagentManager{ tasks: make(map[string]*SubagentTask), provider: provider, defaultModel: defaultModel, + allowedModels: candidates, bus: bus, workspace: workspace, tools: NewToolRegistry(), @@ -55,6 +80,52 @@ func NewSubagentManager( } } +// IsModelAllowed checks if a specific requested model exists in the permitted candidates list. +func (sm *SubagentManager) IsModelAllowed(model string) bool { + // If the user requested the default model directly, that's automatically allowed + if model == sm.defaultModel { + return true + } + + // Otherwise, check against the resolved candidates (primary + fallbacks + explicitly configured) + for _, cand := range sm.allowedModels { + if cand.Model == model { + return true + } + } + return false +} + +// ModelCapabilityHint generates a human-readable summary of allowed models and their tags. +// This is injected into the coordinator's tool descriptions so the LLM can make better routing decisions. +func (sm *SubagentManager) ModelCapabilityHint() string { + if len(sm.allowedModels) == 0 { + return "" + } + + var modelLines []string + for _, cand := range sm.allowedModels { + if len(cand.Tags) == 0 { + modelLines = append(modelLines, fmt.Sprintf(" - %s (general purpose)", cand.Model)) + continue + } + var descs []string + for _, tag := range cand.Tags { + if desc, known := modelTagDescriptions[tag]; known { + descs = append(descs, fmt.Sprintf("%s (%s)", tag, desc)) + } else { + descs = append(descs, tag) + } + } + modelLines = append(modelLines, fmt.Sprintf(" - %s [%s]", cand.Model, strings.Join(descs, ", "))) + } + + hint := "When selecting a 'model' for sub-agents, use ONLY these configured models:\n" + hint += strings.Join(modelLines, "\n") + hint += "\nIf a task requires vision/image analysis, you MUST select a model with the 'vision' tag. If no suitable model is available, omit the 'model' field to use the default." + return hint +} + // SetLLMOptions sets max tokens and temperature for subagent LLM calls. func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) { sm.mu.Lock() diff --git a/pkg/tools/subagent_tool_test.go b/pkg/tools/subagent_tool_test.go index 59bfdffae..e1d87d10b 100644 --- a/pkg/tools/subagent_tool_test.go +++ b/pkg/tools/subagent_tool_test.go @@ -47,7 +47,7 @@ func (m *MockLLMProvider) GetContextWindow() int { func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) manager.SetLLMOptions(2048, 0.6) tool := NewSubagentTool(manager) tool.SetContext("cli", "direct") @@ -74,7 +74,7 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) { // TestSubagentTool_Name verifies tool name func TestSubagentTool_Name(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) if tool.Name() != "subagent" { @@ -85,7 +85,7 @@ func TestSubagentTool_Name(t *testing.T) { // TestSubagentTool_Description verifies tool description func TestSubagentTool_Description(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) desc := tool.Description() @@ -100,7 +100,7 @@ func TestSubagentTool_Description(t *testing.T) { // TestSubagentTool_Parameters verifies tool parameters schema func TestSubagentTool_Parameters(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) params := tool.Parameters() @@ -150,7 +150,7 @@ func TestSubagentTool_Parameters(t *testing.T) { // TestSubagentTool_SetContext verifies context setting func TestSubagentTool_SetContext(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) tool.SetContext("test-channel", "test-chat") @@ -164,7 +164,7 @@ func TestSubagentTool_SetContext(t *testing.T) { func TestSubagentTool_Execute_Success(t *testing.T) { provider := &MockLLMProvider{} msgBus := bus.NewMessageBus() - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus) tool := NewSubagentTool(manager) tool.SetContext("telegram", "chat-123") @@ -220,7 +220,7 @@ func TestSubagentTool_Execute_Success(t *testing.T) { func TestSubagentTool_Execute_NoLabel(t *testing.T) { provider := &MockLLMProvider{} msgBus := bus.NewMessageBus() - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus) tool := NewSubagentTool(manager) ctx := context.Background() @@ -243,7 +243,7 @@ func TestSubagentTool_Execute_NoLabel(t *testing.T) { // TestSubagentTool_Execute_MissingTask tests error handling for missing task func TestSubagentTool_Execute_MissingTask(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil) tool := NewSubagentTool(manager) ctx := context.Background() @@ -294,7 +294,7 @@ func TestSubagentTool_Execute_NilManager(t *testing.T) { func TestSubagentTool_Execute_ContextPassing(t *testing.T) { provider := &MockLLMProvider{} msgBus := bus.NewMessageBus() - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus) tool := NewSubagentTool(manager) // Set context @@ -323,7 +323,7 @@ func TestSubagentTool_ForUserTruncation(t *testing.T) { // Create a mock provider that returns very long content provider := &MockLLMProvider{} msgBus := bus.NewMessageBus() - manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus) + manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus) tool := NewSubagentTool(manager) ctx := context.Background() diff --git a/pkg/tools/team.go b/pkg/tools/team.go index 4290487b0..f3c190aa2 100644 --- a/pkg/tools/team.go +++ b/pkg/tools/team.go @@ -38,7 +38,13 @@ func (t *TeamTool) Name() string { } 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." + base := "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." + if t.manager != nil { + if hint := t.manager.ModelCapabilityHint(); hint != "" { + return base + "\n\n" + hint + } + } + return base } func (t *TeamTool) Parameters() map[string]any { @@ -210,14 +216,17 @@ func upgradeRegistryForConcurrency(original *ToolRegistry) *ToolRegistry { // buildWorkerConfig creates a ToolLoopConfig for a specific team member, // potentially overriding the model based on the member's definition. -func buildWorkerConfig(baseConfig ToolLoopConfig, registry *ToolRegistry, m TeamMember) ToolLoopConfig { +func buildWorkerConfig(baseConfig ToolLoopConfig, registry *ToolRegistry, m TeamMember, manager *SubagentManager) (ToolLoopConfig, error) { cfg := baseConfig cfg.Tools = registry // Heterogeneous Agents: Override model if this team member requested a specific one if m.Model != "" { + if !manager.IsModelAllowed(m.Model) { + return cfg, fmt.Errorf("requested model '%s' is not in the allowed fallback candidates list for this agent workspace", m.Model) + } cfg.Model = m.Model } - return cfg + return cfg, nil } func (t *TeamTool) executeSequential(ctx context.Context, baseConfig ToolLoopConfig, members []TeamMember) *ToolResult { @@ -238,7 +247,13 @@ func (t *TeamTool) executeSequential(ctx context.Context, baseConfig ToolLoopCon {Role: "user", Content: actualTask}, } - workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, m) + workerConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, m, t.manager) + if err != nil { + errStr := fmt.Sprintf("Phase %d (Role: %s) configuration failed: %v", i+1, m.Role, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) + } + loopResult, err := RunToolLoop(ctx, workerConfig, messages, t.originChannel, t.originChatID) if err != nil { errStr := fmt.Sprintf("Phase %d (Role: %s) failed: %v", i+1, m.Role, err) @@ -278,7 +293,12 @@ func (t *TeamTool) executeParallel(ctx context.Context, baseConfig ToolLoopConfi {Role: "user", Content: member.Task}, } - workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, member) + workerConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, member, t.manager) + if err != nil { + resultsChan <- workResult{index: index, role: member.Role, err: err} + return + } + loopResult, err := RunToolLoop(ctx, workerConfig, messages, t.originChannel, t.originChatID) if err != nil { @@ -344,7 +364,13 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, baseConfig Too finalOutput.WriteString(fmt.Sprintf("## Attempt %d\n", attempt)) // 2. Trigger Worker (resumes from its exact previous state!) - workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, worker) + workerConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, worker, t.manager) + if err != nil { + errStr := fmt.Sprintf("Worker configuration failed on attempt %d: %v", attempt, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) + } + workerResult, err := RunToolLoop(ctx, workerConfig, workerMessages, t.originChannel, t.originChatID) if err != nil { errStr := fmt.Sprintf("Worker failed on attempt %d: %v", attempt, err) @@ -365,7 +391,13 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, baseConfig Too {Role: "user", Content: evalContext}, } - evalConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, evaluator) + evalConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, evaluator, t.manager) + if err != nil { + errStr := fmt.Sprintf("Evaluator configuration failed on attempt %d: %v", attempt, err) + finalOutput.WriteString(errStr + "\n") + return ErrorResult(errStr).WithError(err) + } + evalResult, err := RunToolLoop(ctx, evalConfig, evalMessages, t.originChannel, t.originChatID) if err != nil { errStr := fmt.Sprintf("Evaluator failed on attempt %d: %v", attempt, err) @@ -514,7 +546,17 @@ func (t *TeamTool) executeDAG(ctx context.Context, baseConfig ToolLoopConfig, me {Role: "user", Content: actualTask}, } - workerConfig := buildWorkerConfig(baseConfig, baseConfig.Tools, m) + workerConfig, err := buildWorkerConfig(baseConfig, baseConfig.Tools, m, t.manager) + if err != nil { + masterErrMu.Lock() + if masterErr == nil { + masterErr = err + } + masterErrMu.Unlock() + resultChan <- nodeResult{id: id, err: err} + return + } + loopResult, err := RunToolLoop(ctx, workerConfig, messages, t.originChannel, t.originChatID) if err != nil { From 33c759f70a82cc94b08996fee8239f78ee43643a Mon Sep 17 00:00:00 2001 From: Administrator <1280842908@qq.com> Date: Sat, 28 Feb 2026 20:33:40 +0800 Subject: [PATCH 4/5] feat(teams): improve team activation and resilience MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - context.go: add Rule #5 'Team delegation' to system prompt — agents now instructed to proactively use 'team' for multi-step/multi-concern tasks instead of handling them inline - team.go: add 'WHEN TO USE THIS TOOL' activation triggers to tool description; strengthen decomposition rules with domain-agnostic project-manager framing - loop.go: call subagentManager.SetTools(agent.Tools) after full registry is built so sub-agents inherit 'team' tool for recursive hierarchical decomposition - toolloop.go: replace hard token budget failure with soft graceful degradation (50% advisory warning, 0% wrap-up signal + final summary call); add truncation recovery for max_tokens cutoff (finish_reason=truncated) - openai_compat/provider.go: detect truncated JSON tool calls and set FinishReason='truncated' instead of silently storing malformed raw args --- pkg/agent/context.go | 4 +- pkg/agent/loop.go | 6 +++ pkg/providers/openai_compat/provider.go | 13 +++++- pkg/tools/subagent.go | 4 +- pkg/tools/team.go | 24 ++++++++++- pkg/tools/toolloop.go | 54 ++++++++++++++++++++++--- 6 files changed, 94 insertions(+), 11 deletions(-) diff --git a/pkg/agent/context.go b/pkg/agent/context.go index b7c6e1108..d8d063d5e 100644 --- a/pkg/agent/context.go +++ b/pkg/agent/context.go @@ -78,7 +78,9 @@ Your workspace is at: %s 3. **Memory** - When interacting with me if something seems memorable, update %s/memory/MEMORY.md -4. **Context summaries** - Conversation summaries provided as context are approximate references only. They may be incomplete or outdated. Always defer to explicit user instructions over summary content.`, +4. **Context summaries** - Conversation summaries provided as context are approximate references only. They may be incomplete or outdated. Always defer to explicit user instructions over summary content. + +5. **Team delegation** - For any task that is non-trivial, multi-step, or involves distinct concerns (e.g. "convert React to Vue", "build a feature", "analyze and report"), you MUST use the 'team' tool to delegate and parallelize. Do NOT attempt to handle complex tasks inline by calling tools one by one yourself. Decompose first, delegate second, then report the outcome.`, workspacePath, workspacePath, workspacePath, workspacePath, workspacePath) } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 6e9d9d4c2..bc43695a8 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -161,6 +161,12 @@ func registerSharedTools( spawnSubAgentTool := tools.NewSpawnSubAgentTool(subagentManager) agent.Tools.Register(spawnSubAgentTool) + + // Direction 3: Hierarchical Decomposition. + // Share the fully-built registry (which includes team, spawn_sub_agent, etc.) back + // to the subagent manager so that all workers spawned by this agent also inherit + // the full toolset — enabling sub-agents to recursively call 'team' themselves. + subagentManager.SetTools(agent.Tools) } } diff --git a/pkg/providers/openai_compat/provider.go b/pkg/providers/openai_compat/provider.go index 5dab9b03e..f5172d915 100644 --- a/pkg/providers/openai_compat/provider.go +++ b/pkg/providers/openai_compat/provider.go @@ -235,6 +235,7 @@ func parseResponse(body []byte) (*LLMResponse, error) { choice := apiResponse.Choices[0] toolCalls := make([]ToolCall, 0, len(choice.Message.ToolCalls)) + truncated := false for _, tc := range choice.Message.ToolCalls { arguments := make(map[string]any) name := "" @@ -249,8 +250,10 @@ func parseResponse(body []byte) (*LLMResponse, error) { name = tc.Function.Name if tc.Function.Arguments != "" { if err := json.Unmarshal([]byte(tc.Function.Arguments), &arguments); err != nil { + // JSON is malformed (likely truncated due to max_tokens). Log and signal truncation. log.Printf("openai_compat: failed to decode tool call arguments for %q: %v", name, err) - arguments["raw"] = tc.Function.Arguments + truncated = true + continue // Skip this malformed tool call entirely } } } @@ -274,13 +277,19 @@ func parseResponse(body []byte) (*LLMResponse, error) { toolCalls = append(toolCalls, toolCall) } + finishReason := choice.FinishReason + // Propagate truncation: if finish_reason is "length" or we detected bad JSON, mark as truncated. + if truncated || finishReason == "length" { + finishReason = "truncated" + } + return &LLMResponse{ Content: choice.Message.Content, ReasoningContent: choice.Message.ReasoningContent, Reasoning: choice.Message.Reasoning, ReasoningDetails: choice.Message.ReasoningDetails, ToolCalls: toolCalls, - FinishReason: choice.FinishReason, + FinishReason: finishReason, Usage: apiResponse.Usage, }, nil } diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 03eeb8743..05f8654f5 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -15,6 +15,7 @@ import ( // These are set via `"tags": ["vision", "code"]` under each model in the model list. const ( ModelTagVision = "vision" // Supports image/screenshot input (multimodal) + ModelTagImageGen = "image-gen" // Supports image generation output (e.g. DALL-E, Stable Diffusion) ModelTagCode = "code" // Specialized for code generation and analysis ModelTagFast = "fast" // Low-latency model, suited for lightweight tasks ModelTagLongContext = "long-context" // Supports very long context windows (>100k tokens) @@ -24,7 +25,8 @@ const ( // modelTagDescriptions provides LLM-readable explanations of each known tag, // injected at runtime into the tool description to guide model selection. var modelTagDescriptions = map[string]string{ - ModelTagVision: "can analyze images and screenshots", + ModelTagVision: "can analyze images and screenshots (multimodal input)", + ModelTagImageGen: "can generate images from text descriptions (e.g. DALL-E, Stable Diffusion)", ModelTagCode: "specialized in code generation and debugging", ModelTagFast: "fast and lightweight, ideal for simple or high-frequency tasks", ModelTagLongContext: "handles very long inputs (>100k tokens)", diff --git a/pkg/tools/team.go b/pkg/tools/team.go index f3c190aa2..b5f7e898d 100644 --- a/pkg/tools/team.go +++ b/pkg/tools/team.go @@ -38,7 +38,29 @@ func (t *TeamTool) Name() string { } func (t *TeamTool) Description() string { - base := "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." + base := `Compose and execute a team of specialized sub-agents to accomplish a complex task. + +WHEN TO USE THIS TOOL (use proactively — do not attempt to handle these alone): +- The task involves 2 or more distinct areas of concern (e.g. research + writing, coding + testing, data gathering + analysis). +- The task would require more than 5 consecutive tool calls if done alone. +- Any part of the task can be done in parallel to save time. +- The task is large enough that a single agent would likely lose context or quality midway. +- The user asks you to "build", "create", "generate", "analyze", or "convert" something non-trivial. +When in doubt, prefer delegation over doing everything yourself. + +CRITICAL RULES FOR TASK PLANNING: +1. Think like a project manager: analyze the full task first, then design the team structure before spawning anyone. +2. Decompose the task into the smallest independently-ownable units of work. A member should own exactly ONE distinct concern — not a broad compound goal. +3. Identify dependencies between units: if one member's output is required by another, declare it via 'depends_on'. Independent units should run concurrently. +4. Each member's 'task' must be precise and self-contained. Include relevant context (e.g. reference to outputs from dependencies) directly in the task description. +5. Sub-agents are full agents with access to the same tools, including this 'team' tool. If a member's sub-task is itself complex, it may recursively form its own team. + +Strategy guide: +- sequential: each step depends on the full output of the previous step in a strict chain. +- parallel: all tasks are fully independent with no shared inputs or outputs. +- dag: most real-world tasks — some tasks depend on others, some can run concurrently. +- evaluator_optimizer: the output needs iterative critique and revision cycles.` + if t.manager != nil { if hint := t.manager.ModelCapabilityHint(); hint != "" { return base + "\n\n" + hint diff --git a/pkg/tools/toolloop.go b/pkg/tools/toolloop.go index abdd0d762..1377c4c01 100644 --- a/pkg/tools/toolloop.go +++ b/pkg/tools/toolloop.go @@ -76,18 +76,60 @@ func RunToolLoop( return nil, fmt.Errorf("LLM call failed: %w", err) } - // 3.5 Token Budget Enforcement + // 3.5 Token Budget: Soft enforcement with graceful degradation. + // Budget exhaustion is NOT a hard error — workers get a chance to wrap up gracefully. 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, + originalBudget := newBudget + int64(response.Usage.TotalTokens) + + if newBudget <= 0 { + // Budget exhausted: signal the worker to wrap up and return partial result. + logger.WarnCF("toolloop", "Token budget exhausted, injecting wrap-up signal", + map[string]any{ + "deficit": -newBudget, + "iteration": iteration, + }) + finalContent = response.Content + messages = append(messages, providers.Message{ + Role: "assistant", + Content: response.Content, + }) + messages = append(messages, providers.Message{ + Role: "user", + Content: "[SYSTEM] Token budget has been exhausted. Stop all tool calls immediately and return the best result you have completed so far. Do not call any more tools.", + }) + // One final LLM call to get a summary/wrap-up from the model + if finalResp, err := config.Provider.Chat(ctx, messages, nil, config.Model, config.LLMOptions); err == nil { + finalContent = finalResp.Content + } + break + } else if originalBudget > 0 && newBudget < originalBudget/2 { + // Budget below 50%: soft warning injected into next iteration's context. + logger.WarnCF("toolloop", "Token budget below 50%, injecting advisory", + map[string]any{"remaining": newBudget, "iteration": iteration}) + messages = append(messages, providers.Message{ + Role: "user", + Content: "[SYSTEM] Advisory: token budget is running low. Please prioritize completing the most critical parts of your task and avoid unnecessary tool calls.", }) - return nil, fmt.Errorf("Token Budget Exceeded: Team budget completely consumed") } } + // 3.6 Truncation Recovery: LLM response was cut off (max_tokens hit or malformed JSON). + // Inject a recovery message so the LLM knows to retry with a shorter, complete response. + if response.FinishReason == "truncated" { + logger.WarnCF("toolloop", "LLM response was truncated (max_tokens hit), injecting recovery message", + map[string]any{"iteration": iteration}) + messages = append(messages, providers.Message{ + Role: "assistant", + Content: response.Content, + }) + messages = append(messages, providers.Message{ + Role: "user", + Content: "[SYSTEM] Your previous response was cut off because it exceeded the token limit. Please retry by producing a shorter, complete response. If you were about to call a tool, make sure the full JSON arguments are included without truncation.", + }) + continue + } + // 4. If no tool calls, we're done if len(response.ToolCalls) == 0 { finalContent = response.Content From fdc6d3503fdc38fce7e92d9527fdb4da38f4e964 Mon Sep 17 00:00:00 2001 From: Administrator <1280842908@qq.com> Date: Sat, 28 Feb 2026 21:27:04 +0800 Subject: [PATCH 5/5] feat(teams): add produces field for automatic QA reviewer injection - Added Produces string field to TeamMember struct - Added 'produces' property to team tool JSON schema - Added reviewerTaskTemplates map for code/data/document artifact types - Added maybeRunAutoReviewer() helper that auto-injects a QA reviewer agent after all workers complete when any member declares a produces type - Wired reviewer into sequential, parallel, and dag execution strategies - evaluator_optimizer skipped (already has built-in critique loop) --- pkg/tools/team.go | 85 +++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 82 insertions(+), 3 deletions(-) diff --git a/pkg/tools/team.go b/pkg/tools/team.go index b5f7e898d..264305436 100644 --- a/pkg/tools/team.go +++ b/pkg/tools/team.go @@ -23,6 +23,7 @@ type TeamMember struct { Task string Model string // Heterogeneous Agents: Optional specific model for this task DependsOn []string // List of member IDs this member depends on + Produces string // Auto-reviewer: declares artifact type ("code", "data", "document") } func NewTeamTool(manager *SubagentManager) *TeamTool { @@ -109,6 +110,10 @@ func (t *TeamTool) Parameters() map[string]any { "description": "List of 'id' strings this member depends on. Only applicable for 'dag' strategy.", "items": map[string]any{"type": "string"}, }, + "produces": map[string]any{ + "type": "string", + "description": "Declares the type of artifact this member produces. Use 'code' for source code files, 'data' for structured data/JSON/CSV, 'document' for prose documents/reports. When set, the framework automatically appends a QA reviewer step after all workers finish to validate output correctness. Omit if no verification is needed.", + }, }, "required": []string{"role", "task"}, }, @@ -123,6 +128,60 @@ func (t *TeamTool) SetContext(channel, chatID string) { t.originChatID = chatID } +// reviewerTaskTemplates maps a `produces` artifact type to the task prompt +// that the auto-injected QA reviewer will receive. +var reviewerTaskTemplates = map[string]string{ + "code": "You are a code quality reviewer. Read all code files in the workspace that were just written by your predecessors. Check for: syntax errors, incorrect or missing imports, broken logic, type mismatches, and any issues that would cause compilation or runtime failures. List every issue found with the filename and line number if possible. If everything looks correct, respond with 'REVIEW PASSED'.", + "data": "You are a data validation reviewer. Read all output data files (JSON, CSV, YAML, etc.) in the workspace. Check for: invalid format, missing required fields, schema inconsistencies, and malformed values. List every issue found. If everything is valid, respond with 'REVIEW PASSED'.", + "document": "You are a document quality reviewer. Read all output documents in the workspace. Check for: logical inconsistencies, incomplete sections, factual contradictions, and poor structure. List every issue found. If the documents are complete and correct, respond with 'REVIEW PASSED'.", +} + +// maybeRunAutoReviewer inspects TeamMembers for `produces` declarations. +// If any member produced a verifiable artifact type, it runs an automatic +// QA reviewer agent after all workers have completed. +func (t *TeamTool) maybeRunAutoReviewer( + ctx context.Context, + members []TeamMember, + baseConfig ToolLoopConfig, + workerSummary string, +) string { + // Collect unique produces types from all members + producedTypes := make(map[string]bool) + for _, m := range members { + if m.Produces != "" { + producedTypes[m.Produces] = true + } + } + if len(producedTypes) == 0 { + return "" // No verifiable artifacts declared, skip review + } + + // Build reviewer task: combine templates for all declared artifact types + var taskParts []string + for artifactType := range producedTypes { + if tmpl, ok := reviewerTaskTemplates[artifactType]; ok { + taskParts = append(taskParts, tmpl) + } + } + if len(taskParts) == 0 { + return "" // Unknown produces types, skip + } + + reviewerTask := strings.Join(taskParts, "\n\n") + + "\n\nContext from the workers that produced these artifacts:\n" + workerSummary + + reviewerMessages := []providers.Message{ + {Role: "user", Content: reviewerTask}, + } + + reviewerConfig := baseConfig + loopResult, err := RunToolLoop(ctx, reviewerConfig, reviewerMessages, t.originChannel, t.originChatID) + if err != nil { + return fmt.Sprintf("[Auto-Reviewer] Failed to run: %v", err) + } + return "[Auto-Reviewer Result]\n" + loopResult.Content +} + 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") { @@ -177,12 +236,16 @@ func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult } } + producesStr, _ := mMap["produces"].(string) + producesStr = strings.TrimSpace(producesStr) + members = append(members, TeamMember{ ID: id, Role: role, Task: task, Model: modelStr, DependsOn: dependsOn, + Produces: producesStr, }) } @@ -204,13 +267,29 @@ func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult switch strategy { case "sequential": - return t.executeSequential(teamCtx, baseConfig, members) + result := t.executeSequential(teamCtx, baseConfig, members) + if reviewNote := t.maybeRunAutoReviewer(teamCtx, members, baseConfig, result.ForLLM); reviewNote != "" { + result.ForLLM += "\n\n" + reviewNote + result.ForUser += "\n\n" + reviewNote + } + return result case "dag": - return t.executeDAG(teamCtx, baseConfig, members) + result := t.executeDAG(teamCtx, baseConfig, members) + if reviewNote := t.maybeRunAutoReviewer(teamCtx, members, baseConfig, result.ForLLM); reviewNote != "" { + result.ForLLM += "\n\n" + reviewNote + result.ForUser += "\n\n" + reviewNote + } + return result case "evaluator_optimizer": return t.executeEvaluatorOptimizer(teamCtx, baseConfig, members) } - return t.executeParallel(teamCtx, baseConfig, members) + // parallel + result := t.executeParallel(teamCtx, baseConfig, members) + if reviewNote := t.maybeRunAutoReviewer(teamCtx, members, baseConfig, result.ForLLM); reviewNote != "" { + result.ForLLM += "\n\n" + reviewNote + result.ForUser += "\n\n" + reviewNote + } + return result } // upgradeRegistryForConcurrency takes an existing ToolRegistry, clones it,