fix(agent): dispatch per-candidate provider in fallback chain

Previously all agents shared a single LLMProvider instance created from
agents.defaults.model_name. Per-agent model config (agents.list[].model)
only changed the model string passed to Chat() — it never changed which
provider binary was invoked. This caused cross-provider fallback chains
(e.g. gemini-cli falling back to claude-cli) to fail, and made it
impossible to assign different CLI providers to different agents.

Introduces ProviderDispatcher which lazily creates and caches provider
instances keyed by "protocol/modelID". The fallback chain's run closure
now resolves the correct provider via the dispatcher before falling back
to agent.Provider for backward compatibility.

References #1634

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Eric Jacksch 2026-03-16 02:50:10 -04:00
parent f2addff099
commit 2a88370da3
6 changed files with 286 additions and 20 deletions

View file

@ -47,9 +47,10 @@ func agentCmd(message, sessionKey, model string, debug bool) error {
cfg.Agents.Defaults.ModelName = modelID cfg.Agents.Defaults.ModelName = modelID
} }
dispatcher := providers.NewProviderDispatcher(cfg)
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
defer msgBus.Close() defer msgBus.Close()
agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) agentLoop := agent.NewAgentLoop(cfg, msgBus, provider, dispatcher)
defer agentLoop.Close() defer agentLoop.Close()
// Print agent startup info (only for interactive mode) // Print agent startup info (only for interactive mode)

View file

@ -81,8 +81,9 @@ func gatewayCmd(debug bool) error {
cfg.Agents.Defaults.ModelName = modelID cfg.Agents.Defaults.ModelName = modelID
} }
dispatcher := providers.NewProviderDispatcher(cfg)
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) agentLoop := agent.NewAgentLoop(cfg, msgBus, provider, dispatcher)
// Print agent startup info // Print agent startup info
fmt.Println("\n📦 Agent Status:") fmt.Println("\n📦 Agent Status:")

View file

@ -51,6 +51,7 @@ type AgentLoop struct {
mu sync.RWMutex mu sync.RWMutex
// Track active requests for safe provider cleanup // Track active requests for safe provider cleanup
activeRequests sync.WaitGroup activeRequests sync.WaitGroup
dispatcher *providers.ProviderDispatcher
} }
// processOptions configures how a message is processed // processOptions configures how a message is processed
@ -80,6 +81,7 @@ func NewAgentLoop(
cfg *config.Config, cfg *config.Config,
msgBus *bus.MessageBus, msgBus *bus.MessageBus,
provider providers.LLMProvider, provider providers.LLMProvider,
dispatcher *providers.ProviderDispatcher,
) *AgentLoop { ) *AgentLoop {
registry := NewAgentRegistry(cfg, provider) registry := NewAgentRegistry(cfg, provider)
@ -105,6 +107,7 @@ func NewAgentLoop(
summarizing: sync.Map{}, summarizing: sync.Map{},
fallback: fallbackChain, fallback: fallbackChain,
cmdRegistry: commands.NewRegistry(commands.BuiltinDefinitions()), cmdRegistry: commands.NewRegistry(commands.BuiltinDefinitions()),
dispatcher: dispatcher,
} }
return al return al
@ -421,6 +424,11 @@ func (al *AgentLoop) ReloadProviderAndConfig(
al.mu.Unlock() al.mu.Unlock()
// Flush the dispatcher cache so stale providers are evicted on config reload.
if al.dispatcher != nil {
al.dispatcher.Flush(cfg)
}
// Close old provider after releasing the lock // Close old provider after releasing the lock
// This prevents blocking readers while closing // This prevents blocking readers while closing
if oldProvider, ok := extractProvider(oldRegistry); ok { if oldProvider, ok := extractProvider(oldRegistry); ok {
@ -1068,7 +1076,12 @@ func (al *AgentLoop) runLLMIteration(
fbResult, fbErr := al.fallback.Execute( fbResult, fbErr := al.fallback.Execute(
ctx, ctx,
activeCandidates, activeCandidates,
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { func(ctx context.Context, providerName, model string) (*providers.LLMResponse, error) {
if al.dispatcher != nil {
if p, err := al.dispatcher.Get(providerName, model); err == nil {
return p.Chat(ctx, messages, providerToolDefs, model, llmOpts)
}
}
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts) return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
}, },
) )

View file

@ -50,7 +50,7 @@ func newTestAgentLoop(
} }
msgBus = bus.NewMessageBus() msgBus = bus.NewMessageBus()
provider = &mockProvider{} provider = &mockProvider{}
al = NewAgentLoop(cfg, msgBus, provider) al = NewAgentLoop(cfg, msgBus, provider, nil)
return al, cfg, msgBus, provider, func() { os.RemoveAll(tmpDir) } return al, cfg, msgBus, provider, func() { os.RemoveAll(tmpDir) }
} }
@ -65,7 +65,7 @@ func TestRecordLastChannel(t *testing.T) {
if got := al.state.GetLastChannel(); got != testChannel { if got := al.state.GetLastChannel(); got != testChannel {
t.Errorf("Expected channel '%s', got '%s'", testChannel, got) t.Errorf("Expected channel '%s', got '%s'", testChannel, got)
} }
al2 := NewAgentLoop(cfg, msgBus, provider) al2 := NewAgentLoop(cfg, msgBus, provider, nil)
if got := al2.state.GetLastChannel(); got != testChannel { if got := al2.state.GetLastChannel(); got != testChannel {
t.Errorf("Expected persistent channel '%s', got '%s'", testChannel, got) t.Errorf("Expected persistent channel '%s', got '%s'", testChannel, got)
} }
@ -82,7 +82,7 @@ func TestRecordLastChatID(t *testing.T) {
if got := al.state.GetLastChatID(); got != testChatID { if got := al.state.GetLastChatID(); got != testChatID {
t.Errorf("Expected chat ID '%s', got '%s'", testChatID, got) t.Errorf("Expected chat ID '%s', got '%s'", testChatID, got)
} }
al2 := NewAgentLoop(cfg, msgBus, provider) al2 := NewAgentLoop(cfg, msgBus, provider, nil)
if got := al2.state.GetLastChatID(); got != testChatID { if got := al2.state.GetLastChatID(); got != testChatID {
t.Errorf("Expected persistent chat ID '%s', got '%s'", testChatID, got) t.Errorf("Expected persistent chat ID '%s', got '%s'", testChatID, got)
} }
@ -111,7 +111,7 @@ func TestNewAgentLoop_StateInitialized(t *testing.T) {
// Create agent loop // Create agent loop
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &mockProvider{} provider := &mockProvider{}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
// Verify state manager is initialized // Verify state manager is initialized
if al.state == nil { if al.state == nil {
@ -146,7 +146,7 @@ func TestToolRegistry_ToolRegistration(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &mockProvider{} provider := &mockProvider{}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
// Register a custom tool // Register a custom tool
customTool := &mockCustomTool{} customTool := &mockCustomTool{}
@ -203,7 +203,7 @@ func TestToolRegistry_GetDefinitions(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &mockProvider{} provider := &mockProvider{}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
// Register a test tool and verify it shows up in startup info // Register a test tool and verify it shows up in startup info
testTool := &mockCustomTool{} testTool := &mockCustomTool{}
@ -236,7 +236,7 @@ func TestAgentLoop_GetStartupInfo(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &mockProvider{} provider := &mockProvider{}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
info := al.GetStartupInfo() info := al.GetStartupInfo()
@ -283,7 +283,7 @@ func TestAgentLoop_Stop(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &mockProvider{} provider := &mockProvider{}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
// Note: running is only set to true when Run() is called // Note: running is only set to true when Run() is called
// We can't test that without starting the event loop // We can't test that without starting the event loop
@ -403,7 +403,7 @@ func TestProcessMessage_UsesRouteSessionKey(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &simpleMockProvider{response: "ok"} provider := &simpleMockProvider{response: "ok"}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
msg := bus.InboundMessage{ msg := bus.InboundMessage{
Channel: "telegram", Channel: "telegram",
@ -462,7 +462,7 @@ func TestProcessMessage_CommandOutcomes(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &countingMockProvider{response: "LLM reply"} provider := &countingMockProvider{response: "LLM reply"}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
helper := testHelper{al: al} helper := testHelper{al: al}
baseMsg := bus.InboundMessage{ baseMsg := bus.InboundMessage{
@ -539,7 +539,7 @@ func TestProcessMessage_SwitchModelShowModelConsistency(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &countingMockProvider{response: "LLM reply"} provider := &countingMockProvider{response: "LLM reply"}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
helper := testHelper{al: al} helper := testHelper{al: al}
switchResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ switchResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{
@ -596,7 +596,7 @@ func TestToolResult_SilentToolDoesNotSendUserMessage(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &simpleMockProvider{response: "File operation complete"} provider := &simpleMockProvider{response: "File operation complete"}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
helper := testHelper{al: al} helper := testHelper{al: al}
// ReadFileTool returns SilentResult, which should not send user message // ReadFileTool returns SilentResult, which should not send user message
@ -638,7 +638,7 @@ func TestToolResult_UserFacingToolDoesSendMessage(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &simpleMockProvider{response: "Command output: hello world"} provider := &simpleMockProvider{response: "Command output: hello world"}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
helper := testHelper{al: al} helper := testHelper{al: al}
// ExecTool returns UserResult, which should send user message // ExecTool returns UserResult, which should send user message
@ -717,7 +717,7 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) {
successResp: "Recovered from context error", successResp: "Recovered from context error",
} }
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
// Inject some history to simulate a full context // Inject some history to simulate a full context
sessionKey := "test-session-context" sessionKey := "test-session-context"
@ -803,7 +803,7 @@ func TestProcessDirectWithChannel_TriggersMCPInitialization(t *testing.T) {
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
provider := &mockProvider{} provider := &mockProvider{}
al := NewAgentLoop(cfg, msgBus, provider) al := NewAgentLoop(cfg, msgBus, provider, nil)
defer al.Close() defer al.Close()
if al.mcp.hasManager() { if al.mcp.hasManager() {
@ -845,7 +845,7 @@ func TestTargetReasoningChannelID_AllChannels(t *testing.T) {
}, },
} }
al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}) al := NewAgentLoop(cfg, bus.NewMessageBus(), &mockProvider{}, nil)
chManager, err := channels.NewManager(&config.Config{}, bus.NewMessageBus(), nil) chManager, err := channels.NewManager(&config.Config{}, bus.NewMessageBus(), nil)
if err != nil { if err != nil {
t.Fatalf("Failed to create channel manager: %v", err) t.Fatalf("Failed to create channel manager: %v", err)
@ -915,7 +915,7 @@ func TestHandleReasoning(t *testing.T) {
}, },
} }
msgBus := bus.NewMessageBus() msgBus := bus.NewMessageBus()
return NewAgentLoop(cfg, msgBus, &mockProvider{}), msgBus return NewAgentLoop(cfg, msgBus, &mockProvider{}, nil), msgBus
} }
t.Run("skips when any required field is empty", func(t *testing.T) { t.Run("skips when any required field is empty", func(t *testing.T) {

88
pkg/providers/dispatch.go Normal file
View file

@ -0,0 +1,88 @@
// PicoClaw - Ultra-lightweight personal AI agent
// License: MIT
//
// Copyright (c) 2026 PicoClaw contributors
package providers
import (
"fmt"
"sync"
"github.com/sipeed/picoclaw/pkg/config"
)
// ProviderDispatcher creates and caches per-model LLMProvider instances.
// This fixes the bug where all agents share a single provider instance: when a
// fallback chain selects a different protocol/model pair than the agent's default
// provider, the dispatcher creates and caches a dedicated provider for that pair.
//
// Cache key is "protocol/modelID" which matches the Model field in ModelConfig.
// Thread-safe: uses sync.RWMutex with read-locking for cache hits.
type ProviderDispatcher struct {
mu sync.RWMutex
cache map[string]LLMProvider
cfg *config.Config
}
// NewProviderDispatcher creates a new dispatcher with the given config.
func NewProviderDispatcher(cfg *config.Config) *ProviderDispatcher {
return &ProviderDispatcher{
cache: make(map[string]LLMProvider),
cfg: cfg,
}
}
// Get returns a cached or newly created provider for the given protocol+modelID pair.
// It finds the ModelConfig by iterating cfg.ModelList and matching where
// ModelConfig.Model == protocol+"/"+modelID.
// Returns an error if no matching ModelConfig is found or provider creation fails.
func (d *ProviderDispatcher) Get(protocol, modelID string) (LLMProvider, error) {
key := protocol + "/" + modelID
// Fast path: read-lock to check cache.
d.mu.RLock()
if p, ok := d.cache[key]; ok {
d.mu.RUnlock()
return p, nil
}
d.mu.RUnlock()
// Slow path: find config and create provider under write-lock.
d.mu.Lock()
defer d.mu.Unlock()
// Double-check after acquiring write-lock.
if p, ok := d.cache[key]; ok {
return p, nil
}
// Find the matching ModelConfig entry.
var matched *config.ModelConfig
for i := range d.cfg.ModelList {
if d.cfg.ModelList[i].Model == key {
matched = &d.cfg.ModelList[i]
break
}
}
if matched == nil {
return nil, fmt.Errorf("dispatcher: no model_list entry with model=%q", key)
}
provider, _, err := CreateProviderFromConfig(matched)
if err != nil {
return nil, fmt.Errorf("dispatcher: creating provider for %q: %w", key, err)
}
d.cache[key] = provider
return provider, nil
}
// Flush clears the provider cache and updates the config reference.
// Call this after a config reload so the dispatcher picks up new settings.
func (d *ProviderDispatcher) Flush(cfg *config.Config) {
d.mu.Lock()
defer d.mu.Unlock()
d.cache = make(map[string]LLMProvider)
d.cfg = cfg
}

View file

@ -0,0 +1,163 @@
// PicoClaw - Ultra-lightweight personal AI agent
// License: MIT
//
// Copyright (c) 2026 PicoClaw contributors
package providers_test
import (
"context"
"sync"
"testing"
"github.com/sipeed/picoclaw/pkg/config"
"github.com/sipeed/picoclaw/pkg/providers"
)
// minimalCfg builds a *config.Config with one ModelConfig entry that will
// successfully create an antigravity provider (no API key required).
func minimalCfg(modelKey string) *config.Config {
return &config.Config{
ModelList: []config.ModelConfig{
{
ModelName: "test-alias",
Model: modelKey,
},
},
}
}
// TestProviderDispatcher_Get_CachesInstance verifies that calling Get twice with
// the same protocol+modelID returns the exact same provider instance.
func TestProviderDispatcher_Get_CachesInstance(t *testing.T) {
cfg := minimalCfg("antigravity/test")
d := providers.NewProviderDispatcher(cfg)
p1, err := d.Get("antigravity", "test")
if err != nil {
t.Fatalf("first Get: unexpected error: %v", err)
}
if p1 == nil {
t.Fatal("first Get: returned nil provider")
}
p2, err := d.Get("antigravity", "test")
if err != nil {
t.Fatalf("second Get: unexpected error: %v", err)
}
if p1 != p2 {
t.Errorf("expected cached provider instance, got different pointers: %p vs %p", p1, p2)
}
}
// TestProviderDispatcher_Get_UnknownProtocol verifies that Get returns an error
// when no ModelConfig entry matches the requested protocol+modelID.
func TestProviderDispatcher_Get_UnknownProtocol(t *testing.T) {
cfg := minimalCfg("antigravity/test")
d := providers.NewProviderDispatcher(cfg)
_, err := d.Get("unknown-protocol", "no-such-model")
if err == nil {
t.Fatal("expected error for unknown protocol/model, got nil")
}
}
// TestProviderDispatcher_Flush verifies that Flush clears the cache so that a
// subsequent Get creates a new provider instance rather than returning the old one.
func TestProviderDispatcher_Flush(t *testing.T) {
cfg := minimalCfg("antigravity/test")
d := providers.NewProviderDispatcher(cfg)
p1, err := d.Get("antigravity", "test")
if err != nil {
t.Fatalf("pre-flush Get: %v", err)
}
// Flush with the same config (simulating a reload).
d.Flush(cfg)
p2, err := d.Get("antigravity", "test")
if err != nil {
t.Fatalf("post-flush Get: %v", err)
}
if p1 == p2 {
t.Error("expected new provider instance after Flush, but got the same pointer")
}
}
// TestProviderDispatcher_Get_ThreadSafe exercises concurrent Gets to verify
// there are no data races. Run with: go test -race ./pkg/providers/...
func TestProviderDispatcher_Get_ThreadSafe(t *testing.T) {
cfg := minimalCfg("antigravity/concurrent")
d := providers.NewProviderDispatcher(cfg)
const goroutines = 20
var wg sync.WaitGroup
wg.Add(goroutines)
for i := 0; i < goroutines; i++ {
go func() {
defer wg.Done()
p, err := d.Get("antigravity", "concurrent")
if err != nil {
t.Errorf("concurrent Get: %v", err)
return
}
// Exercise the provider slightly to ensure no race on the cached value.
_ = p.GetDefaultModel()
}()
}
wg.Wait()
}
// TestProviderDispatcher_Get_FlushRace exercises concurrent Gets and Flushes
// together to verify the mutex correctly protects both operations.
func TestProviderDispatcher_Get_FlushRace(t *testing.T) {
cfg := minimalCfg("antigravity/race")
d := providers.NewProviderDispatcher(cfg)
var wg sync.WaitGroup
// Half goroutines call Get, half call Flush.
for i := 0; i < 10; i++ {
wg.Add(1)
go func() {
defer wg.Done()
_, _ = d.Get("antigravity", "race")
}()
wg.Add(1)
go func() {
defer wg.Done()
d.Flush(cfg)
}()
}
wg.Wait()
}
// Compile-time check: antigravity provider satisfies LLMProvider.
var _ providers.LLMProvider = (func() providers.LLMProvider {
p, _, _ := providers.CreateProviderFromConfig(&config.ModelConfig{Model: "antigravity/x"})
return p
})()
// TestAntigravityProvider_Chat is a lightweight smoke test confirming the
// antigravity provider (used in dispatcher tests) satisfies the interface.
func TestAntigravityProvider_Chat(t *testing.T) {
cfg := minimalCfg("antigravity/smoke")
d := providers.NewProviderDispatcher(cfg)
p, err := d.Get("antigravity", "smoke")
if err != nil {
t.Fatalf("Get: %v", err)
}
// The antigravity provider's Chat is a no-op stub; just ensure it doesn't panic.
resp, err := p.Chat(context.Background(), nil, nil, "smoke", nil)
// antigravity may return nil response + nil error or an error; either is fine.
_ = resp
_ = err
}