fix(provider): implement proper provider-per-candidate fallback architecture

This commit fixes a critical bug in the model fallback chain where the
`provider` parameter was being completely ignored, causing all fallback
candidates to use the same provider instance configured for the primary
model. This led to requests being sent to the wrong API endpoints.

## Root Cause

The fallback callback at loop.go:522-527 was using `agent.Provider.Chat()`
for ALL candidates, ignoring the `provider` parameter passed by the fallback
chain. For example:
- Primary: cerebras/gpt-oss-120b → agent.Provider configured for api.cerebras.ai
- Fallback: ollama/llama3.2 → STILL uses cerebras provider → wrong endpoint!

Additionally, the `prompt_cache_key` parameter (an OpenAI-only feature) was
being sent to all providers, causing 422 errors from providers like Cerebras
that don't support this parameter.

## Changes

### Core Architecture Fix
1. **ProviderRegistry** (pkg/providers/registry.go): New registry that creates
   and caches provider instances per provider name (e.g., "cerebras", "ollama")
2. **AgentInstance** (pkg/agent/instance.go): Added ProviderRegistry field
3. **Fallback callback** (pkg/agent/loop.go:522-544): Now gets correct provider
   per candidate via agent.ProviderRegistry.GetProvider(providerName)

### Prompt Cache Fix
1. **Agent loop** (pkg/agent/loop.go): Removed `prompt_cache_key` from generic
   options map (lines 526, 543, 1009, 1060)
2. **OpenAI compat provider** (pkg/providers/openai_compat/provider.go:158-162):
   Changed from blacklist (exclude Gemini) to whitelist (only api.openai.com)

### Test Updates
- Updated tests to pass new ProviderRegistry parameter

## Impact

- Fallback chains with mixed providers (cerebras + ollama + openai) now work
- Providers like Cerebras no longer receive unsupported `prompt_cache_key`
- Each fallback candidate uses its correctly configured provider

Fixes: Model fallback chain sending requests to wrong API endpoints
Related: 422 errors from Cerebras due to unsupported prompt_cache_key
This commit is contained in:
Vishnuvardhan Reddy 2026-02-26 19:14:52 +00:00
parent 5f6a6b5829
commit af8aa29e05
7 changed files with 174 additions and 55 deletions

View file

@ -24,7 +24,8 @@ type AgentInstance struct {
MaxTokens int MaxTokens int
Temperature float64 Temperature float64
ContextWindow int ContextWindow int
Provider providers.LLMProvider Provider providers.LLMProvider // Default provider for backward compatibility
ProviderRegistry *providers.ProviderRegistry // Registry for multi-provider fallback
Sessions *session.SessionManager Sessions *session.SessionManager
ContextBuilder *ContextBuilder ContextBuilder *ContextBuilder
Tools *tools.ToolRegistry Tools *tools.ToolRegistry
@ -39,6 +40,7 @@ func NewAgentInstance(
defaults *config.AgentDefaults, defaults *config.AgentDefaults,
cfg *config.Config, cfg *config.Config,
provider providers.LLMProvider, provider providers.LLMProvider,
providerRegistry *providers.ProviderRegistry,
) *AgentInstance { ) *AgentInstance {
workspace := resolveAgentWorkspace(agentCfg, defaults) workspace := resolveAgentWorkspace(agentCfg, defaults)
os.MkdirAll(workspace, 0o755) os.MkdirAll(workspace, 0o755)
@ -109,6 +111,7 @@ func NewAgentInstance(
Temperature: temperature, Temperature: temperature,
ContextWindow: maxTokens, ContextWindow: maxTokens,
Provider: provider, Provider: provider,
ProviderRegistry: providerRegistry,
Sessions: sessionsManager, Sessions: sessionsManager,
ContextBuilder: contextBuilder, ContextBuilder: contextBuilder,
Tools: toolsRegistry, Tools: toolsRegistry,

View file

@ -133,7 +133,7 @@ func TestNewAgentInstance_ModelResolution(t *testing.T) {
// Use the existing mockProvider from mock_provider_test.go // Use the existing mockProvider from mock_provider_test.go
provider := &mockProvider{} provider := &mockProvider{}
instance := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) instance := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider, nil)
require.NotNil(t, instance) require.NotNil(t, instance)

View file

@ -29,7 +29,7 @@ func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) {
cfg.Agents.Defaults.Temperature = &configuredTemp cfg.Agents.Defaults.Temperature = &configuredTemp
provider := &mockProvider{} provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider, nil)
if agent.MaxTokens != 1234 { if agent.MaxTokens != 1234 {
t.Fatalf("MaxTokens = %d, want %d", agent.MaxTokens, 1234) t.Fatalf("MaxTokens = %d, want %d", agent.MaxTokens, 1234)
@ -61,7 +61,7 @@ func TestNewAgentInstance_DefaultsTemperatureWhenZero(t *testing.T) {
cfg.Agents.Defaults.Temperature = &configuredTemp cfg.Agents.Defaults.Temperature = &configuredTemp
provider := &mockProvider{} provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider, nil)
if agent.Temperature != 0.0 { if agent.Temperature != 0.0 {
t.Fatalf("Temperature = %f, want %f", agent.Temperature, 0.0) t.Fatalf("Temperature = %f, want %f", agent.Temperature, 0.0)
@ -87,7 +87,7 @@ func TestNewAgentInstance_DefaultsTemperatureWhenUnset(t *testing.T) {
} }
provider := &mockProvider{} provider := &mockProvider{}
agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider) agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, provider, nil)
if agent.Temperature != 0.7 { if agent.Temperature != 0.7 {
t.Fatalf("Temperature = %f, want %f", agent.Temperature, 0.7) t.Fatalf("Temperature = %f, want %f", agent.Temperature, 0.7)

View file

@ -519,11 +519,26 @@ func (al *AgentLoop) runLLMIteration(
callLLM := func() (*providers.LLMResponse, error) { callLLM := func() (*providers.LLMResponse, error) {
if len(agent.Candidates) > 1 && al.fallback != nil { if len(agent.Candidates) > 1 && al.fallback != nil {
fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates, fbResult, fbErr := al.fallback.Execute(ctx, agent.Candidates,
func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) { func(ctx context.Context, providerName, model string) (*providers.LLMResponse, error) {
return agent.Provider.Chat(ctx, messages, providerToolDefs, model, map[string]any{ // Get the correct provider for this candidate
// This is critical: different candidates may use different providers
// (e.g., cerebras, ollama, openai), and each requires its own
// provider instance with the correct API endpoint configuration.
var provider providers.LLMProvider
if agent.ProviderRegistry != nil {
p, err := agent.ProviderRegistry.GetProvider(providerName)
if err != nil {
return nil, fmt.Errorf("failed to get provider for %q: %w", providerName, err)
}
provider = p
} else {
// Fallback to default provider if registry not available
// (this shouldn't happen in normal operation)
provider = agent.Provider
}
return provider.Chat(ctx, messages, providerToolDefs, model, map[string]any{
"max_tokens": agent.MaxTokens, "max_tokens": agent.MaxTokens,
"temperature": agent.Temperature, "temperature": agent.Temperature,
"prompt_cache_key": agent.ID,
}) })
}, },
) )
@ -540,7 +555,6 @@ func (al *AgentLoop) runLLMIteration(
return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{ return agent.Provider.Chat(ctx, messages, providerToolDefs, agent.Model, map[string]any{
"max_tokens": agent.MaxTokens, "max_tokens": agent.MaxTokens,
"temperature": agent.Temperature, "temperature": agent.Temperature,
"prompt_cache_key": agent.ID,
}) })
} }
@ -1006,7 +1020,6 @@ func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) {
map[string]any{ map[string]any{
"max_tokens": 1024, "max_tokens": 1024,
"temperature": 0.3, "temperature": 0.3,
"prompt_cache_key": agent.ID,
}, },
) )
if err == nil { if err == nil {
@ -1057,7 +1070,6 @@ func (al *AgentLoop) summarizeBatch(
map[string]any{ map[string]any{
"max_tokens": 1024, "max_tokens": 1024,
"temperature": 0.3, "temperature": 0.3,
"prompt_cache_key": agent.ID,
}, },
) )
if err != nil { if err != nil {

View file

@ -26,20 +26,23 @@ func NewAgentRegistry(
resolver: routing.NewRouteResolver(cfg), resolver: routing.NewRouteResolver(cfg),
} }
// Create provider registry from config for proper fallback support
providerRegistry := providers.NewProviderRegistry(cfg)
agentConfigs := cfg.Agents.List agentConfigs := cfg.Agents.List
if len(agentConfigs) == 0 { if len(agentConfigs) == 0 {
implicitAgent := &config.AgentConfig{ implicitAgent := &config.AgentConfig{
ID: "main", ID: "main",
Default: true, Default: true,
} }
instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, provider) instance := NewAgentInstance(implicitAgent, &cfg.Agents.Defaults, cfg, provider, providerRegistry)
registry.agents["main"] = instance registry.agents["main"] = instance
logger.InfoCF("agent", "Created implicit main agent (no agents.list configured)", nil) logger.InfoCF("agent", "Created implicit main agent (no agents.list configured)", nil)
} else { } else {
for i := range agentConfigs { for i := range agentConfigs {
ac := &agentConfigs[i] ac := &agentConfigs[i]
id := routing.NormalizeAgentID(ac.ID) id := routing.NormalizeAgentID(ac.ID)
instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, provider) instance := NewAgentInstance(ac, &cfg.Agents.Defaults, cfg, provider, providerRegistry)
registry.agents[id] = instance registry.agents[id] = instance
logger.InfoCF("agent", "Registered agent", logger.InfoCF("agent", "Registered agent",
map[string]any{ map[string]any{

View file

@ -153,10 +153,11 @@ func (p *Provider) Chat(
// with the same key and reuse prefix KV cache across calls. // with the same key and reuse prefix KV cache across calls.
// The key is typically the agent ID — stable per agent, shared across requests. // The key is typically the agent ID — stable per agent, shared across requests.
// See: https://platform.openai.com/docs/guides/prompt-caching // See: https://platform.openai.com/docs/guides/prompt-caching
// Prompt caching is only supported by OpenAI-native endpoints. // IMPORTANT: Prompt caching is ONLY supported by OpenAI-native endpoints (api.openai.com).
// Gemini and other providers reject unknown fields, so skip for non-OpenAI APIs. // Gemini, Cerebras, and other providers reject unknown fields, so we must skip them.
if cacheKey, ok := options["prompt_cache_key"].(string); ok && cacheKey != "" { if cacheKey, ok := options["prompt_cache_key"].(string); ok && cacheKey != "" {
if !strings.Contains(p.apiBase, "generativelanguage.googleapis.com") { // Only send prompt_cache_key to actual OpenAI endpoints
if strings.Contains(p.apiBase, "api.openai.com") {
requestBody["prompt_cache_key"] = cacheKey requestBody["prompt_cache_key"] = cacheKey
} }
} }

100
pkg/providers/registry.go Normal file
View file

@ -0,0 +1,100 @@
// PicoClaw - Ultra-lightweight personal AI agent
// License: MIT
//
// Copyright (c) 2026 PicoClaw contributors
package providers
import (
"fmt"
"strings"
"sync"
"github.com/sipeed/picoclaw/pkg/config"
)
// ProviderRegistry manages provider instances for different provider names.
// It lazily creates providers on-demand and caches them for reuse.
// This is essential for proper fallback behavior where different candidates
// may use different providers (e.g., cerebras, ollama, openai).
type ProviderRegistry struct {
cfg *config.Config
modelList []config.ModelConfig
providers map[string]LLMProvider
mu sync.RWMutex
}
// NewProviderRegistry creates a new provider registry with the given config.
func NewProviderRegistry(cfg *config.Config) *ProviderRegistry {
var modelList []config.ModelConfig
if cfg != nil {
modelList = cfg.ModelList
}
return &ProviderRegistry{
cfg: cfg,
modelList: modelList,
providers: make(map[string]LLMProvider),
}
}
// GetProvider returns a provider for the given provider name (e.g., "openai", "cerebras", "ollama").
// If the provider has already been created, it returns the cached instance.
// Otherwise, it creates a new provider from the model list config.
func (pr *ProviderRegistry) GetProvider(providerName string) (LLMProvider, error) {
// Normalize provider name (lowercase)
providerName = strings.ToLower(providerName)
// Check cache first
pr.mu.RLock()
if provider, ok := pr.providers[providerName]; ok {
pr.mu.RUnlock()
return provider, nil
}
pr.mu.RUnlock()
// Find the model config for this provider
var modelCfg *config.ModelConfig
for i := range pr.modelList {
cfg := &pr.modelList[i]
// Extract provider from model string (e.g., "cerebras/gpt-oss-120b" -> "cerebras")
protocol, _ := ExtractProtocol(cfg.Model)
if strings.EqualFold(protocol, providerName) {
modelCfg = cfg
break
}
}
// If not found in model list, try creating with protocol-only config
if modelCfg == nil {
modelCfg = &config.ModelConfig{
Model: providerName + "/dummy", // Protocol is what matters
}
}
// Create the provider
provider, _, err := CreateProviderFromConfig(modelCfg)
if err != nil {
return nil, fmt.Errorf("failed to create provider for %q: %w", providerName, err)
}
// Cache the provider
pr.mu.Lock()
pr.providers[providerName] = provider
pr.mu.Unlock()
return provider, nil
}
// GetDefaultProvider returns the provider for the default/primary protocol.
// This is used for backward compatibility with code that expects a single provider.
func (pr *ProviderRegistry) GetDefaultProvider() (LLMProvider, error) {
if len(pr.modelList) == 0 {
// No model list configured, return error
// Caller should handle this by using a default provider
return nil, fmt.Errorf("no model list configured")
}
// Use the first model's provider as default
protocol, _ := ExtractProtocol(pr.modelList[0].Model)
return pr.GetProvider(protocol)
}