feat(team): centralize team tool config with allowed_models, limits, and toggles

This commit is contained in:
Administrator 2026-03-12 10:04:38 +08:00
parent 55dd32267a
commit f5f991686d
9 changed files with 173 additions and 64 deletions

View file

@ -155,7 +155,7 @@ func NewAgentInstance(
Primary: model,
Fallbacks: fallbacks,
}
resolveFromModelList := func(raw string) (string, []string, bool) {
resolveFromModelList := func(raw string) (string, bool) {
ensureProtocol := func(model string) string {
model = strings.TrimSpace(model)
if model == "" {
@ -169,12 +169,12 @@ func NewAgentInstance(
raw = strings.TrimSpace(raw)
if raw == "" {
return "", nil, false
return "", false
}
if cfg != nil {
if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" {
return ensureProtocol(mc.Model), mc.Tags, true
return ensureProtocol(mc.Model), true
}
for i := range cfg.ModelList {
@ -183,16 +183,16 @@ func NewAgentInstance(
continue
}
if fullModel == raw {
return ensureProtocol(fullModel), cfg.ModelList[i].Tags, true
return ensureProtocol(fullModel), true
}
_, modelID := providers.ExtractProtocol(fullModel)
if modelID == raw {
return ensureProtocol(fullModel), cfg.ModelList[i].Tags, true
return ensureProtocol(fullModel), true
}
}
}
return "", nil, false
return "", false
}
candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList)

View file

@ -220,7 +220,7 @@ func registerSharedTools(
}
// Spawn tool with allowlist checker
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, msgBus)
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, cfg.Tools.Team, msgBus)
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
spawnTool := tools.NewSpawnTool(subagentManager)
currentAgentID := agentID
@ -230,10 +230,14 @@ func registerSharedTools(
agent.Tools.Register(spawnTool)
teamTool := tools.NewTeamTool(subagentManager)
agent.Tools.Register(teamTool)
if cfg.Tools.IsToolEnabled("team") {
agent.Tools.Register(teamTool)
}
spawnSubAgentTool := tools.NewSpawnSubAgentTool(subagentManager)
agent.Tools.Register(spawnSubAgentTool)
if cfg.Tools.IsToolEnabled("spawn_sub_agent") {
agent.Tools.Register(spawnSubAgentTool)
}
// Direction 3: Hierarchical Decomposition.
// Share the fully-built registry (which includes team, spawn_sub_agent, etc.) back
@ -242,7 +246,7 @@ func registerSharedTools(
subagentManager.SetTools(agent.Tools)
if cfg.Tools.IsToolEnabled("spawn") {
if cfg.Tools.IsToolEnabled("subagent") {
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, msgBus)
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Candidates, agent.Workspace, cfg.Tools.Team, msgBus)
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
spawnTool := tools.NewSpawnTool(subagentManager)
currentAgentID := agentID

View file

@ -74,6 +74,22 @@ func (f *FlexibleStringSlice) UnmarshalText(text []byte) error {
return nil
}
type TeamModelConfig struct {
Name string `json:"name" yaml:"name"`
Tags []string `json:"tags,omitempty" yaml:"tags,omitempty"`
}
type TeamToolsConfig struct {
ToolConfig ` envPrefix:"PICOCLAW_TOOLS_TEAM_"`
MaxMembers int `json:"max_members" env:"PICOCLAW_TOOLS_TEAM_MAX_MEMBERS"`
MaxTeamTokens int `json:"max_team_tokens" env:"PICOCLAW_TOOLS_TEAM_MAX_TOKENS"`
MaxEvaluatorLoops int `json:"max_evaluator_loops" env:"PICOCLAW_TOOLS_TEAM_MAX_EVALUATOR_LOOPS"`
MaxTimeoutMinutes int `json:"max_timeout_minutes" env:"PICOCLAW_TOOLS_TEAM_MAX_TIMEOUT_MINUTES"`
DisableAutoReviewer bool `json:"disable_auto_reviewer" env:"PICOCLAW_TOOLS_TEAM_DISABLE_AUTO_REVIEWER"`
AllowedStrategies []string `json:"allowed_strategies" env:"PICOCLAW_TOOLS_TEAM_ALLOWED_STRATEGIES"`
AllowedModels []TeamModelConfig `json:"allowed_models" env:"-"`
}
type Config struct {
Agents AgentsConfig `json:"agents"`
Bindings []AgentBinding `json:"bindings,omitempty"`
@ -591,7 +607,6 @@ 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")
Tags []string `json:"tags,omitempty"` // Model capability labels like 'vision'
// HTTP-based providers
APIBase string `json:"api_base,omitempty"` // API endpoint URL
@ -748,6 +763,8 @@ type ToolsConfig struct {
Spawn ToolConfig `json:"spawn" envPrefix:"PICOCLAW_TOOLS_SPAWN_"`
SPI ToolConfig `json:"spi" envPrefix:"PICOCLAW_TOOLS_SPI_"`
Subagent ToolConfig `json:"subagent" envPrefix:"PICOCLAW_TOOLS_SUBAGENT_"`
SpawnSubAgent ToolConfig `json:"spawn_sub_agent" envPrefix:"PICOCLAW_TOOLS_SPAWN_SUB_AGENT_"`
Team TeamToolsConfig `json:"team" envPrefix:"PICOCLAW_TOOLS_TEAM_"`
WebFetch ToolConfig `json:"web_fetch" envPrefix:"PICOCLAW_TOOLS_WEB_FETCH_"`
WriteFile ToolConfig `json:"write_file" envPrefix:"PICOCLAW_TOOLS_WRITE_FILE_"`
}

View file

@ -16,7 +16,6 @@ type FallbackChain struct {
type FallbackCandidate struct {
Provider string
Model string
Tags []string
}
// FallbackResult contains the successful response and metadata about all attempts.
@ -50,19 +49,17 @@ func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandid
func ResolveCandidatesWithLookup(
cfg ModelConfig,
defaultProvider string,
lookup func(raw string) (resolved string, tags []string, ok bool),
lookup func(raw string) (resolved 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, tags, ok := lookup(candidateRaw); ok {
if resolved, ok := lookup(candidateRaw); ok {
candidateRaw = resolved
modelTags = tags
}
}
@ -78,7 +75,6 @@ func ResolveCandidatesWithLookup(
candidates = append(candidates, FallbackCandidate{
Provider: ref.Provider,
Model: ref.Model,
Tags: modelTags,
})
}

View file

@ -459,11 +459,11 @@ func TestResolveCandidatesWithLookup_AliasResolvesToNestedModel(t *testing.T) {
Fallbacks: nil,
}
lookup := func(raw string) (string, []string, bool) {
lookup := func(raw string) (string, bool) {
if raw == "step-3.5-flash" {
return "openrouter/stepfun/step-3.5-flash:free", nil, true
return "openrouter/stepfun/step-3.5-flash:free", true
}
return "", nil, false
return "", 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, []string, bool) {
lookup := func(raw string) (string, bool) {
if raw == "step-3.5-flash" {
return "openrouter/stepfun/step-3.5-flash:free", nil, true
return "openrouter/stepfun/step-3.5-flash:free", true
}
return "", nil, false
return "", false
}
candidates := ResolveCandidatesWithLookup(cfg, "", lookup)
@ -503,11 +503,11 @@ func TestResolveCandidatesWithLookup_AliasWithoutProtocolUsesDefaultProvider(t *
Fallbacks: nil,
}
lookup := func(raw string) (string, []string, bool) {
lookup := func(raw string) (string, bool) {
if raw == "glm-5" {
return "glm-5", nil, true
return "glm-5", true
}
return "", nil, false
return "", false
}
candidates := ResolveCandidatesWithLookup(cfg, "openai", lookup)

View file

@ -4,11 +4,13 @@ import (
"context"
"strings"
"testing"
"github.com/sipeed/picoclaw/pkg/config"
)
func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
tool := NewSpawnTool(manager)
ctx := context.Background()
@ -42,7 +44,7 @@ func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
func TestSpawnTool_Execute_ValidTask(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
tool := NewSpawnTool(manager)
ctx := context.Background()

View file

@ -8,6 +8,7 @@ import (
"time"
"github.com/sipeed/picoclaw/pkg/bus"
"github.com/sipeed/picoclaw/pkg/config"
"github.com/sipeed/picoclaw/pkg/providers"
)
@ -54,6 +55,7 @@ type SubagentManager struct {
bus *bus.MessageBus
workspace string
tools *ToolRegistry
teamConfig config.TeamToolsConfig
maxIterations int
maxTokens int
temperature float64
@ -67,6 +69,7 @@ func NewSubagentManager(
defaultModel string,
candidates []providers.FallbackCandidate,
workspace string,
teamConfig config.TeamToolsConfig,
bus *bus.MessageBus,
) *SubagentManager {
return &SubagentManager{
@ -74,6 +77,7 @@ func NewSubagentManager(
provider: provider,
defaultModel: defaultModel,
allowedModels: candidates,
teamConfig: teamConfig,
bus: bus,
workspace: workspace,
tools: NewToolRegistry(),
@ -89,7 +93,21 @@ func (sm *SubagentManager) IsModelAllowed(model string) bool {
return true
}
// Otherwise, check against the resolved candidates (primary + fallbacks + explicitly configured)
// 1. Check against explicitly allowed models in team config
for _, cand := range sm.teamConfig.AllowedModels {
if cand.Name == model {
return true
}
}
// 2. Otherwise, check against the resolved candidates (primary + fallbacks + explicitly configured)
// If teamConfig.AllowedModels is set, we strictly enforce it and DO NOT fall back to candidates
// unless the candidate model has tags that overlap with AllowedTags. But since AllowedTags
// was not implemented yet, just check fallback for backwards compatibility if teamConfig is empty.
if len(sm.teamConfig.AllowedModels) > 0 {
return false
}
for _, cand := range sm.allowedModels {
if cand.Model == model {
return true
@ -107,23 +125,22 @@ func (sm *SubagentManager) ModelCapabilityHint() string {
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, ", ")))
modelLines = append(modelLines, fmt.Sprintf(" - %s (general purpose)", cand.Model))
}
hint := "When selecting a 'model' for sub-agents, use ONLY these configured models:\n"
hint += strings.Join(modelLines, "\n")
if len(sm.teamConfig.AllowedModels) > 0 {
for _, cand := range sm.teamConfig.AllowedModels {
tagsStr := ""
if len(cand.Tags) > 0 {
tagsStr = fmt.Sprintf(" [%s]", strings.Join(cand.Tags, ", "))
}
hint += fmt.Sprintf(" - %s%s\n", cand.Name, tagsStr)
}
} else {
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
}

View file

@ -6,6 +6,7 @@ import (
"testing"
"github.com/sipeed/picoclaw/pkg/bus"
"github.com/sipeed/picoclaw/pkg/config"
"github.com/sipeed/picoclaw/pkg/providers"
)
@ -47,7 +48,7 @@ func (m *MockLLMProvider) GetContextWindow() int {
func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", nil)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
manager.SetLLMOptions(2048, 0.6)
tool := NewSubagentTool(manager)
@ -73,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", nil, "/tmp/test", nil)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
tool := NewSubagentTool(manager)
if tool.Name() != "subagent" {
@ -84,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", nil, "/tmp/test", nil)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
tool := NewSubagentTool(manager)
desc := tool.Description()
@ -99,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", nil, "/tmp/test", nil)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
tool := NewSubagentTool(manager)
params := tool.Parameters()
@ -149,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", nil, "/tmp/test", nil)
// manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
// tool := NewSubagentTool(manager)
// tool.SetContext("test-channel", "test-chat")
@ -163,7 +164,7 @@ func TestSubagentTool_Parameters(t *testing.T) {
func TestSubagentTool_Execute_Success(t *testing.T) {
provider := &MockLLMProvider{}
msgBus := bus.NewMessageBus()
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", msgBus)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, msgBus)
tool := NewSubagentTool(manager)
ctx := WithToolContext(context.Background(), "telegram", "chat-123")
@ -218,7 +219,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", nil, "/tmp/test", msgBus)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, msgBus)
tool := NewSubagentTool(manager)
ctx := context.Background()
@ -241,7 +242,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", nil, "/tmp/test", nil)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, nil)
tool := NewSubagentTool(manager)
ctx := context.Background()
@ -292,7 +293,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", nil, "/tmp/test", msgBus)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, msgBus)
tool := NewSubagentTool(manager)
channel := "test-channel"
@ -318,7 +319,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", nil, "/tmp/test", msgBus)
manager := NewSubagentManager(provider, "test-model", nil, "/tmp/test", config.TeamToolsConfig{}, msgBus)
tool := NewSubagentTool(manager)
ctx := context.Background()

View file

@ -167,6 +167,16 @@ func (t *TeamTool) maybeRunAutoReviewer(
return "" // Unknown produces types, skip
}
// 3. Skip Auto-Reviewer if disabled by config
sm := t.manager
sm.mu.RLock()
disabled := sm.teamConfig.DisableAutoReviewer
sm.mu.RUnlock()
if disabled {
return ""
}
reviewerTask := strings.Join(taskParts, "\n\n") +
"\n\nContext from the workers that produced these artifacts:\n" + workerSummary
@ -184,8 +194,37 @@ func (t *TeamTool) maybeRunAutoReviewer(
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'")
if !ok {
return ErrorResult("strategy is required")
}
if t.manager == nil {
return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil"))
}
sm := t.manager
sm.mu.RLock()
teamConfig := sm.teamConfig
sm.mu.RUnlock()
// 1. Validate Strategy
validStrategy := false
if len(teamConfig.AllowedStrategies) > 0 {
for _, s := range teamConfig.AllowedStrategies {
if strategy == s {
validStrategy = true
break
}
}
} else {
// Default allowed strategies if not configured
if strategy == "sequential" || strategy == "parallel" || strategy == "dag" || strategy == "evaluator_optimizer" {
validStrategy = true
}
}
if !validStrategy {
return ErrorResult(fmt.Sprintf("strategy '%s' is not allowed by configuration", strategy))
}
membersRaw, ok := args["members"].([]any)
@ -193,17 +232,38 @@ func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult
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))
// 2. Validate Max Members
if teamConfig.MaxMembers > 0 && len(membersRaw) > teamConfig.MaxMembers {
return ErrorResult(fmt.Sprintf("Team exceeds maximum allowed members (%d). You requested %d members.", teamConfig.MaxMembers, len(membersRaw)))
}
if t.manager == nil {
return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil"))
maxTokensFloat, ok := args["max_team_tokens"].(float64)
// Enforce hard budgets from config
effectiveMaxTokens := int64(0)
if teamConfig.MaxTeamTokens > 0 {
effectiveMaxTokens = int64(teamConfig.MaxTeamTokens)
}
if ok && maxTokensFloat > 0 {
requestedTokens := int64(maxTokensFloat)
// If LLM requested tokens but config enforces a smaller hard limit, clamp it
if effectiveMaxTokens > 0 && requestedTokens > effectiveMaxTokens {
effectiveMaxTokens = requestedTokens // LLM asked for more, but we clamp to config
// Wait, the clamping logic: if requested > max_team_tokens, clamp effectively shrinks it to max_team_tokens.
effectiveMaxTokens = int64(teamConfig.MaxTeamTokens)
} else if effectiveMaxTokens == 0 || requestedTokens < effectiveMaxTokens {
effectiveMaxTokens = requestedTokens // LLM asked for less budget, let them be conservative
}
}
var budget *atomic.Int64
if effectiveMaxTokens > 0 {
budget = &atomic.Int64{}
budget.Store(effectiveMaxTokens)
}
var members []TeamMember
for i, mRaw := range membersRaw {
mMap, ok := mRaw.(map[string]any)
@ -256,8 +316,11 @@ func (t *TeamTool) Execute(ctx context.Context, args map[string]any) *ToolResult
}
// 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)
timeoutDur := 15 * time.Minute
if teamConfig.MaxTimeoutMinutes > 0 {
timeoutDur = time.Duration(teamConfig.MaxTimeoutMinutes) * time.Minute
}
teamCtx, cancel := context.WithTimeout(ctx, timeoutDur)
defer cancel()
// If strategy is parallel or dag, we must upgrade the file tools to be concurrent-safe (locking)
@ -460,7 +523,16 @@ func (t *TeamTool) executeEvaluatorOptimizer(ctx context.Context, baseConfig Too
{Role: "user", Content: worker.Task},
}
sm := t.manager
sm.mu.RLock()
teamConfig := sm.teamConfig
sm.mu.RUnlock()
maxLoops := 5
if teamConfig.MaxEvaluatorLoops > 0 {
maxLoops = teamConfig.MaxEvaluatorLoops
}
for attempt := 1; attempt <= maxLoops; attempt++ {
finalOutput.WriteString(fmt.Sprintf("## Attempt %d\n", attempt))