fix(agent): preserve cooldown key scoping on model switch

This commit is contained in:
Alix-007 2026-03-25 14:40:34 +08:00
parent 83310d9933
commit 7c5808d645
2 changed files with 87 additions and 0 deletions

View file

@ -3409,6 +3409,7 @@ func (al *AgentLoop) buildCommandsRuntime(agent *AgentInstance, opts *processOpt
if len(nextCandidates) == 0 {
return "", fmt.Errorf("model %q did not resolve to any provider candidates", value)
}
nextCandidates = applyCooldownKeys(cfg, nextCandidates)
oldModel := agent.Model
oldProvider := agent.Provider

View file

@ -1556,6 +1556,92 @@ func TestProcessMessage_SwitchModelShowModelConsistency(t *testing.T) {
}
}
func TestProcessMessage_SwitchModelPreservesPerModelCooldownKeys(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "agent-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,
Provider: "openai",
ModelName: "local",
ModelFallbacks: []string{"shared-provider"},
MaxTokens: 4096,
MaxToolIterations: 10,
},
},
ModelList: []*config.ModelConfig{
{
ModelName: "local",
Model: "openai/local-model",
APIBase: "https://local.example.invalid/v1",
},
{
ModelName: "router-a",
Model: "litellm/openai/gpt-4o-mini",
APIBase: "https://litellm.example.invalid/v1",
CooldownStrategy: "per-model",
},
{
ModelName: "shared-provider",
Model: "litellm/openai/gpt-4.1",
APIBase: "https://litellm.example.invalid/v1",
},
},
}
cfg.WithSecurity(&config.SecurityConfig{
ModelList: map[string]config.ModelSecurityEntry{
"local": {
APIKeys: []string{"test-key"},
},
"router-a": {
APIKeys: []string{"test-key"},
},
"shared-provider": {
APIKeys: []string{"test-key"},
},
},
})
msgBus := bus.NewMessageBus()
provider := &countingMockProvider{response: "LLM reply"}
al := NewAgentLoop(cfg, msgBus, provider)
helper := testHelper{al: al}
switchResp := helper.executeAndGetResponse(t, context.Background(), bus.InboundMessage{
Channel: "telegram",
SenderID: "user1",
ChatID: "chat1",
Content: "/switch model to router-a",
Peer: bus.Peer{
Kind: "direct",
ID: "user1",
},
})
if !strings.Contains(switchResp, "Switched model from local to router-a") {
t.Fatalf("unexpected /switch reply: %q", switchResp)
}
agent := al.GetRegistry().GetDefaultAgent()
if agent == nil {
t.Fatal("default agent is nil")
}
if len(agent.Candidates) != 2 {
t.Fatalf("len(Candidates) = %d, want 2", len(agent.Candidates))
}
if got := agent.Candidates[0].CooldownKey; got != "litellm/openai/gpt-4o-mini" {
t.Fatalf("candidate[0] cooldown key = %q, want %q", got, "litellm/openai/gpt-4o-mini")
}
if got := agent.Candidates[1].CooldownKey; got != "litellm" {
t.Fatalf("candidate[1] cooldown key = %q, want %q", got, "litellm")
}
}
func TestProcessMessage_SwitchModelRejectsUnknownAlias(t *testing.T) {
tmpDir, err := os.MkdirTemp("", "agent-test-*")
if err != nil {