From 83310d9933ad366331ca43f59dd432d7fa7ee339 Mon Sep 17 00:00:00 2001 From: Alix-007 <267018309+Alix-007@users.noreply.github.com> Date: Wed, 25 Mar 2026 14:30:29 +0800 Subject: [PATCH] fix(agent): keep multikey cooldown isolation in candidate keying --- pkg/agent/instance.go | 52 ++++++++++++++++++++++++++++++++++++-- pkg/agent/instance_test.go | 43 ++++++++++++++++++++++++++++--- pkg/config/config.go | 10 ++++---- 3 files changed, 94 insertions(+), 11 deletions(-) diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index c3453c45a..da4914212 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -6,6 +6,7 @@ import ( "os" "path/filepath" "regexp" + "strconv" "strings" "github.com/sipeed/picoclaw/pkg/config" @@ -239,7 +240,7 @@ func applyCooldownKeys(cfg *config.Config, candidates []providers.FallbackCandid copy(resolved, candidates) for i := range resolved { - resolved[i].CooldownKey = resolveCooldownKey(strategyLookup, resolved[i]) + resolved[i].CooldownKey = resolveCooldownKey(strategyLookup, resolved[i], resolved) } return resolved @@ -267,7 +268,11 @@ func buildCooldownStrategyLookup(cfg *config.Config) map[string]string { return strategyLookup } -func resolveCooldownKey(strategyLookup map[string]string, candidate providers.FallbackCandidate) string { +func resolveCooldownKey( + strategyLookup map[string]string, + candidate providers.FallbackCandidate, + candidates []providers.FallbackCandidate, +) string { if strings.TrimSpace(candidate.Provider) == "" { if strings.TrimSpace(candidate.Model) == "" { return "" @@ -276,6 +281,9 @@ func resolveCooldownKey(strategyLookup map[string]string, candidate providers.Fa } candidateKey := providers.ModelKey(candidate.Provider, candidate.Model) + if belongsToMultiKeySet(candidate, candidates) { + return candidateKey + } if strategyLookup[candidateKey] == "model" { return candidateKey } @@ -283,6 +291,46 @@ func resolveCooldownKey(strategyLookup map[string]string, candidate providers.Fa return candidate.Provider } +func belongsToMultiKeySet( + candidate providers.FallbackCandidate, + candidates []providers.FallbackCandidate, +) bool { + provider := providers.NormalizeProvider(candidate.Provider) + model := strings.ToLower(strings.TrimSpace(candidate.Model)) + if provider == "" || model == "" { + return false + } + + baseModel, isReplica := multiKeyBaseModel(model) + if isReplica { + return true + } + + for _, other := range candidates { + if providers.NormalizeProvider(other.Provider) != provider { + continue + } + otherBase, otherIsReplica := multiKeyBaseModel(other.Model) + if otherIsReplica && otherBase == baseModel { + return true + } + } + + return false +} + +func multiKeyBaseModel(model string) (string, bool) { + normalized := strings.ToLower(strings.TrimSpace(model)) + idx := strings.LastIndex(normalized, "__key_") + if idx <= 0 { + return normalized, false + } + if _, err := strconv.Atoi(normalized[idx+len("__key_"):]); err != nil { + return normalized, false + } + return normalized[:idx], true +} + // resolveAgentWorkspace determines the workspace directory for an agent. func resolveAgentWorkspace(agentCfg *config.AgentConfig, defaults *config.AgentDefaults) string { if agentCfg != nil && strings.TrimSpace(agentCfg.Workspace) != "" { diff --git a/pkg/agent/instance_test.go b/pkg/agent/instance_test.go index 97c32e4f9..b23ff40b0 100644 --- a/pkg/agent/instance_test.go +++ b/pkg/agent/instance_test.go @@ -172,11 +172,11 @@ func TestNewAgentInstance_ResolvePerModelCooldownKeys(t *testing.T) { Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ Workspace: tmpDir, - Model: "router-a", + ModelName: "router-a", ModelFallbacks: []string{"router-b", "shared-provider"}, }, }, - ModelList: []config.ModelConfig{ + ModelList: []*config.ModelConfig{ { ModelName: "router-a", Model: "litellm/openai/gpt-4o-mini", @@ -218,7 +218,7 @@ func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) { Agents: config.AgentsConfig{ Defaults: config.AgentDefaults{ Workspace: tmpDir, - Model: "primary-model", + ModelName: "primary-model", Routing: &config.RoutingConfig{ Enabled: true, LightModel: "light-model", @@ -226,7 +226,7 @@ func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) { }, }, }, - ModelList: []config.ModelConfig{ + ModelList: []*config.ModelConfig{ { ModelName: "primary-model", Model: "litellm/openai/gpt-4o-mini", @@ -255,6 +255,41 @@ func TestNewAgentInstance_ResolveLightModelPerModelCooldownKeys(t *testing.T) { } } +func TestNewAgentInstance_ResolveMultiKeyCooldownKeys(t *testing.T) { + tmpDir := t.TempDir() + + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "glm-main", + ModelFallbacks: []string{"glm-replica"}, + }, + }, + ModelList: []*config.ModelConfig{ + { + ModelName: "glm-main", + Model: "zhipu/glm-4.7", + }, + { + ModelName: "glm-replica", + Model: "zhipu/glm-4.7__key_1", + }, + }, + } + + agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{}) + if len(agent.Candidates) != 2 { + t.Fatalf("len(Candidates) = %d, want 2", len(agent.Candidates)) + } + if got := agent.Candidates[0].CooldownKey; got != "zhipu/glm-4.7" { + t.Fatalf("candidate[0] cooldown key = %q, want %q", got, "zhipu/glm-4.7") + } + if got := agent.Candidates[1].CooldownKey; got != "zhipu/glm-4.7__key_1" { + t.Fatalf("candidate[1] cooldown key = %q, want %q", got, "zhipu/glm-4.7__key_1") + } +} + func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) { workspace := t.TempDir() mediaDir := media.TempDir() diff --git a/pkg/config/config.go b/pkg/config/config.go index 652651a26..cca48bc99 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -667,10 +667,10 @@ type ModelConfig struct { Workspace string `json:"workspace,omitempty"` // Workspace path for CLI-based providers // Optional optimizations - RPM int `json:"rpm,omitempty"` // Requests per minute limit - MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") - RequestTimeout int `json:"request_timeout,omitempty"` - ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive + RPM int `json:"rpm,omitempty"` // Requests per minute limit + MaxTokensField string `json:"max_tokens_field,omitempty"` // Field name for max tokens (e.g., "max_completion_tokens") + RequestTimeout int `json:"request_timeout,omitempty"` + ThinkingLevel string `json:"thinking_level,omitempty"` // Extended thinking: off|low|medium|high|xhigh|adaptive // CooldownStrategy controls the scope of cooldown tracking for this model. // - "provider" (default): cooldown is shared across all models in the provider // - "model" or "per-model": cooldown is isolated to this specific model @@ -696,7 +696,6 @@ func (c *ModelConfig) APIKey() string { func (c *ModelConfig) IsVirtual() bool { return c.isVirtual } -} // Validate checks if the ModelConfig has all required fields. func (c *ModelConfig) Validate() error { @@ -737,6 +736,7 @@ func (c *ModelConfig) SetAPIKey(value string) { c.APIKeys = append(c.APIKeys, NewSecureString(value)) } } + type GatewayConfig struct { Host string `json:"host" env:"PICOCLAW_GATEWAY_HOST"` Port int `json:"port" env:"PICOCLAW_GATEWAY_PORT"`