diff --git a/cmd/picoclaw/internal/agent/helpers.go b/cmd/picoclaw/internal/agent/helpers.go index c3ddbb77f..648352cc4 100644 --- a/cmd/picoclaw/internal/agent/helpers.go +++ b/cmd/picoclaw/internal/agent/helpers.go @@ -47,9 +47,10 @@ func agentCmd(message, sessionKey, model string, debug bool) error { cfg.Agents.Defaults.ModelName = modelID } + dispatcher := providers.NewProviderDispatcher(cfg) msgBus := bus.NewMessageBus() defer msgBus.Close() - agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) + agentLoop := agent.NewAgentLoop(cfg, msgBus, provider, dispatcher) defer agentLoop.Close() // Print agent startup info (only for interactive mode) diff --git a/cmd/picoclaw/internal/gateway/helpers.go b/cmd/picoclaw/internal/gateway/helpers.go index 3562f03ef..8ef1095ea 100644 --- a/cmd/picoclaw/internal/gateway/helpers.go +++ b/cmd/picoclaw/internal/gateway/helpers.go @@ -81,8 +81,9 @@ func gatewayCmd(debug bool) error { cfg.Agents.Defaults.ModelName = modelID } + dispatcher := providers.NewProviderDispatcher(cfg) msgBus := bus.NewMessageBus() - agentLoop := agent.NewAgentLoop(cfg, msgBus, provider) + agentLoop := agent.NewAgentLoop(cfg, msgBus, provider, dispatcher) // Print agent startup info fmt.Println("\n📦 Agent Status:") diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index f20a56b9c..cbd177708 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -51,6 +51,7 @@ type AgentLoop struct { mu sync.RWMutex // Track active requests for safe provider cleanup activeRequests sync.WaitGroup + dispatcher *providers.ProviderDispatcher } // processOptions configures how a message is processed @@ -80,6 +81,7 @@ func NewAgentLoop( cfg *config.Config, msgBus *bus.MessageBus, provider providers.LLMProvider, + dispatcher *providers.ProviderDispatcher, ) *AgentLoop { registry := NewAgentRegistry(cfg, provider) @@ -105,6 +107,7 @@ func NewAgentLoop( summarizing: sync.Map{}, fallback: fallbackChain, cmdRegistry: commands.NewRegistry(commands.BuiltinDefinitions()), + dispatcher: dispatcher, } return al @@ -421,6 +424,11 @@ func (al *AgentLoop) ReloadProviderAndConfig( 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 // This prevents blocking readers while closing if oldProvider, ok := extractProvider(oldRegistry); ok { @@ -1068,7 +1076,12 @@ func (al *AgentLoop) runLLMIteration( fbResult, fbErr := al.fallback.Execute( ctx, 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) }, ) diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index a6604e87f..b57b635a1 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -50,7 +50,7 @@ func newTestAgentLoop( } msgBus = bus.NewMessageBus() provider = &mockProvider{} - al = NewAgentLoop(cfg, msgBus, provider) + al = NewAgentLoop(cfg, msgBus, provider, nil) 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 { 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 { 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 { 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 { t.Errorf("Expected persistent chat ID '%s', got '%s'", testChatID, got) } @@ -111,7 +111,7 @@ func TestNewAgentLoop_StateInitialized(t *testing.T) { // Create agent loop msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) // Verify state manager is initialized if al.state == nil { @@ -146,7 +146,7 @@ func TestToolRegistry_ToolRegistration(t *testing.T) { msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) // Register a custom tool customTool := &mockCustomTool{} @@ -203,7 +203,7 @@ func TestToolRegistry_GetDefinitions(t *testing.T) { msgBus := bus.NewMessageBus() 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 testTool := &mockCustomTool{} @@ -236,7 +236,7 @@ func TestAgentLoop_GetStartupInfo(t *testing.T) { msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) info := al.GetStartupInfo() @@ -283,7 +283,7 @@ func TestAgentLoop_Stop(t *testing.T) { msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) // Note: running is only set to true when Run() is called // We can't test that without starting the event loop @@ -403,7 +403,7 @@ func TestProcessMessage_UsesRouteSessionKey(t *testing.T) { msgBus := bus.NewMessageBus() provider := &simpleMockProvider{response: "ok"} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) msg := bus.InboundMessage{ Channel: "telegram", @@ -462,7 +462,7 @@ func TestProcessMessage_CommandOutcomes(t *testing.T) { msgBus := bus.NewMessageBus() provider := &countingMockProvider{response: "LLM reply"} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) helper := testHelper{al: al} baseMsg := bus.InboundMessage{ @@ -539,7 +539,7 @@ func TestProcessMessage_SwitchModelShowModelConsistency(t *testing.T) { msgBus := bus.NewMessageBus() provider := &countingMockProvider{response: "LLM reply"} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) helper := testHelper{al: al} switchResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{ @@ -596,7 +596,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, provider, nil) helper := testHelper{al: al} // ReadFileTool returns SilentResult, which should not send user message @@ -638,7 +638,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, provider, nil) helper := testHelper{al: al} // ExecTool returns UserResult, which should send user message @@ -717,7 +717,7 @@ func TestAgentLoop_ContextExhaustionRetry(t *testing.T) { successResp: "Recovered from context error", } - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) // Inject some history to simulate a full context sessionKey := "test-session-context" @@ -803,7 +803,7 @@ func TestProcessDirectWithChannel_TriggersMCPInitialization(t *testing.T) { msgBus := bus.NewMessageBus() provider := &mockProvider{} - al := NewAgentLoop(cfg, msgBus, provider) + al := NewAgentLoop(cfg, msgBus, provider, nil) defer al.Close() 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) if err != nil { t.Fatalf("Failed to create channel manager: %v", err) @@ -915,7 +915,7 @@ func TestHandleReasoning(t *testing.T) { }, } 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) { diff --git a/pkg/providers/dispatch.go b/pkg/providers/dispatch.go new file mode 100644 index 000000000..b345bcac1 --- /dev/null +++ b/pkg/providers/dispatch.go @@ -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 +} diff --git a/pkg/providers/dispatch_test.go b/pkg/providers/dispatch_test.go new file mode 100644 index 000000000..8c3d3fa73 --- /dev/null +++ b/pkg/providers/dispatch_test.go @@ -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 +}