From 469f11c0bbd4b95616f0a12087e364cf07de4aa1 Mon Sep 17 00:00:00 2001
From: ShutovKS <i@shutovks.ru>
Date: Tue, 14 Jul 2026 02:07:03 +0300
Subject: [PATCH] fix: prune expired model cooldowns from session state

state.failedUntil accumulated one entry per failed model for the lifetime
of a session and was never cleaned up. Expired cooldowns are now pruned on
each session.error, keeping the map bounded.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 src/plugin.ts      | 11 ++++++++++-
 test/index.test.ts | 26 ++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/src/plugin.ts b/src/plugin.ts
index 7d25f19..b1d30fc 100644
--- a/src/plugin.ts
+++ b/src/plugin.ts
@@ -276,6 +276,12 @@ function isEnabled(config: Config, state: SessionState | undefined): boolean {
   return state?.enabledOverride ?? config.enabled
 }
 
+function pruneExpiredCooldowns(state: SessionState, now: number): void {
+  for (const [model, until] of state.failedUntil) {
+    if (until <= now) state.failedUntil.delete(model)
+  }
+}
+
 function nextModel(config: Config, state: SessionState): string | undefined {
   const now = Date.now()
   return config.fallbackModels.find((model) => {
@@ -441,10 +447,13 @@ export function createModelFallbackPlugin(input: RuntimeInput, options: Options
       if (state.awaitingModel && eventModel && eventModel !== state.awaitingModel) return
       state.awaitingModel = undefined
 
+      const now = Date.now()
+      pruneExpiredCooldowns(state, now)
+
       const failedModel = eventModel ?? state.currentModel
       if (failedModel) {
         state.currentModel = failedModel
-        state.failedUntil.set(failedModel, Date.now() + config.cooldownMs)
+        state.failedUntil.set(failedModel, now + config.cooldownMs)
       }
 
       const model = nextModel(config, state)
diff --git a/test/index.test.ts b/test/index.test.ts
index 9e091ab..d789951 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -492,6 +492,32 @@ describe("model fallback plugin", () => {
     expect(prompts).toHaveLength(2)
   })
 
+  test("#given a model cooldown has expired #when a later error fires #then the model is reusable and stale cooldowns are pruned", async () => {
+    const { runtime, prompts } = createRuntime()
+    const plugin = createModelFallbackPlugin(runtime, {
+      fallback_models: ["openai/gpt-5.1-codex", "anthropic/claude-sonnet"],
+      max_attempts: 10,
+      cooldown_ms: 1,
+    })
+    await emitSessionCreated(plugin)
+
+    // opus fails -> switch to codex
+    await emitSessionError(plugin, { providerID: "anthropic", modelID: "claude-opus" })
+    expect(prompts[0]?.body.model).toEqual({ providerID: "openai", modelID: "gpt-5.1-codex" })
+
+    // codex fails -> switch to sonnet
+    await emitSessionError(plugin, { providerID: "openai", modelID: "gpt-5.1-codex" })
+    expect(prompts[1]?.body.model).toEqual({ providerID: "anthropic", modelID: "claude-sonnet" })
+
+    // wait for the 1ms cooldowns to expire so codex is selectable again
+    await new Promise((resolve) => setTimeout(resolve, 5))
+
+    // sonnet fails -> codex cooldown expired, so it is reused
+    await emitSessionError(plugin, { providerID: "anthropic", modelID: "claude-sonnet" })
+    expect(prompts).toHaveLength(3)
+    expect(prompts[2]?.body.model).toEqual({ providerID: "openai", modelID: "gpt-5.1-codex" })
+  })
+
   test("#given first fallback is also unavailable #when retrying a failed model #then it skips to the next fallback", async () => {
     const { runtime, prompts } = createRuntime()
     const plugin = createModelFallbackPlugin(runtime, {