diff --git a/cmd/picoclaw/internal/agent/helpers.go b/cmd/picoclaw/internal/agent/helpers.go index f754abc65..94bc8be98 100644 --- a/cmd/picoclaw/internal/agent/helpers.go +++ b/cmd/picoclaw/internal/agent/helpers.go @@ -37,19 +37,16 @@ func agentCmd(message, sessionKey, model string, debug bool) error { cfg.Agents.Defaults.ModelName = model } - provider, modelID, err := providers.CreateProvider(cfg) + // Build model registry — same startup path as gatewayCmd + modelRegistry, err := providers.NewModelRegistry(cfg) if err != nil { - return fmt.Errorf("error creating provider: %w", err) - } - - // Use the resolved model ID from provider creation - if modelID != "" { - cfg.Agents.Defaults.ModelName = modelID + return fmt.Errorf("error creating model registry: %w", err) } + defer modelRegistry.Close() msgBus := bus.NewMessageBus() defer msgBus.Close() - agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) + agentLoop := agent.NewAgentLoop(cfg, msgBus, modelRegistry) // Print agent startup info (only for interactive mode) startupInfo := agentLoop.GetStartupInfo() diff --git a/cmd/picoclaw/internal/gateway/helpers.go b/cmd/picoclaw/internal/gateway/helpers.go index 747f7d44e..c2972ae91 100644 --- a/cmd/picoclaw/internal/gateway/helpers.go +++ b/cmd/picoclaw/internal/gateway/helpers.go @@ -49,18 +49,15 @@ func gatewayCmd(debug bool) error { return fmt.Errorf("error loading config: %w", err) } - provider, modelID, err := providers.CreateProvider(cfg) + // Build model registry — the single source of truth for all providers. + // Providers are cached internally; models with identical configs share one. + modelRegistry, err := providers.NewModelRegistry(cfg) if err != nil { - return fmt.Errorf("error creating provider: %w", err) - } - - // Use the resolved model ID from provider creation - if modelID != "" { - cfg.Agents.Defaults.ModelName = modelID + return fmt.Errorf("error creating model registry: %w", err) } msgBus := bus.NewMessageBus() - agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) + agentLoop := agent.NewAgentLoop(cfg, msgBus, modelRegistry) // Print agent startup info fmt.Println("\n📦 Agent Status:") @@ -188,9 +185,7 @@ func gatewayCmd(debug bool) error { <-sigChan fmt.Println("\nShutting down...") - if cp, ok := provider.(providers.StatefulProvider); ok { - cp.Close() - } + modelRegistry.Close() cancel() msgBus.Close() diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index ed25f537f..c8348e432 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -7,8 +7,10 @@ import ( "path/filepath" "regexp" "strings" + "sync" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/routing" "github.com/sipeed/picoclaw/pkg/session" @@ -18,9 +20,13 @@ import ( // AgentInstance represents a fully configured agent with its own workspace, // session manager, context builder, and tool registry. type AgentInstance struct { + // mu protects Model, Candidates and SubagentMgr which can be changed at + // runtime by the /switch model command while the LLM loop is running. + mu sync.RWMutex + ID string Name string - Model string + Model string // registry key (model name) Fallbacks []string Workspace string MaxIterations int @@ -29,32 +35,38 @@ type AgentInstance struct { ContextWindow int SummarizeMessageThreshold int SummarizeTokenPercent int - Provider providers.LLMProvider Sessions *session.SessionManager ContextBuilder *ContextBuilder Tools *tools.ToolRegistry Subagents *config.SubagentsConfig SkillsFilter []string Candidates []providers.FallbackCandidate + SubagentMgr *tools.SubagentManager // updated on /switch model + AllowReadPaths []*regexp.Regexp // compiled path whitelist for reads + AllowWritePaths []*regexp.Regexp // compiled path whitelist for writes } // NewAgentInstance creates an agent instance from config. +// The ModelRegistry is used to resolve the provider and model ID. func NewAgentInstance( agentCfg *config.AgentConfig, defaults *config.AgentDefaults, cfg *config.Config, - provider providers.LLMProvider, + modelRegistry *providers.ModelRegistry, ) *AgentInstance { workspace := resolveAgentWorkspace(agentCfg, defaults) os.MkdirAll(workspace, 0o755) - model := resolveAgentModel(agentCfg, defaults) + modelName := resolveAgentModel(agentCfg, defaults) fallbacks := resolveAgentFallbacks(agentCfg, defaults) + // Model stays as the model name (registry key); provider is resolved at call time. + model := modelName + restrict := defaults.RestrictToWorkspace readRestrict := restrict && !defaults.AllowReadOutsideWorkspace - // Compile path whitelist patterns from config. + // Compile path whitelist patterns once and store on the instance. allowReadPaths := compilePatterns(cfg.Tools.AllowReadPaths) allowWritePaths := compilePatterns(cfg.Tools.AllowWritePaths) @@ -113,53 +125,28 @@ func NewAgentInstance( summarizeTokenPercent = 75 } - // Resolve fallback candidates - modelCfg := providers.ModelConfig{ - Primary: model, - Fallbacks: fallbacks, + // Build fallback candidates from the registry using model names as keys. + allModels := append([]string{modelName}, fallbacks...) + seen := make(map[string]bool) + var candidates []providers.FallbackCandidate + for _, name := range allModels { + name = strings.TrimSpace(name) + if name == "" || seen[name] { + continue + } + seen[name] = true + if _, ok := modelRegistry.Get(name); ok { + candidates = append(candidates, providers.FallbackCandidate{Model: name}) + } else { + logger.WarnCF("agent", "Model not found in registry, skipping", + map[string]any{"model": name, "agent_id": agentID}) + } } - resolveFromModelList := func(raw string) (string, bool) { - ensureProtocol := func(model string) string { - model = strings.TrimSpace(model) - if model == "" { - return "" - } - if strings.Contains(model, "/") { - return model - } - return "openai/" + model - } - - raw = strings.TrimSpace(raw) - if raw == "" { - return "", false - } - - if cfg != nil { - if mc, err := cfg.GetModelConfig(raw); err == nil && mc != nil && strings.TrimSpace(mc.Model) != "" { - return ensureProtocol(mc.Model), true - } - - for i := range cfg.ModelList { - fullModel := strings.TrimSpace(cfg.ModelList[i].Model) - if fullModel == "" { - continue - } - if fullModel == raw { - return ensureProtocol(fullModel), true - } - _, modelID := providers.ExtractProtocol(fullModel) - if modelID == raw { - return ensureProtocol(fullModel), true - } - } - } - - return "", false + if len(candidates) == 0 { + logger.WarnCF("agent", "Agent has no valid candidates; will use registry default", + map[string]any{"agent_id": agentID, "configured_model": modelName}) } - candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList) - return &AgentInstance{ ID: agentID, Name: agentName, @@ -172,13 +159,14 @@ func NewAgentInstance( ContextWindow: maxTokens, SummarizeMessageThreshold: summarizeMessageThreshold, SummarizeTokenPercent: summarizeTokenPercent, - Provider: provider, Sessions: sessionsManager, ContextBuilder: contextBuilder, Tools: toolsRegistry, Subagents: subagents, SkillsFilter: skillsFilter, Candidates: candidates, + AllowReadPaths: allowReadPaths, + AllowWritePaths: allowWritePaths, } } @@ -237,3 +225,27 @@ func expandHome(path string) string { } return path } + +// getModelSnapshot returns a consistent snapshot of the mutable routing fields. +// Use this at the start of each LLM iteration rather than reading fields directly. +func (a *AgentInstance) getModelSnapshot() (model string, candidates []providers.FallbackCandidate) { + a.mu.RLock() + defer a.mu.RUnlock() + model = a.Model + candidates = make([]providers.FallbackCandidate, len(a.Candidates)) + copy(candidates, a.Candidates) + return +} + +// switchModel atomically updates Model. +// Returns the previous model name. +func (a *AgentInstance) switchModel(model string) string { + a.mu.Lock() + defer a.mu.Unlock() + old := a.Model + a.Model = model + if a.SubagentMgr != nil { + a.SubagentMgr.UpdateModel(model) + } + return old +} diff --git a/pkg/agent/instance_test.go b/pkg/agent/instance_test.go index 4f41ecd1c..dc6e5bf80 100644 --- a/pkg/agent/instance_test.go +++ b/pkg/agent/instance_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" ) func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) { @@ -29,7 +30,7 @@ func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) { cfg.Agents.Defaults.Temperature = &configuredTemp provider := &mockProvider{} - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, mockRegistry(provider)) if agent.MaxTokens != 1234 { t.Fatalf("MaxTokens = %d, want %d", agent.MaxTokens, 1234) @@ -61,7 +62,7 @@ func TestNewAgentInstance_DefaultsTemperatureWhenZero(t *testing.T) { cfg.Agents.Defaults.Temperature = &configuredTemp provider := &mockProvider{} - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, mockRegistry(provider)) if agent.Temperature != 0.0 { t.Fatalf("Temperature = %f, want %f", agent.Temperature, 0.0) @@ -87,7 +88,7 @@ func TestNewAgentInstance_DefaultsTemperatureWhenUnset(t *testing.T) { } provider := &mockProvider{} - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, mockRegistry(provider)) if agent.Temperature != 0.7 { t.Fatalf("Temperature = %f, want %f", agent.Temperature, 0.7) @@ -95,68 +96,75 @@ func TestNewAgentInstance_DefaultsTemperatureWhenUnset(t *testing.T) { } func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) { - tests := []struct { - name string - aliasName string - modelName string - apiBase string - wantProvider string - wantModel string - }{ - { - name: "alias with provider prefix", - aliasName: "step-3.5-flash", - modelName: "openrouter/stepfun/step-3.5-flash:free", - apiBase: "https://openrouter.ai/api/v1", - wantProvider: "openrouter", - wantModel: "stepfun/step-3.5-flash:free", + tmpDir, err := os.MkdirTemp("", "agent-instance-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "step-3.5-flash", + }, }, - { - name: "alias without provider prefix", - aliasName: "glm-5", - modelName: "glm-5", - apiBase: "https://api.z.ai/api/coding/paas/v4", - wantProvider: "openai", - wantModel: "glm-5", + ModelList: []config.ModelConfig{ + { + ModelName: "step-3.5-flash", + Model: "openrouter/stepfun/step-3.5-flash:free", + APIBase: "https://openrouter.ai/api/v1", + }, }, } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - tmpDir, err := os.MkdirTemp("", "agent-instance-test-*") - if err != nil { - t.Fatalf("Failed to create temp dir: %v", err) - } - defer os.RemoveAll(tmpDir) + reg, err := providers.NewModelRegistry(cfg) + if err != nil { + t.Fatalf("NewModelRegistry: %v", err) + } + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, reg) - cfg := &config.Config{ - Agents: config.AgentsConfig{ - Defaults: config.AgentDefaults{ - Workspace: tmpDir, - Model: tt.aliasName, - }, - }, - ModelList: []config.ModelConfig{ - { - ModelName: tt.aliasName, - Model: tt.modelName, - APIBase: tt.apiBase, - }, - }, - } - - provider := &mockProvider{} - agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) - - if len(agent.Candidates) != 1 { - t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) - } - if agent.Candidates[0].Provider != tt.wantProvider { - t.Fatalf("candidate provider = %q, want %q", agent.Candidates[0].Provider, tt.wantProvider) - } - if agent.Candidates[0].Model != tt.wantModel { - t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, tt.wantModel) - } - }) + if len(agent.Candidates) != 1 { + t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) + } + if agent.Candidates[0].Model != "step-3.5-flash" { + t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "step-3.5-flash") + } +} + +func TestNewAgentInstance_ResolveCandidatesFromModelListAliasWithoutProtocol(t *testing.T) { + tmpDir, err := os.MkdirTemp("", "agent-instance-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "glm-5", + }, + }, + ModelList: []config.ModelConfig{ + { + ModelName: "glm-5", + Model: "glm-5", + APIBase: "https://api.z.ai/api/coding/paas/v4", + }, + }, + } + + reg, err := providers.NewModelRegistry(cfg) + if err != nil { + t.Fatalf("NewModelRegistry: %v", err) + } + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, reg) + + if len(agent.Candidates) != 1 { + t.Fatalf("len(Candidates) = %d, want 1", len(agent.Candidates)) + } + if agent.Candidates[0].Model != "glm-5" { + t.Fatalf("candidate model = %q, want %q", agent.Candidates[0].Model, "glm-5") } } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index da43bf177..0319276d0 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -37,6 +37,7 @@ type AgentLoop struct { bus *bus.MessageBus cfg *config.Config registry *AgentRegistry + modelRegistry *providers.ModelRegistry state *state.Manager running atomic.Bool summarizing sync.Map @@ -60,34 +61,38 @@ type processOptions struct { const defaultResponse = "I've completed processing but have no response to give. Increase `max_tool_iterations` in config.json." -func NewAgentLoop( - cfg *config.Config, - msgBus *bus.MessageBus, - provider providers.LLMProvider, -) *AgentLoop { - registry := NewAgentRegistry(cfg, provider) +// NewAgentLoop creates an agent loop. The ModelRegistry is the single source +// of truth for all model→provider mappings. +func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, modelRegistry *providers.ModelRegistry) *AgentLoop { + agentRegistry := NewAgentRegistry(cfg, modelRegistry) // Register shared tools to all agents - registerSharedTools(cfg, msgBus, registry, provider) + registerSharedTools(cfg, msgBus, agentRegistry, modelRegistry) // Set up shared fallback chain cooldown := providers.NewCooldownTracker() - fallbackChain := providers.NewFallbackChain(cooldown) + fallbackChain := providers.NewFallbackChain(cooldown).WithProviderKeyFn(func(model string) string { + if entry, ok := modelRegistry.Get(model); ok { + return entry.ProviderKey + } + return model + }) // Create state manager using default agent's workspace for channel recording - defaultAgent := registry.GetDefaultAgent() + defaultAgent := agentRegistry.GetDefaultAgent() var stateManager *state.Manager if defaultAgent != nil { stateManager = state.NewManager(defaultAgent.Workspace) } return &AgentLoop{ - bus: msgBus, - cfg: cfg, - registry: registry, - state: stateManager, - summarizing: sync.Map{}, - fallback: fallbackChain, + bus: msgBus, + cfg: cfg, + registry: agentRegistry, + modelRegistry: modelRegistry, + state: stateManager, + summarizing: sync.Map{}, + fallback: fallbackChain, } } @@ -96,7 +101,7 @@ func registerSharedTools( cfg *config.Config, msgBus *bus.MessageBus, registry *AgentRegistry, - provider providers.LLMProvider, + modelRegistry *providers.ModelRegistry, ) { for _, agentID := range registry.ListAgentIDs() { agent, ok := registry.GetAgent(agentID) @@ -161,9 +166,10 @@ func registerSharedTools( agent.Tools.Register(tools.NewFindSkillsTool(registryMgr, searchCache)) agent.Tools.Register(tools.NewInstallSkillTool(registryMgr, agent.Workspace)) - // Spawn tool with allowlist checker - subagentManager := tools.NewSubagentManager(provider, agent.Model, agent.Workspace, msgBus) + // Spawn tool with allowlist checker — each agent uses its own provider + subagentManager := tools.NewSubagentManager(modelRegistry, agent.Model, agent.Workspace, msgBus) subagentManager.SetLLMOptions(agent.MaxTokens, agent.Temperature) + agent.SubagentMgr = subagentManager spawnTool := tools.NewSpawnTool(subagentManager) currentAgentID := agentID spawnTool.SetAllowlistChecker(func(targetAgentID string) bool { @@ -317,6 +323,23 @@ func (al *AgentLoop) Stop() { al.running.Store(false) } +// agentEntry resolves the registry entry for an agent's current model. +// All LLM calls should go through this to ensure provider is always derived from the registry. +func (al *AgentLoop) agentEntry(agent *AgentInstance) (*providers.ModelEntry, error) { + agent.mu.RLock() + model := agent.Model + agent.mu.RUnlock() + if al.modelRegistry != nil { + if entry, ok := al.modelRegistry.Get(model); ok { + return entry, nil + } + if entry, ok := al.modelRegistry.GetDefault(); ok { + return entry, nil + } + } + return nil, fmt.Errorf("no provider found for model %q", model) +} + func (al *AgentLoop) RegisterTool(tool tools.Tool) { for _, agentID := range al.registry.ListAgentIDs() { if agent, ok := al.registry.GetAgent(agentID); ok { @@ -450,6 +473,18 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) return al.processSystemMessage(ctx, msg) } + // Reset message-tool sentInRound before any early-return paths (including + // handleCommand) so that a stale true from the previous LLM round never + // causes the command response to be silently dropped in the Run loop's + // alreadySent check (which would leave the typing indicator stuck on). + if defaultAgent := al.registry.GetDefaultAgent(); defaultAgent != nil { + if tool, ok := defaultAgent.Tools.Get("message"); ok { + if mt, ok := tool.(tools.ContextualTool); ok { + mt.SetContext(msg.Channel, msg.ChatID) + } + } + } + // Check for commands if response, handled := al.handleCommand(ctx, msg); handled { return response, nil @@ -731,6 +766,10 @@ func (al *AgentLoop) runLLMIteration( for iteration < agent.MaxIterations { iteration++ + // Snapshot mutable routing fields once per iteration so that a concurrent + // /switch model command cannot cause a mid-iteration inconsistency. + currentModel, currentCandidates := agent.getModelSnapshot() + logger.DebugCF("agent", "LLM iteration", map[string]any{ "agent_id": agent.ID, @@ -746,7 +785,7 @@ func (al *AgentLoop) runLLMIteration( map[string]any{ "agent_id": agent.ID, "iteration": iteration, - "model": agent.Model, + "model": currentModel, "messages_count": len(messages), "tools_count": len(providerToolDefs), "max_tokens": agent.MaxTokens, @@ -767,38 +806,39 @@ func (al *AgentLoop) runLLMIteration( var err error callLLM := func() (*providers.LLMResponse, error) { - if len(agent.Candidates) > 1 && al.fallback != nil { - fbResult, fbErr := al.fallback.Execute( - ctx, - agent.Candidates, - func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { - return agent.Provider.Chat( - ctx, - messages, - providerToolDefs, - model, - map[string]any{ - "max_tokens": agent.MaxTokens, - "temperature": agent.Temperature, - "prompt_cache_key": agent.ID, - }, - ) + if len(currentCandidates) > 1 && al.fallback != nil { + fbResult, fbErr := al.fallback.Execute(ctx, currentCandidates, + func(ctx context.Context, model string) (*providers.LLMResponse, error) { + entry, ok := al.modelRegistry.Get(model) + if !ok { + return nil, fmt.Errorf("model %q not found in registry", model) + } + return entry.Provider.Chat(ctx, messages, providerToolDefs, entry.ModelID, map[string]any{ + "max_tokens": agent.MaxTokens, + "temperature": agent.Temperature, + "prompt_cache_key": agent.ID, + }) }, ) if fbErr != nil { return nil, fbErr } - if fbResult.Provider != "" && len(fbResult.Attempts) > 0 { - logger.InfoCF( - "agent", - fmt.Sprintf("Fallback: succeeded with %s/%s after %d attempts", - fbResult.Provider, fbResult.Model, len(fbResult.Attempts)+1), - map[string]any{"agent_id": agent.ID, "iteration": iteration}, - ) + if fbResult.Model != "" && len(fbResult.Attempts) > 0 { + logger.InfoCF("agent", fmt.Sprintf("Fallback: succeeded with %s/%s after %d attempts", + fbResult.Provider, fbResult.Model, len(fbResult.Attempts)+1), + map[string]any{"agent_id": agent.ID, "iteration": iteration}) } return fbResult.Response, nil } - return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{ + entry, ok := al.modelRegistry.Get(currentModel) + if !ok { + var defOk bool + entry, defOk = al.modelRegistry.GetDefault() + if !defOk { + return nil, fmt.Errorf("no provider found for model %q", currentModel) + } + } + return entry.Provider.Chat(ctx, messages, providerToolDefs, entry.ModelID, map[string]any{ "max_tokens": agent.MaxTokens, "temperature": agent.Temperature, "prompt_cache_key": agent.ID, @@ -1286,17 +1326,23 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { s1, s2, ) - resp, err := agent.Provider.Chat( - ctx, - []providers.Message{{Role: "user", Content: mergePrompt}}, - nil, - agent.Model, - map[string]any{ - "max_tokens": 1024, - "temperature": 0.3, - "prompt_cache_key": agent.ID, - }, - ) + resp, err := func() (*providers.LLMResponse, error) { + entry, err := al.agentEntry(agent) + if err != nil { + return nil, err + } + return entry.Provider.Chat( + ctx, + []providers.Message{{Role: "user", Content: mergePrompt}}, + nil, + entry.ModelID, + map[string]any{ + "max_tokens": 1024, + "temperature": 0.3, + "prompt_cache_key": agent.ID, + }, + ) + }() if err == nil { finalSummary = resp.Content } else { @@ -1339,11 +1385,15 @@ func (al *AgentLoop) summarizeBatch( } prompt := sb.String() - response, err := agent.Provider.Chat( + entry, err := al.agentEntry(agent) + if err != nil { + return "", err + } + response, err := entry.Provider.Chat( ctx, []providers.Message{{Role: "user", Content: prompt}}, nil, - agent.Model, + entry.ModelID, map[string]any{ "max_tokens": 1024, "temperature": 0.3, @@ -1409,7 +1459,22 @@ func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage) } switch args[0] { case "models": - return "Available models: configured in config.json per agent", true + defaultAgent := al.registry.GetDefaultAgent() + currentModel := "" + if defaultAgent != nil { + currentModel, _ = defaultAgent.getModelSnapshot() + } + var names []string + for _, m := range al.modelRegistry.ModelNames() { + if currentModel == m { + m += " ✅" + } + names = append(names, m) + } + if len(names) == 0 { + return "No models configured in model_list", true + } + return fmt.Sprintf("Available models:\n• %s\n\nUse /switch model to ", strings.Join(names, "\n• ")), true case "channels": if al.channelManager == nil { return "Channel manager not initialized", true @@ -1439,8 +1504,13 @@ func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage) if defaultAgent == nil { return "No default agent configured", true } - oldModel := defaultAgent.Model - defaultAgent.Model = value + if al.modelRegistry == nil { + return "Model registry not available", true + } + if _, ok := al.modelRegistry.Get(value); !ok { + return fmt.Sprintf("Unknown model '%s'. Available: %s", value, strings.Join(al.modelRegistry.ModelNames(), ", ")), true + } + oldModel := defaultAgent.switchModel(value) return fmt.Sprintf("Switched model from %s to %s", oldModel, value), true case "channel": if al.channelManager == nil { diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 023286f02..94169f6f7 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -5,15 +5,12 @@ import ( "fmt" "os" "path/filepath" - "slices" - "strings" "testing" "time" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/channels" "github.com/sipeed/picoclaw/pkg/config" - "github.com/sipeed/picoclaw/pkg/media" "github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/tools" ) @@ -29,15 +26,21 @@ func (f *fakeChannel) IsAllowed(string) bool { func (f *fakeChannel) IsAllowedSender(sender bus.SenderInfo) bool { return true } func (f *fakeChannel) ReasoningChannelID() string { return f.id } -func newTestAgentLoop( - t *testing.T, -) (al *AgentLoop, cfg *config.Config, msgBus *bus.MessageBus, provider *mockProvider, cleanup func()) { - t.Helper() +// mockRegistry wraps a mockProvider into a minimal ModelRegistry for tests. +func mockRegistry(p providers.LLMProvider) *providers.ModelRegistry { + return providers.NewModelRegistryFromProvider(p, "mock-model") +} + +func TestRecordLastChannel(t *testing.T) { + // Create temp workspace tmpDir, err := os.MkdirTemp("", "agent-test-*") if err != nil { t.Fatalf("Failed to create temp dir: %v", err) } - cfg = &config.Config{ + defer os.RemoveAll(tmpDir) + + // Create test config + cfg := &config.Config{ Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ Workspace: tmpDir, @@ -47,43 +50,74 @@ func newTestAgentLoop( }, }, } - msgBus = bus.NewMessageBus() - provider = &mockProvider{} - al = NewAgentLoop(cfg, msgBus, provider) - return al, cfg, msgBus, provider, func() { os.RemoveAll(tmpDir) } -} -func TestRecordLastChannel(t *testing.T) { - al, cfg, msgBus, provider, cleanup := newTestAgentLoop(t) - defer cleanup() + // Create agent loop + msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) + // Test RecordLastChannel testChannel := "test-channel" - if err := al.RecordLastChannel(testChannel); err != nil { + err = al.RecordLastChannel(testChannel) + if err != nil { t.Fatalf("RecordLastChannel failed: %v", err) } - if got := al.state.GetLastChannel(); got != testChannel { - t.Errorf("Expected channel '%s', got '%s'", testChannel, got) + + // Verify channel was saved + lastChannel := al.state.GetLastChannel() + if lastChannel != testChannel { + t.Errorf("Expected channel '%s', got '%s'", testChannel, lastChannel) } - al2 := NewAgentLoop(cfg, msgBus, provider) - if got := al2.state.GetLastChannel(); got != testChannel { - t.Errorf("Expected persistent channel '%s', got '%s'", testChannel, got) + + // Verify persistence by creating a new agent loop + al2 := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) + if al2.state.GetLastChannel() != testChannel { + t.Errorf("Expected persistent channel '%s', got '%s'", testChannel, al2.state.GetLastChannel()) } } func TestRecordLastChatID(t *testing.T) { - al, cfg, msgBus, provider, cleanup := newTestAgentLoop(t) - defer cleanup() + // Create temp workspace + tmpDir, err := os.MkdirTemp("", "agent-test-*") + if err != nil { + t.Fatalf("Failed to create temp dir: %v", err) + } + defer os.RemoveAll(tmpDir) + // Create test config + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + Model: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + + // Create agent loop + msgBus := bus.NewMessageBus() + provider := &mockProvider{} + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) + + // Test RecordLastChatID testChatID := "test-chat-id-123" - if err := al.RecordLastChatID(testChatID); err != nil { + err = al.RecordLastChatID(testChatID) + if err != nil { t.Fatalf("RecordLastChatID failed: %v", err) } - if got := al.state.GetLastChatID(); got != testChatID { - t.Errorf("Expected chat ID '%s', got '%s'", testChatID, got) + + // Verify chat ID was saved + lastChatID := al.state.GetLastChatID() + if lastChatID != testChatID { + t.Errorf("Expected chat ID '%s', got '%s'", testChatID, lastChatID) } - al2 := NewAgentLoop(cfg, msgBus, provider) - if got := al2.state.GetLastChatID(); got != testChatID { - t.Errorf("Expected persistent chat ID '%s', got '%s'", testChatID, got) + + // Verify persistence by creating a new agent loop + al2 := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) + if al2.state.GetLastChatID() != testChatID { + t.Errorf("Expected persistent chat ID '%s', got '%s'", testChatID, al2.state.GetLastChatID()) } } @@ -110,7 +144,7 @@ func TestNewAgentLoop_StateInitialized(t *testing.T) { // Create agent loop msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) // Verify state manager is initialized if al.state == nil { @@ -145,7 +179,7 @@ func TestToolRegistry_ToolRegistration(t *testing.T) { msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) // Register a custom tool customTool := &mockCustomTool{} @@ -158,7 +192,13 @@ func TestToolRegistry_ToolRegistration(t *testing.T) { toolsList := toolsInfo["names"].([]string) // Check that our custom tool name is in the list - found := slices.Contains(toolsList, "mock_custom") + found := false + for _, name := range toolsList { + if name == "mock_custom" { + found = true + break + } + } if !found { t.Error("Expected custom tool to be registered") } @@ -185,7 +225,7 @@ func TestToolContext_Updates(t *testing.T) { msgBus := bus.NewMessageBus() provider := &simpleMockProvider{response: "OK"} - _ = NewAgentLoop(cfg, msgBus, provider) + _ = NewAgentLoop(cfg, msgBus, mockRegistry(provider)) // Verify that ContextualTool interface is defined and can be implemented // This test validates the interface contract exists @@ -216,7 +256,7 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) // Register a test tool and verify it shows up in startup info testTool := &mockCustomTool{} @@ -227,7 +267,13 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { toolsList := toolsInfo["names"].([]string) // Check that our custom tool name is in the list - found := slices.Contains(toolsList, "mock_custom") + found := false + for _, name := range toolsList { + if name == "mock_custom" { + found = true + break + } + } if !found { t.Error("Expected custom tool to be registered") } @@ -254,7 +300,7 @@ func TestAgentLoop_GetStartupInfo(t *testing.T) { msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) info := al.GetStartupInfo() @@ -301,7 +347,7 @@ func TestAgentLoop_Stop(t *testing.T) { msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) // Note: running is only set to true when Run() is called // We can't test that without starting the event loop @@ -429,7 +475,7 @@ func TestToolResult_SilentToolDoesNotSendUserMessage(t *testing.T) { msgBus := bus.NewMessageBus() provider := &simpleMockProvider{response: "File operation complete"} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) helper := testHelper{al: al} // ReadFileTool returns SilentResult, which should not send user message @@ -471,7 +517,7 @@ func TestToolResult_UserFacingToolDoesSendMessage(t *testing.T) { msgBus := bus.NewMessageBus() provider := &simpleMockProvider{response: "Command output: hello world"} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) helper := testHelper{al: al} // ExecTool returns UserResult, which should send user message @@ -550,7 +596,7 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { successResp: "Recovered from context error", } - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, mockRegistry(provider)) // Inject some history to simulate a full context sessionKey := "test-session-context" @@ -621,7 +667,7 @@ func TestTargetReasoningChannelID_AllChannels(t *testing.T) { }, } - al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) + al := NewAgentLoop(cfg, bus.NewMessageBus(), mockRegistry(&mockProvider{})) chManager, err := channels.NewManager(&config.Config{}, bus.NewMessageBus(), nil) if err != nil { t.Fatalf("Failed to create channel manager: %v", err) @@ -691,7 +737,7 @@ func TestHandleReasoning(t *testing.T) { }, } msgBus := bus.NewMessageBus() - return NewAgentLoop(cfg, msgBus, &mockProvider{}), msgBus + return NewAgentLoop(cfg, msgBus, mockRegistry(&mockProvider{})), msgBus } t.Run("skips when any required field is empty", func(t *testing.T) { @@ -810,142 +856,3 @@ func TestHandleReasoning(t *testing.T) { } }) } - -func TestResolveMediaRefs_ResolvesToBase64(t *testing.T) { - store := media.NewFileMediaStore() - dir := t.TempDir() - - // Create a minimal valid PNG (8-byte header is enough for filetype detection) - pngPath := filepath.Join(dir, "test.png") - // PNG magic: 0x89 P N G \r \n 0x1A \n + minimal IHDR - pngHeader := []byte{ - 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature - 0x00, 0x00, 0x00, 0x0D, // IHDR length - 0x49, 0x48, 0x44, 0x52, // "IHDR" - 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, // 1x1 RGB - 0x00, 0x00, 0x00, // no interlace - 0x90, 0x77, 0x53, 0xDE, // CRC - } - if err := os.WriteFile(pngPath, pngHeader, 0o644); err != nil { - t.Fatal(err) - } - ref, err := store.Store(pngPath, media.MediaMeta{}, "test") - if err != nil { - t.Fatal(err) - } - - messages := []providers.Message{ - {Role: "user", Content: "describe this", Media: []string{ref}}, - } - result := resolveMediaRefs(messages, store, config.DefaultMaxMediaSize) - - if len(result[0].Media) != 1 { - t.Fatalf("expected 1 resolved media, got %d", len(result[0].Media)) - } - if !strings.HasPrefix(result[0].Media[0], "data:image/png;base64,") { - t.Fatalf("expected data:image/png;base64, prefix, got %q", result[0].Media[0][:40]) - } -} - -func TestResolveMediaRefs_SkipsOversizedFile(t *testing.T) { - store := media.NewFileMediaStore() - dir := t.TempDir() - - bigPath := filepath.Join(dir, "big.png") - // Write PNG header + padding to exceed limit - data := make([]byte, 1024+1) // 1KB + 1 byte - copy(data, []byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}) - if err := os.WriteFile(bigPath, data, 0o644); err != nil { - t.Fatal(err) - } - ref, _ := store.Store(bigPath, media.MediaMeta{}, "test") - - messages := []providers.Message{ - {Role: "user", Content: "hi", Media: []string{ref}}, - } - // Use a tiny limit (1KB) so the file is oversized - result := resolveMediaRefs(messages, store, 1024) - - if len(result[0].Media) != 0 { - t.Fatalf("expected 0 media (oversized), got %d", len(result[0].Media)) - } -} - -func TestResolveMediaRefs_SkipsUnknownType(t *testing.T) { - store := media.NewFileMediaStore() - dir := t.TempDir() - - txtPath := filepath.Join(dir, "readme.txt") - if err := os.WriteFile(txtPath, []byte("hello world"), 0o644); err != nil { - t.Fatal(err) - } - ref, _ := store.Store(txtPath, media.MediaMeta{}, "test") - - messages := []providers.Message{ - {Role: "user", Content: "hi", Media: []string{ref}}, - } - result := resolveMediaRefs(messages, store, config.DefaultMaxMediaSize) - - if len(result[0].Media) != 0 { - t.Fatalf("expected 0 media (unknown type), got %d", len(result[0].Media)) - } -} - -func TestResolveMediaRefs_PassesThroughNonMediaRefs(t *testing.T) { - messages := []providers.Message{ - {Role: "user", Content: "hi", Media: []string{"https://example.com/img.png"}}, - } - result := resolveMediaRefs(messages, nil, config.DefaultMaxMediaSize) - - if len(result[0].Media) != 1 || result[0].Media[0] != "https://example.com/img.png" { - t.Fatalf("expected passthrough of non-media:// URL, got %v", result[0].Media) - } -} - -func TestResolveMediaRefs_DoesNotMutateOriginal(t *testing.T) { - store := media.NewFileMediaStore() - dir := t.TempDir() - pngPath := filepath.Join(dir, "test.png") - pngHeader := []byte{ - 0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, - 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, - 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, - 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, 0xDE, - } - os.WriteFile(pngPath, pngHeader, 0o644) - ref, _ := store.Store(pngPath, media.MediaMeta{}, "test") - - original := []providers.Message{ - {Role: "user", Content: "hi", Media: []string{ref}}, - } - originalRef := original[0].Media[0] - - resolveMediaRefs(original, store, config.DefaultMaxMediaSize) - - if original[0].Media[0] != originalRef { - t.Fatal("resolveMediaRefs mutated original message slice") - } -} - -func TestResolveMediaRefs_UsesMetaContentType(t *testing.T) { - store := media.NewFileMediaStore() - dir := t.TempDir() - - // File with JPEG content but stored with explicit content type - jpegPath := filepath.Join(dir, "photo") - jpegHeader := []byte{0xFF, 0xD8, 0xFF, 0xE0} // JPEG magic bytes - os.WriteFile(jpegPath, jpegHeader, 0o644) - ref, _ := store.Store(jpegPath, media.MediaMeta{ContentType: "image/jpeg"}, "test") - - messages := []providers.Message{ - {Role: "user", Content: "hi", Media: []string{ref}}, - } - result := resolveMediaRefs(messages, store, config.DefaultMaxMediaSize) - - if len(result[0].Media) != 1 { - t.Fatalf("expected 1 media, got %d", len(result[0].Media)) - } - if !strings.HasPrefix(result[0].Media[0], "data:image/jpeg;base64,") { - t.Fatalf("expected jpeg prefix, got %q", result[0].Media[0][:30]) - } -} diff --git a/pkg/agent/registry.go b/pkg/agent/registry.go index 77b846832..0fe18afa2 100644 --- a/pkg/agent/registry.go +++ b/pkg/agent/registry.go @@ -17,9 +17,10 @@ type AgentRegistry struct { } // NewAgentRegistry creates a registry from config, instantiating all agents. +// Each agent resolves its own provider + modelID from the ModelRegistry. func NewAgentRegistry( cfg *config.Config, - provider providers.LLMProvider, + modelRegistry *providers.ModelRegistry, ) *AgentRegistry { registry := &AgentRegistry{ agents: make(map[string]*AgentInstance), @@ -32,14 +33,14 @@ func NewAgentRegistry( ID: "main", Default: true, } - instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, provider) + instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, modelRegistry) registry.agents["main"] = instance logger.InfoCF("agent", "Created implicit main agent (no agents.list configured)", nil) } else { for i := range agentConfigs { ac := &agentConfigs[i] id := routing.NormalizeAgentID(ac.ID) - instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, provider) + instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, modelRegistry) registry.agents[id] = instance logger.InfoCF("agent", "Registered agent", map[string]any{ diff --git a/pkg/agent/registry_test.go b/pkg/agent/registry_test.go index 518bb441f..346bcb7ee 100644 --- a/pkg/agent/registry_test.go +++ b/pkg/agent/registry_test.go @@ -24,6 +24,14 @@ func (m *mockRegistryProvider) GetDefaultModel() string { return "mock-model" } +// mockModelRegistry creates a ModelRegistry backed by a mockRegistryProvider. +// Maps "gpt-4" (the default in testCfg) and "claude-opus" (used in model tests). +func mockModelRegistry() *providers.ModelRegistry { + p := &mockRegistryProvider{} + reg := providers.NewModelRegistryFromProvider(p, "gpt-4") + return reg +} + func testCfg(agents []config.AgentConfig) *config.Config { return &config.Config{ Agents: config.AgentsConfig{ @@ -40,7 +48,7 @@ func testCfg(agents []config.AgentConfig) *config.Config { func TestNewAgentRegistry_ImplicitMain(t *testing.T) { cfg := testCfg(nil) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) ids := registry.ListAgentIDs() if len(ids) != 1 || ids[0] != "main" { @@ -61,7 +69,7 @@ func TestNewAgentRegistry_ExplicitAgents(t *testing.T) { {ID: "sales", Default: true, Name: "Sales Bot"}, {ID: "support", Name: "Support Bot"}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) ids := registry.ListAgentIDs() if len(ids) != 2 { @@ -86,7 +94,7 @@ func TestAgentRegistry_GetAgent_Normalize(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "my-agent", Default: true}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) agent, ok := registry.GetAgent("My-Agent") if !ok || agent == nil { @@ -102,7 +110,7 @@ func TestAgentRegistry_GetDefaultAgent(t *testing.T) { {ID: "alpha"}, {ID: "beta", Default: true}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) // GetDefaultAgent first checks for "main", then returns any agent := registry.GetDefaultAgent() @@ -124,7 +132,7 @@ func TestAgentRegistry_CanSpawnSubagent(t *testing.T) { {ID: "child2"}, {ID: "restricted"}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) if !registry.CanSpawnSubagent("parent", "child1") { t.Error("expected parent to be allowed to spawn child1") @@ -151,7 +159,7 @@ func TestAgentRegistry_CanSpawnSubagent_Wildcard(t *testing.T) { }, {ID: "any-agent"}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) if !registry.CanSpawnSubagent("admin", "any-agent") { t.Error("expected wildcard to allow spawning any agent") @@ -166,7 +174,7 @@ func TestAgentInstance_Model(t *testing.T) { cfg := testCfg([]config.AgentConfig{ {ID: "custom", Default: true, Model: model}, }) - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) agent, _ := registry.GetAgent("custom") if agent.Model != "claude-opus" { @@ -179,7 +187,7 @@ func TestAgentInstance_FallbackInheritance(t *testing.T) { {ID: "inherit", Default: true}, }) cfg.Agents.Defaults.ModelFallbacks = []string{"openai/gpt-4o-mini", "anthropic/haiku"} - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) agent, _ := registry.GetAgent("inherit") if len(agent.Fallbacks) != 2 { @@ -196,7 +204,7 @@ func TestAgentInstance_FallbackExplicitEmpty(t *testing.T) { {ID: "no-fallback", Default: true, Model: model}, }) cfg.Agents.Defaults.ModelFallbacks = []string{"should-not-inherit"} - registry := NewAgentRegistry(cfg, &mockRegistryProvider{}) + registry := NewAgentRegistry(cfg, mockModelRegistry()) agent, _ := registry.GetAgent("no-fallback") if len(agent.Fallbacks) != 0 { diff --git a/pkg/providers/fallback.go b/pkg/providers/fallback.go index 7ba563b66..96c1a70d5 100644 --- a/pkg/providers/fallback.go +++ b/pkg/providers/fallback.go @@ -9,13 +9,13 @@ import ( // FallbackChain orchestrates model fallback across multiple candidates. type FallbackChain struct { - cooldown *CooldownTracker + cooldown *CooldownTracker + providerKeyFn func(model string) string // optional: derives a stable provider key for cooldown tracking } -// FallbackCandidate represents one model/provider to try. +// FallbackCandidate represents one model to try, identified by its registry name. type FallbackCandidate struct { - Provider string - Model string + Model string } // FallbackResult contains the successful response and metadata about all attempts. @@ -41,6 +41,22 @@ func NewFallbackChain(cooldown *CooldownTracker) *FallbackChain { return &FallbackChain{cooldown: cooldown} } +// WithProviderKeyFn sets a function that maps a model name to a stable provider key +// used for cooldown tracking. Multiple models on the same provider share one cooldown bucket. +// If not set, the model name itself is used as the key. +func (fc *FallbackChain) WithProviderKeyFn(fn func(model string) string) *FallbackChain { + fc.providerKeyFn = fn + return fc +} + +// providerKey returns the cooldown key for a given model. +func (fc *FallbackChain) providerKey(model string) string { + if fc.providerKeyFn != nil { + return fc.providerKeyFn(model) + } + return model +} + // ResolveCandidates parses model config into a deduplicated candidate list. func ResolveCandidates(cfg ModelConfig, defaultProvider string) []FallbackCandidate { return ResolveCandidatesWithLookup(cfg, defaultProvider, nil) @@ -72,8 +88,7 @@ func ResolveCandidatesWithLookup( } seen[key] = true candidates = append(candidates, FallbackCandidate{ - Provider: ref.Provider, - Model: ref.Model, + Model: ref.Model, }) } @@ -101,7 +116,7 @@ func ResolveCandidatesWithLookup( func (fc *FallbackChain) Execute( ctx context.Context, candidates []FallbackCandidate, - run func(ctx context.Context, provider, model string) (*LLMResponse, error), + run func(ctx context.Context, model string) (*LLMResponse, error), ) (*FallbackResult, error) { if len(candidates) == 0 { return nil, fmt.Errorf("fallback: no candidates configured") @@ -117,17 +132,19 @@ func (fc *FallbackChain) Execute( return nil, context.Canceled } + pk := fc.providerKey(candidate.Model) + // Check cooldown. - if !fc.cooldown.IsAvailable(candidate.Provider) { - remaining := fc.cooldown.CooldownRemaining(candidate.Provider) + if !fc.cooldown.IsAvailable(pk) { + remaining := fc.cooldown.CooldownRemaining(pk) result.Attempts = append(result.Attempts, FallbackAttempt{ - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Skipped: true, Reason: FailoverRateLimit, Error: fmt.Errorf( "provider %s in cooldown (%s remaining)", - candidate.Provider, + pk, remaining.Round(time.Second), ), }) @@ -136,14 +153,14 @@ func (fc *FallbackChain) Execute( // Execute the run function. start := time.Now() - resp, err := run(ctx, candidate.Provider, candidate.Model) + resp, err := run(ctx, candidate.Model) elapsed := time.Since(start) if err == nil { // Success. - fc.cooldown.MarkSuccess(candidate.Provider) + fc.cooldown.MarkSuccess(pk) result.Response = resp - result.Provider = candidate.Provider + result.Provider = pk result.Model = candidate.Model return result, nil } @@ -151,7 +168,7 @@ func (fc *FallbackChain) Execute( // Context cancellation: abort immediately, no fallback. if ctx.Err() == context.Canceled { result.Attempts = append(result.Attempts, FallbackAttempt{ - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Error: err, Duration: elapsed, @@ -160,24 +177,24 @@ func (fc *FallbackChain) Execute( } // Classify the error. - failErr := ClassifyError(err, candidate.Provider, candidate.Model) + failErr := ClassifyError(err, pk, candidate.Model) if failErr == nil { // Unclassifiable error: do not fallback, return immediately. result.Attempts = append(result.Attempts, FallbackAttempt{ - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Error: err, Duration: elapsed, }) return nil, fmt.Errorf("fallback: unclassified error from %s/%s: %w", - candidate.Provider, candidate.Model, err) + pk, candidate.Model, err) } // Non-retriable error: abort immediately. if !failErr.IsRetriable() { result.Attempts = append(result.Attempts, FallbackAttempt{ - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Error: failErr, Reason: failErr.Reason, @@ -187,9 +204,9 @@ func (fc *FallbackChain) Execute( } // Retriable error: mark failure and continue to next candidate. - fc.cooldown.MarkFailure(candidate.Provider, failErr.Reason) + fc.cooldown.MarkFailure(pk, failErr.Reason) result.Attempts = append(result.Attempts, FallbackAttempt{ - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Error: failErr, Reason: failErr.Reason, @@ -212,7 +229,7 @@ func (fc *FallbackChain) Execute( func (fc *FallbackChain) ExecuteImage( ctx context.Context, candidates []FallbackCandidate, - run func(ctx context.Context, provider, model string) (*LLMResponse, error), + run func(ctx context.Context, model string) (*LLMResponse, error), ) (*FallbackResult, error) { if len(candidates) == 0 { return nil, fmt.Errorf("image fallback: no candidates configured") @@ -227,20 +244,22 @@ func (fc *FallbackChain) ExecuteImage( return nil, context.Canceled } + pk := fc.providerKey(candidate.Model) + start := time.Now() - resp, err := run(ctx, candidate.Provider, candidate.Model) + resp, err := run(ctx, candidate.Model) elapsed := time.Since(start) if err == nil { result.Response = resp - result.Provider = candidate.Provider + result.Provider = pk result.Model = candidate.Model return result, nil } if ctx.Err() == context.Canceled { result.Attempts = append(result.Attempts, FallbackAttempt{ - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Error: err, Duration: elapsed, @@ -252,7 +271,7 @@ func (fc *FallbackChain) ExecuteImage( errMsg := strings.ToLower(err.Error()) if IsImageDimensionError(errMsg) || IsImageSizeError(errMsg) { result.Attempts = append(result.Attempts, FallbackAttempt{ - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Error: err, Reason: FailoverFormat, @@ -260,7 +279,7 @@ func (fc *FallbackChain) ExecuteImage( }) return nil, &FailoverError{ Reason: FailoverFormat, - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Wrapped: err, } @@ -268,7 +287,7 @@ func (fc *FallbackChain) ExecuteImage( // Any other error: record and try next. result.Attempts = append(result.Attempts, FallbackAttempt{ - Provider: candidate.Provider, + Provider: pk, Model: candidate.Model, Error: err, Duration: elapsed, diff --git a/pkg/providers/fallback_test.go b/pkg/providers/fallback_test.go index 1783ebcb5..4a6e07e64 100644 --- a/pkg/providers/fallback_test.go +++ b/pkg/providers/fallback_test.go @@ -7,12 +7,12 @@ import ( "time" ) -func makeCandidate(provider, model string) FallbackCandidate { - return FallbackCandidate{Provider: provider, Model: model} +func makeCandidate(model string) FallbackCandidate { + return FallbackCandidate{Model: model} } -func successRun(content string) func(ctx context.Context, provider, model string) (*LLMResponse, error) { - return func(ctx context.Context, provider, model string) (*LLMResponse, error) { +func successRun(content string) func(ctx context.Context, model string) (*LLMResponse, error) { + return func(ctx context.Context, model string) (*LLMResponse, error) { return &LLMResponse{Content: content, FinishReason: "stop"}, nil } } @@ -21,7 +21,7 @@ func TestFallback_SingleCandidate_Success(t *testing.T) { ct := NewCooldownTracker() fc := NewFallbackChain(ct) - candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")} + candidates := []FallbackCandidate{makeCandidate("gpt-4")} result, err := fc.Execute(context.Background(), candidates, successRun("hello")) if err != nil { t.Fatalf("unexpected error: %v", err) @@ -29,8 +29,8 @@ func TestFallback_SingleCandidate_Success(t *testing.T) { if result.Response.Content != "hello" { t.Errorf("content = %q, want hello", result.Response.Content) } - if result.Provider != "openai" || result.Model != "gpt-4" { - t.Errorf("provider/model = %s/%s, want openai/gpt-4", result.Provider, result.Model) + if result.Model != "gpt-4" { + t.Errorf("model = %s, want gpt-4", result.Model) } } @@ -39,12 +39,12 @@ func TestFallback_SecondCandidateSuccess(t *testing.T) { fc := NewFallbackChain(ct) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4"), - makeCandidate("anthropic", "claude-opus"), + makeCandidate("gpt-4"), + makeCandidate("claude-opus"), } attempt := 0 - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { attempt++ if attempt == 1 { return nil, errors.New("rate limit exceeded") @@ -56,8 +56,8 @@ func TestFallback_SecondCandidateSuccess(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } - if result.Provider != "anthropic" { - t.Errorf("provider = %q, want anthropic", result.Provider) + if result.Model != "claude-opus" { + t.Errorf("model = %q, want claude-opus", result.Model) } if result.Response.Content != "from claude" { t.Errorf("content = %q, want 'from claude'", result.Response.Content) @@ -72,12 +72,12 @@ func TestFallback_AllFail(t *testing.T) { fc := NewFallbackChain(ct) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4"), - makeCandidate("anthropic", "claude"), - makeCandidate("groq", "llama"), + makeCandidate("gpt-4"), + makeCandidate("claude"), + makeCandidate("llama"), } - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { return nil, errors.New("rate limit exceeded") } @@ -100,12 +100,12 @@ func TestFallback_ContextCanceled(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4"), - makeCandidate("anthropic", "claude"), + makeCandidate("gpt-4"), + makeCandidate("claude"), } attempt := 0 - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { attempt++ if attempt == 1 { cancel() // cancel context @@ -126,12 +126,12 @@ func TestFallback_NonRetriableError(t *testing.T) { fc := NewFallbackChain(ct) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4"), - makeCandidate("anthropic", "claude"), + makeCandidate("gpt-4"), + makeCandidate("claude"), } attempt := 0 - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { attempt++ return nil, errors.New("string should match pattern") } @@ -157,17 +157,17 @@ func TestFallback_CooldownSkip(t *testing.T) { ct, _ := newTestTracker(now) fc := NewFallbackChain(ct) - // Put openai in cooldown - ct.MarkFailure("openai", FailoverRateLimit) + // Put gpt-4 in cooldown + ct.MarkFailure("gpt-4", FailoverRateLimit) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4"), - makeCandidate("anthropic", "claude"), + makeCandidate("gpt-4"), + makeCandidate("claude"), } - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { - if provider == "openai" { - t.Error("should not call openai (in cooldown)") + run := func(ctx context.Context, model string) (*LLMResponse, error) { + if model == "gpt-4" { + t.Error("should not call gpt-4 (in cooldown)") } return &LLMResponse{Content: "claude response", FinishReason: "stop"}, nil } @@ -176,8 +176,8 @@ func TestFallback_CooldownSkip(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } - if result.Provider != "anthropic" { - t.Errorf("provider = %q, want anthropic", result.Provider) + if result.Model != "claude" { + t.Errorf("model = %q, want claude", result.Model) } // Should have 1 skipped attempt skipped := 0 @@ -195,17 +195,17 @@ func TestFallback_AllInCooldown(t *testing.T) { ct := NewCooldownTracker() fc := NewFallbackChain(ct) - // Put all providers in cooldown - ct.MarkFailure("openai", FailoverRateLimit) - ct.MarkFailure("anthropic", FailoverBilling) + // Put all models in cooldown + ct.MarkFailure("gpt-4", FailoverRateLimit) + ct.MarkFailure("claude", FailoverBilling) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4"), - makeCandidate("anthropic", "claude"), + makeCandidate("gpt-4"), + makeCandidate("claude"), } _, err := fc.Execute(context.Background(), candidates, - func(ctx context.Context, provider, model string) (*LLMResponse, error) { + func(ctx context.Context, model string) (*LLMResponse, error) { t.Error("should not call any provider (all in cooldown)") return nil, nil }) @@ -234,7 +234,7 @@ func TestFallback_EmptyFallbacks(t *testing.T) { ct := NewCooldownTracker() fc := NewFallbackChain(ct) - candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")} + candidates := []FallbackCandidate{makeCandidate("gpt-4")} result, err := fc.Execute(context.Background(), candidates, successRun("ok")) if err != nil { t.Fatalf("unexpected error: %v", err) @@ -249,12 +249,12 @@ func TestFallback_UnclassifiedError(t *testing.T) { fc := NewFallbackChain(ct) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4"), - makeCandidate("anthropic", "claude"), + makeCandidate("gpt-4"), + makeCandidate("claude"), } attempt := 0 - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { attempt++ return nil, errors.New("completely unknown internal error") } @@ -272,13 +272,13 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) { ct := NewCooldownTracker() fc := NewFallbackChain(ct) - candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4")} + candidates := []FallbackCandidate{makeCandidate("gpt-4")} attempt := 0 - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { attempt++ if attempt == 1 { - ct.MarkFailure("openai", FailoverRateLimit) // simulate failure tracked elsewhere + ct.MarkFailure("gpt-4", FailoverRateLimit) // simulate failure tracked elsewhere } return &LLMResponse{Content: "ok", FinishReason: "stop"}, nil } @@ -287,7 +287,7 @@ func TestFallback_SuccessResetsCooldown(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } - if !ct.IsAvailable("openai") { + if !ct.IsAvailable("gpt-4") { t.Error("success should reset cooldown") } } @@ -298,7 +298,7 @@ func TestImageFallback_Success(t *testing.T) { ct := NewCooldownTracker() fc := NewFallbackChain(ct) - candidates := []FallbackCandidate{makeCandidate("openai", "gpt-4o")} + candidates := []FallbackCandidate{makeCandidate("gpt-4o")} result, err := fc.ExecuteImage(context.Background(), candidates, successRun("image result")) if err != nil { t.Fatalf("unexpected error: %v", err) @@ -313,12 +313,12 @@ func TestImageFallback_DimensionError(t *testing.T) { fc := NewFallbackChain(ct) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4o"), - makeCandidate("anthropic", "claude"), + makeCandidate("gpt-4o"), + makeCandidate("claude"), } attempt := 0 - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { attempt++ return nil, errors.New("image dimensions exceed max 4096x4096") } @@ -337,12 +337,12 @@ func TestImageFallback_SizeError(t *testing.T) { fc := NewFallbackChain(ct) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4o"), - makeCandidate("anthropic", "claude"), + makeCandidate("gpt-4o"), + makeCandidate("claude"), } attempt := 0 - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { attempt++ return nil, errors.New("image exceeds 20 mb") } @@ -361,12 +361,12 @@ func TestImageFallback_RetryOnOtherErrors(t *testing.T) { fc := NewFallbackChain(ct) candidates := []FallbackCandidate{ - makeCandidate("openai", "gpt-4o"), - makeCandidate("anthropic", "claude-sonnet"), + makeCandidate("gpt-4o"), + makeCandidate("claude-sonnet"), } attempt := 0 - run := func(ctx context.Context, provider, model string) (*LLMResponse, error) { + run := func(ctx context.Context, model string) (*LLMResponse, error) { attempt++ if attempt == 1 { return nil, errors.New("rate limit exceeded") @@ -378,8 +378,8 @@ func TestImageFallback_RetryOnOtherErrors(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } - if result.Provider != "anthropic" { - t.Errorf("provider = %q, want anthropic", result.Provider) + if result.Model != "claude-sonnet" { + t.Errorf("model = %q, want claude-sonnet", result.Model) } } @@ -406,14 +406,14 @@ func TestResolveCandidates_Simple(t *testing.T) { t.Fatalf("candidates = %d, want 3", len(candidates)) } - if candidates[0].Provider != "openai" || candidates[0].Model != "gpt-4" { - t.Errorf("candidate[0] = %s/%s, want openai/gpt-4", candidates[0].Provider, candidates[0].Model) + if candidates[0].Model != "gpt-4" { + t.Errorf("candidate[0].Model = %q, want gpt-4", candidates[0].Model) } - if candidates[1].Provider != "anthropic" || candidates[1].Model != "claude-opus" { - t.Errorf("candidate[1] = %s/%s, want anthropic/claude-opus", candidates[1].Provider, candidates[1].Model) + if candidates[1].Model != "claude-opus" { + t.Errorf("candidate[1].Model = %q, want claude-opus", candidates[1].Model) } - if candidates[2].Provider != "groq" || candidates[2].Model != "llama-3" { - t.Errorf("candidate[2] = %s/%s, want groq/llama-3", candidates[2].Provider, candidates[2].Model) + if candidates[2].Model != "llama-3" { + t.Errorf("candidate[2].Model = %q, want llama-3", candidates[2].Model) } } @@ -470,9 +470,6 @@ func TestResolveCandidatesWithLookup_AliasResolvesToNestedModel(t *testing.T) { if len(candidates) != 1 { t.Fatalf("candidates = %d, want 1", len(candidates)) } - if candidates[0].Provider != "openrouter" { - t.Fatalf("provider = %q, want openrouter", candidates[0].Provider) - } if candidates[0].Model != "stepfun/step-3.5-flash:free" { t.Fatalf("model = %q, want stepfun/step-3.5-flash:free", candidates[0].Model) } @@ -514,9 +511,6 @@ func TestResolveCandidatesWithLookup_AliasWithoutProtocolUsesDefaultProvider(t * if len(candidates) != 1 { t.Fatalf("candidates = %d, want 1", len(candidates)) } - if candidates[0].Provider != "openai" { - t.Fatalf("provider = %q, want openai", candidates[0].Provider) - } if candidates[0].Model != "glm-5" { t.Fatalf("model = %q, want glm-5", candidates[0].Model) } diff --git a/pkg/providers/registry.go b/pkg/providers/registry.go new file mode 100644 index 000000000..52abde80f --- /dev/null +++ b/pkg/providers/registry.go @@ -0,0 +1,180 @@ +// PicoClaw - Ultra-lightweight personal AI agent +// License: MIT +// +// Copyright (c) 2026 PicoClaw contributors + +package providers + +import ( + "fmt" + + "github.com/sipeed/picoclaw/pkg/config" +) + +// ModelEntry holds the resolved provider and model ID for a configured model. +type ModelEntry struct { + Provider LLMProvider + ModelID string // Protocol-stripped model ID sent to Chat() + ProviderKey string // Protocol string used for cooldown tracking (e.g. "openai", "anthropic") +} + +// ModelRegistry is the single source of truth for all configured models. +// It maps model_name → (provider, modelID). Providers with identical configs +// (protocol + api_base + api_key + auth_method) share a single instance. +type ModelRegistry struct { + models map[string]*ModelEntry + defaultModelName string +} + +// NewModelRegistry builds a registry from config, creating providers for +// every model_list entry. It is built once at startup and passed through +// to the agent layer. +func NewModelRegistry(cfg *config.Config) (*ModelRegistry, error) { + // Ensure model_list is populated from legacy providers config + if cfg.HasProvidersConfig() { + providerModels := config.ConvertProvidersToModelList(cfg) + existingModelNames := make(map[string]bool) + for _, m := range cfg.ModelList { + existingModelNames[m.ModelName] = true + } + for _, pm := range providerModels { + if !existingModelNames[pm.ModelName] { + cfg.ModelList = append(cfg.ModelList, pm) + } + } + } + + if len(cfg.ModelList) == 0 { + return nil, fmt.Errorf("no providers configured. Please add entries to model_list in your config") + } + + reg := &ModelRegistry{ + models: make(map[string]*ModelEntry), + defaultModelName: cfg.Agents.Defaults.GetModelName(), + } + + // Provider cache: models sharing protocol+apiBase+apiKey+authMethod reuse one provider. + type providerKey struct { + protocol string + apiBase string + apiKey string + authMethod string + } + cache := make(map[providerKey]LLMProvider) + + for i := range cfg.ModelList { + mc := &cfg.ModelList[i] + if mc.Model == "" || mc.ModelName == "" { + continue + } + + protocol, modelID := ExtractProtocol(mc.Model) + + apiBase := mc.APIBase + if apiBase == "" { + apiBase = getDefaultAPIBase(protocol) + } + + key := providerKey{ + protocol: protocol, + apiBase: apiBase, + apiKey: mc.APIKey, + authMethod: mc.AuthMethod, + } + + provider, cached := cache[key] + if !cached { + var err error + provider, _, err = CreateProviderFromConfig(mc) + if err != nil { + return nil, fmt.Errorf("failed to create provider for model %q: %w", mc.ModelName, err) + } + cache[key] = provider + } + + providerKey := protocol + if providerKey == "" { + providerKey = "openai" + } + reg.models[mc.ModelName] = &ModelEntry{ + Provider: provider, + ModelID: modelID, + ProviderKey: providerKey, + } + } + + if len(reg.models) == 0 { + return nil, fmt.Errorf("no valid models in model_list") + } + + // Resolve default model name to model ID so the agent gets the right value + if entry, ok := reg.models[reg.defaultModelName]; ok { + _ = entry // default exists, good + } + + return reg, nil +} + +// Get returns the entry for a model_name. +func (r *ModelRegistry) Get(modelName string) (*ModelEntry, bool) { + entry, ok := r.models[modelName] + return entry, ok +} + +// GetDefault returns the entry for the configured default model. +func (r *ModelRegistry) GetDefault() (*ModelEntry, bool) { + return r.Get(r.defaultModelName) +} + +// DefaultModelName returns the configured default model name. +func (r *ModelRegistry) DefaultModelName() string { + return r.defaultModelName +} + +// ModelNames returns all registered model names. +func (r *ModelRegistry) ModelNames() []string { + names := make([]string, 0, len(r.models)) + for name := range r.models { + names = append(names, name) + } + return names +} + +// NewModelRegistryFromProvider wraps a single provider into a minimal +// ModelRegistry. Useful for tests and the CLI agent path where a full +// config-driven registry isn't needed. +// +// modelID may include a protocol prefix (e.g. "anthropic/claude-3-5-sonnet"). +// The protocol is stripped to form both the registry key (user-facing model name) +// and the ModelID passed to Chat(). +func NewModelRegistryFromProvider(provider LLMProvider, modelID string) *ModelRegistry { + protocol, stripped := ExtractProtocol(modelID) + name := stripped + if name == "" { + name = "default" + } + providerKey := protocol + if providerKey == "" { + providerKey = "openai" + } + return &ModelRegistry{ + models: map[string]*ModelEntry{ + name: {Provider: provider, ModelID: stripped, ProviderKey: providerKey}, + }, + defaultModelName: name, + } +} + +// Close closes any stateful providers in the registry. +func (r *ModelRegistry) Close() { + closed := make(map[LLMProvider]bool) + for _, entry := range r.models { + if closed[entry.Provider] { + continue + } + if sp, ok := entry.Provider.(StatefulProvider); ok { + sp.Close() + } + closed[entry.Provider] = true + } +} diff --git a/pkg/tools/spawn_test.go b/pkg/tools/spawn_test.go index 0646c82a9..98b2f140a 100644 --- a/pkg/tools/spawn_test.go +++ b/pkg/tools/spawn_test.go @@ -4,11 +4,13 @@ import ( "context" "strings" "testing" + + "github.com/sipeed/picoclaw/pkg/providers" ) func TestSpawnTool_Execute_EmptyTask(t *testing.T) { provider := &MockLLMProvider{} - manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil) + manager := NewSubagentManager(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/tmp/test", 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", "/tmp/test", nil) + manager := NewSubagentManager(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/tmp/test", nil) tool := NewSpawnTool(manager) ctx := context.Background() diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index 69f1a49a2..3717d10e9 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -23,11 +23,11 @@ type SubagentTask struct { } type SubagentManager struct { - tasks map[string]*SubagentTask - mu sync.RWMutex - provider providers.LLMProvider - defaultModel string - bus *bus.MessageBus + tasks map[string]*SubagentTask + mu sync.RWMutex + registry *providers.ModelRegistry + defaultModelName string // registry key; looked up at spawn time + bus *bus.MessageBus workspace string tools *ToolRegistry maxIterations int @@ -39,19 +39,19 @@ type SubagentManager struct { } func NewSubagentManager( - provider providers.LLMProvider, - defaultModel, workspace string, + registry *providers.ModelRegistry, + defaultModelName, workspace string, bus *bus.MessageBus, ) *SubagentManager { return &SubagentManager{ - tasks: make(map[string]*SubagentTask), - provider: provider, - defaultModel: defaultModel, - bus: bus, - workspace: workspace, - tools: NewToolRegistry(), - maxIterations: 10, - nextID: 1, + tasks: make(map[string]*SubagentTask), + registry: registry, + defaultModelName: defaultModelName, + bus: bus, + workspace: workspace, + tools: NewToolRegistry(), + maxIterations: 10, + nextID: 1, } } @@ -65,7 +65,13 @@ func (sm *SubagentManager) SetLLMOptions(maxTokens int, temperature float64) { sm.hasTemperature = true } -// SetTools sets the tool registry for subagent execution. +// UpdateModel updates the model used for new subagent spawns. +// Called when the user switches models at runtime. +func (sm *SubagentManager) UpdateModel(modelName string) { + sm.mu.Lock() + defer sm.mu.Unlock() + sm.defaultModelName = modelName +} // If not set, subagent will have access to the provided tools. func (sm *SubagentManager) SetTools(tools *ToolRegistry) { sm.mu.Lock() @@ -151,8 +157,21 @@ After completing the task, provide a clear summary of what was done.` temperature := sm.temperature hasMaxTokens := sm.hasMaxTokens hasTemperature := sm.hasTemperature + registry := sm.registry + modelName := sm.defaultModelName sm.mu.RUnlock() + entry, ok := registry.Get(modelName) + if !ok { + if entry, ok = registry.GetDefault(); !ok { + sm.mu.Lock() + task.Status = "failed" + task.Result = fmt.Sprintf("model %q not found in registry", modelName) + sm.mu.Unlock() + return + } + } + var llmOptions map[string]any if hasMaxTokens || hasTemperature { llmOptions = map[string]any{} @@ -165,8 +184,8 @@ After completing the task, provide a clear summary of what was done.` } loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - Model: sm.defaultModel, + Provider: entry.Provider, + Model: entry.ModelID, Tools: tools, MaxIterations: maxIter, LLMOptions: llmOptions, @@ -328,8 +347,17 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe temperature := sm.temperature hasMaxTokens := sm.hasMaxTokens hasTemperature := sm.hasTemperature + registry := sm.registry + modelName := sm.defaultModelName sm.mu.RUnlock() + entry, ok := registry.Get(modelName) + if !ok { + if entry, ok = registry.GetDefault(); !ok { + return ErrorResult(fmt.Sprintf("model %q not found in registry", modelName)).WithError(fmt.Errorf("model not found")) + } + } + var llmOptions map[string]any if hasMaxTokens || hasTemperature { llmOptions = map[string]any{} @@ -342,8 +370,8 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe } loopResult, err := RunToolLoop(ctx, ToolLoopConfig{ - Provider: sm.provider, - Model: sm.defaultModel, + Provider: entry.Provider, + Model: entry.ModelID, Tools: tools, MaxIterations: maxIter, LLMOptions: llmOptions, diff --git a/pkg/tools/subagent_tool_test.go b/pkg/tools/subagent_tool_test.go index 59bfdffae..4bf8d1698 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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/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(providers.NewModelRegistryFromProvider(provider, "test-model"), "test-model", "/tmp/test", msgBus) tool := NewSubagentTool(manager) ctx := context.Background()