From 7c5808d645cc32d072d64b119a33b4b2d0283aa4 Mon Sep 17 00:00:00 2001
From: Alix-007 <267018309+Alix-007@users.noreply.github.com>
Date: Wed, 25 Mar 2026 14:40:34 +0800
Subject: [PATCH] fix(agent): preserve cooldown key scoping on model switch

---
 pkg/agent/loop.go      |  1 +
 pkg/agent/loop_test.go | 86 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index ef2951365..e7229bb5a 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -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
diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go
index 25d20c689..b7be4b3ed 100644
--- a/pkg/agent/loop_test.go
+++ b/pkg/agent/loop_test.go
@@ -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 {