Merge fix/subagent-provider-dispatch: per-agent subagent dispatch, allowlist self-spawn, response attribution

This commit is contained in:
Eric Jacksch 2026-03-19 23:16:00 -04:00
commit fe32c4cf2b
6 changed files with 205 additions and 87 deletions

View file

@ -85,13 +85,13 @@ func NewAgentLoop(
) *AgentLoop {
registry := NewAgentRegistry(cfg, provider)
// Register shared tools to all agents
registerSharedTools(cfg, msgBus, registry, provider)
// Set up shared fallback chain
cooldown := providers.NewCooldownTracker()
fallbackChain := providers.NewFallbackChain(cooldown)
// Register shared tools to all agents
registerSharedTools(cfg, msgBus, registry, provider, dispatcher, fallbackChain)
// Create state manager using default agent's workspace for channel recording
defaultAgent := registry.GetDefaultAgent()
var stateManager *state.Manager
@ -119,6 +119,8 @@ func registerSharedTools(
msgBus *bus.MessageBus,
registry *AgentRegistry,
provider providers.LLMProvider,
dispatcher *providers.ProviderDispatcher,
fallbackChain *providers.FallbackChain,
) {
for _, agentID := range registry.ListAgentIDs() {
agent, ok := registry.GetAgent(agentID)
@ -228,10 +230,31 @@ func registerSharedTools(
// Spawn tool with allowlist checker
if cfg.Tools.IsToolEnabled("spawn") {
if cfg.Tools.IsToolEnabled("subagent") {
subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace)
currentAgentID := agentID
// Build a resolver so the subagent manager can look up any target
// agent's candidates without importing the agent package from tools.
candidateResolver := func(targetAgentID string) ([]providers.FallbackCandidate, bool) {
target, ok := registry.GetAgent(targetAgentID)
if !ok {
return nil, false
}
if len(target.Candidates) == 0 {
return nil, false
}
return target.Candidates, true
}
subagentManager := tools.NewSubagentManager(tools.SubagentManagerConfig{
Provider: provider,
DefaultModel: agent.Model,
Workspace: agent.Workspace,
Dispatcher: dispatcher,
Fallback: fallbackChain,
SelfCandidates: agent.Candidates,
CallerAgentID: currentAgentID,
CandidateResolver: candidateResolver,
})
subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature)
spawnTool := tools.NewSpawnTool(subagentManager)
currentAgentID := agentID
spawnTool.SetAllowlistChecker(func(targetAgentID string) bool {
return registry.CanSpawnSubagent(currentAgentID, targetAgentID)
})
@ -407,8 +430,10 @@ func (al *AgentLoop) ReloadProviderAndConfig(
return fmt.Errorf("context canceled after registry creation: %w", err)
}
// Ensure shared tools are re-registered on the new registry
registerSharedTools(cfg, al.bus, registry, provider)
// Ensure shared tools are re-registered on the new registry.
// Build a fresh fallback chain for the new registry's subagent managers.
newFallbackChain := providers.NewFallbackChain(providers.NewCooldownTracker())
registerSharedTools(cfg, al.bus, registry, provider, al.dispatcher, newFallbackChain)
// Atomically swap the config and registry under write lock
// This ensures readers see a consistent pair

View file

@ -72,10 +72,16 @@ func (t *SpawnTool) execute(ctx context.Context, args map[string]any, cb AsyncCa
label, _ := args["label"].(string)
agentID, _ := args["agent_id"].(string)
// Check allowlist if targeting a specific agent
if agentID != "" && t.allowlistCheck != nil {
if !t.allowlistCheck(agentID) {
return ErrorResult(fmt.Sprintf("not allowed to spawn agent '%s'", agentID))
// Check allowlist for both self-spawn (agentID == "") and targeted spawns.
// For self-spawn we pass the manager's callerAgentID so that operators can
// restrict agents to only spawning specific peers (not themselves).
if t.allowlistCheck != nil {
checkID := agentID
if checkID == "" && t.manager != nil {
checkID = t.manager.callerAgentID
}
if checkID != "" && !t.allowlistCheck(checkID) {
return ErrorResult(fmt.Sprintf("not allowed to spawn agent '%s'", checkID))
}
}

View file

@ -8,7 +8,11 @@ import (
func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager := NewSubagentManager(SubagentManagerConfig{
Provider: provider,
DefaultModel: "test-model",
Workspace: "/tmp/test",
})
tool := NewSpawnTool(manager)
ctx := context.Background()
@ -42,7 +46,11 @@ func TestSpawnTool_Execute_EmptyTask(t *testing.T) {
func TestSpawnTool_Execute_ValidTask(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager := NewSubagentManager(SubagentManagerConfig{
Provider: provider,
DefaultModel: "test-model",
Workspace: "/tmp/test",
})
tool := NewSpawnTool(manager)
ctx := context.Background()

View file

@ -3,12 +3,22 @@ package tools
import (
"context"
"fmt"
"strings"
"sync"
"time"
"github.com/sipeed/picoclaw/pkg/providers"
)
// attributedContent prepends "AgentName: " to content when agentID is non-empty.
func attributedContent(agentID, content string) string {
if strings.TrimSpace(agentID) == "" {
return content
}
name := strings.ToUpper(agentID[:1]) + agentID[1:]
return "**" + name + ":** " + content
}
type SubagentTask struct {
ID string
Task string
@ -34,23 +44,97 @@ type SubagentManager struct {
hasMaxTokens bool
hasTemperature bool
nextID int
// Per-agent dispatch fields. When dispatcher is set, subagent LLM calls
// are routed through the dispatcher using selfCandidates (for self-spawns)
// or the target agent's candidates (resolved via candidateResolver).
dispatcher *providers.ProviderDispatcher
fallback *providers.FallbackChain
selfCandidates []providers.FallbackCandidate
callerAgentID string
candidateResolver func(agentID string) ([]providers.FallbackCandidate, bool)
}
func NewSubagentManager(
provider providers.LLMProvider,
defaultModel, workspace string,
) *SubagentManager {
// SubagentManagerConfig holds all configuration for constructing a SubagentManager.
type SubagentManagerConfig struct {
// Provider is the fallback LLM provider used when dispatcher lookup fails.
Provider providers.LLMProvider
// DefaultModel is the model name used when no candidates are resolved.
DefaultModel string
// Workspace is the agent's working directory.
Workspace string
// Dispatcher dispatches LLM calls per-candidate. Optional.
Dispatcher *providers.ProviderDispatcher
// Fallback is the chain used when multiple candidates are configured. Optional.
Fallback *providers.FallbackChain
// SelfCandidates are the calling agent's model candidates for self-spawns.
SelfCandidates []providers.FallbackCandidate
// CallerAgentID is the ID of the agent that owns this manager (for allowlist).
CallerAgentID string
// CandidateResolver resolves model candidates for a named target agent.
// Returns false if the agent is unknown.
CandidateResolver func(agentID string) ([]providers.FallbackCandidate, bool)
}
func NewSubagentManager(cfg SubagentManagerConfig) *SubagentManager {
return &SubagentManager{
tasks: make(map[string]*SubagentTask),
provider: provider,
defaultModel: defaultModel,
workspace: workspace,
tools: NewToolRegistry(),
maxIterations: 10,
nextID: 1,
tasks: make(map[string]*SubagentTask),
provider: cfg.Provider,
defaultModel: cfg.DefaultModel,
workspace: cfg.Workspace,
tools: NewToolRegistry(),
maxIterations: 10,
nextID: 1,
dispatcher: cfg.Dispatcher,
fallback: cfg.Fallback,
selfCandidates: cfg.SelfCandidates,
callerAgentID: cfg.CallerAgentID,
candidateResolver: cfg.CandidateResolver,
}
}
// resolveLoopConfig builds a ToolLoopConfig for the given target agent ID.
// When agentID is empty it is treated as a self-spawn and uses selfCandidates.
// When agentID is non-empty and a candidateResolver is set, it resolves the
// target agent's candidates. Falls back to provider+defaultModel when no
// dispatch metadata is available.
func (sm *SubagentManager) resolveLoopConfig(agentID string) ToolLoopConfig {
candidates := sm.selfCandidates
if agentID != "" && sm.candidateResolver != nil {
if resolved, ok := sm.candidateResolver(agentID); ok {
candidates = resolved
}
}
cfg := ToolLoopConfig{
Provider: sm.provider,
Model: sm.defaultModel,
MaxIterations: sm.maxIterations,
Tools: sm.tools,
Dispatcher: sm.dispatcher,
Fallback: sm.fallback,
Candidates: candidates,
}
// Override Model from first candidate when available.
if len(candidates) > 0 {
cfg.Model = candidates[0].Model
}
if sm.hasMaxTokens || sm.hasTemperature {
opts := map[string]any{}
if sm.hasMaxTokens {
opts["max_tokens"] = sm.maxTokens
}
if sm.hasTemperature {
opts["temperature"] = sm.temperature
}
cfg.LLMOptions = opts
}
return cfg
}
// SetLLMOptions sets max tokens and temperature for subagent LLM calls.
func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) {
sm.mu.Lock()
@ -139,34 +223,11 @@ After completing the task, provide a clear summary of what was done.`
default:
}
// Run tool loop with access to tools
sm.mu.RLock()
tools := sm.tools
maxIter := sm.maxIterations
maxTokens := sm.maxTokens
temperature := sm.temperature
hasMaxTokens := sm.hasMaxTokens
hasTemperature := sm.hasTemperature
loopCfg := sm.resolveLoopConfig(task.AgentID)
sm.mu.RUnlock()
var llmOptions map[string]any
if hasMaxTokens || hasTemperature {
llmOptions = map[string]any{}
if hasMaxTokens {
llmOptions["max_tokens"] = maxTokens
}
if hasTemperature {
llmOptions["temperature"] = temperature
}
}
loopResult, err := RunToolLoop(ctx, ToolLoopConfig{
Provider: sm.provider,
Model: sm.defaultModel,
Tools: tools,
MaxIterations: maxIter,
LLMOptions: llmOptions,
}, messages, task.OriginChannel, task.OriginChatID)
loopResult, err := RunToolLoop(ctx, loopCfg, messages, task.OriginChannel, task.OriginChatID)
sm.mu.Lock()
var result *ToolResult
@ -204,7 +265,7 @@ After completing the task, provide a clear summary of what was done.`
loopResult.Iterations,
loopResult.Content,
),
ForUser: loopResult.Content,
ForUser: attributedContent(task.AgentID, loopResult.Content),
Silent: false,
IsError: false,
Async: false,
@ -275,6 +336,7 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe
}
label, _ := args["label"].(string)
agentID, _ := args["agent_id"].(string)
if t.manager == nil {
return ErrorResult("Subagent manager not configured").WithError(fmt.Errorf("manager is nil"))
@ -292,28 +354,13 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe
},
}
// Use RunToolLoop to execute with tools (same as async SpawnTool)
// Use RunToolLoop to execute with tools (same as async SpawnTool).
// SubagentTool always performs a self-spawn (no explicit agent_id).
sm := t.manager
sm.mu.RLock()
tools := sm.tools
maxIter := sm.maxIterations
maxTokens := sm.maxTokens
temperature := sm.temperature
hasMaxTokens := sm.hasMaxTokens
hasTemperature := sm.hasTemperature
loopCfg := sm.resolveLoopConfig("")
sm.mu.RUnlock()
var llmOptions map[string]any
if hasMaxTokens || hasTemperature {
llmOptions = map[string]any{}
if hasMaxTokens {
llmOptions["max_tokens"] = maxTokens
}
if hasTemperature {
llmOptions["temperature"] = temperature
}
}
// Fall back to "cli"/"direct" for non-conversation callers (e.g., CLI, tests)
// to preserve the same defaults as the original NewSubagentTool constructor.
channel := ToolChannel(ctx)
@ -325,13 +372,7 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe
chatID = "direct"
}
loopResult, err := RunToolLoop(ctx, ToolLoopConfig{
Provider: sm.provider,
Model: sm.defaultModel,
Tools: tools,
MaxIterations: maxIter,
LLMOptions: llmOptions,
}, messages, channel, chatID)
loopResult, err := RunToolLoop(ctx, loopCfg, messages, channel, chatID)
if err != nil {
return ErrorResult(fmt.Sprintf("Subagent execution failed: %v", err)).WithError(err)
}
@ -353,7 +394,7 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe
return &ToolResult{
ForLLM: llmContent,
ForUser: userContent,
ForUser: attributedContent(agentID, userContent),
Silent: false,
IsError: false,
Async: false,

View file

@ -46,7 +46,7 @@ func (m *MockLLMProvider) GetContextWindow() int {
func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
manager.SetLLMOptions(2048, 0.6)
tool := NewSubagentTool(manager)
@ -72,7 +72,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")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
tool := NewSubagentTool(manager)
if tool.Name() != "subagent" {
@ -83,7 +83,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")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
tool := NewSubagentTool(manager)
desc := tool.Description()
@ -98,7 +98,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")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
tool := NewSubagentTool(manager)
params := tool.Parameters()
@ -148,7 +148,7 @@ func TestSubagentTool_Parameters(t *testing.T) {
// TestSubagentTool_Execute_Success tests successful execution
func TestSubagentTool_Execute_Success(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
tool := NewSubagentTool(manager)
ctx := WithToolContext(context.Background(), "telegram", "chat-123")
@ -202,7 +202,7 @@ func TestSubagentTool_Execute_Success(t *testing.T) {
// TestSubagentTool_Execute_NoLabel tests execution without label
func TestSubagentTool_Execute_NoLabel(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
tool := NewSubagentTool(manager)
ctx := context.Background()
@ -225,7 +225,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")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
tool := NewSubagentTool(manager)
ctx := context.Background()
@ -275,7 +275,7 @@ func TestSubagentTool_Execute_NilManager(t *testing.T) {
// TestSubagentTool_Execute_ContextPassing verifies context is properly used
func TestSubagentTool_Execute_ContextPassing(t *testing.T) {
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
tool := NewSubagentTool(manager)
channel := "test-channel"
@ -300,7 +300,7 @@ func TestSubagentTool_Execute_ContextPassing(t *testing.T) {
func TestSubagentTool_ForUserTruncation(t *testing.T) {
// Create a mock provider that returns very long content
provider := &MockLLMProvider{}
manager := NewSubagentManager(provider, "test-model", "/tmp/test")
manager := NewSubagentManager(SubagentManagerConfig{Provider: provider, DefaultModel: "test-model", Workspace: "/tmp/test"})
tool := NewSubagentTool(manager)
ctx := context.Background()

View file

@ -24,6 +24,14 @@ type ToolLoopConfig struct {
Tools *ToolRegistry
MaxIterations int
LLMOptions map[string]any
// Per-candidate dispatch: when Dispatcher and Candidates are set, RunToolLoop
// selects the provider for each LLM call through the dispatcher using the
// candidate list, falling back through candidates on error via Fallback.
// If Fallback is nil, only the first candidate is tried.
Dispatcher *providers.ProviderDispatcher
Candidates []providers.FallbackCandidate
Fallback *providers.FallbackChain
}
// ToolLoopResult contains the result of running the tool loop.
@ -63,8 +71,38 @@ func RunToolLoop(
if llmOpts == nil {
llmOpts = map[string]any{}
}
// 3. Call LLM
response, err := config.Provider.Chat(ctx, messages, providerToolDefs, config.Model, llmOpts)
// 3. Call LLM — use per-candidate dispatch when configured, else fall back
// to the plain Provider/Model pair.
var response *providers.LLMResponse
var err error
if config.Dispatcher != nil && len(config.Candidates) > 0 {
if config.Fallback != nil && len(config.Candidates) > 1 {
fbResult, fbErr := config.Fallback.Execute(
ctx,
config.Candidates,
func(ctx context.Context, providerName, model string) (*providers.LLMResponse, error) {
if p, perr := config.Dispatcher.Get(providerName, model); perr == nil {
return p.Chat(ctx, messages, providerToolDefs, model, llmOpts)
}
return config.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
},
)
if fbErr != nil {
err = fbErr
} else {
response = fbResult.Response
}
} else {
first := config.Candidates[0]
if p, perr := config.Dispatcher.Get(first.Provider, first.Model); perr == nil {
response, err = p.Chat(ctx, messages, providerToolDefs, first.Model, llmOpts)
} else {
response, err = config.Provider.Chat(ctx, messages, providerToolDefs, first.Model, llmOpts)
}
}
} else {
response, err = config.Provider.Chat(ctx, messages, providerToolDefs, config.Model, llmOpts)
}
if err != nil {
logger.ErrorCF("toolloop", "LLM call failed",
map[string]any{