From 67bd324987f89c395a371b36285f55ac460176d6 Mon Sep 17 00:00:00 2001
From: YeonGyu-Kim <code.yeon.gyu@gmail.com>
Date: Tue, 19 May 2026 12:29:59 +0900
Subject: [PATCH] fix(runtime-fallback): keep pending retry state on gate skip

---
 src/hooks/runtime-fallback/auto-retry.test.ts | 78 +++++++++++++++++++
 src/hooks/runtime-fallback/auto-retry.ts      | 28 +++++--
 2 files changed, 100 insertions(+), 6 deletions(-)
 create mode 100644 src/hooks/runtime-fallback/auto-retry.test.ts

diff --git a/src/hooks/runtime-fallback/auto-retry.test.ts b/src/hooks/runtime-fallback/auto-retry.test.ts
new file mode 100644
index 000000000..b8842b1a3
--- /dev/null
+++ b/src/hooks/runtime-fallback/auto-retry.test.ts
@@ -0,0 +1,78 @@
+import { describe, expect, test } from "bun:test"
+
+import { createAutoRetryHelpers } from "./auto-retry"
+import { createFallbackState } from "./fallback-state"
+import type { HookDeps, RuntimeFallbackPluginInput } from "./types"
+
+function createContext(promptCalls: { count: number }): RuntimeFallbackPluginInput {
+  const session = {
+    abort: async () => ({}),
+    messages: async () => ({
+      data: [
+        {
+          info: { role: "user" },
+          parts: [{ type: "text", text: "retry this" }],
+        },
+      ],
+    }),
+    promptAsync: async () => {
+      promptCalls.count += 1
+      return {}
+    },
+    status: async () => ({ data: { "session-auto-retry": { type: "busy" } } }),
+  }
+  return {
+    client: {
+      session,
+      tui: {
+        showToast: async () => ({}),
+      },
+    },
+    directory: "/test/dir",
+  }
+}
+
+function createDeps(promptCalls: { count: number }): HookDeps {
+  return {
+    ctx: createContext(promptCalls),
+    config: {
+      enabled: true,
+      retry_on_errors: [429, 503, 529],
+      max_fallback_attempts: 3,
+      cooldown_seconds: 60,
+      timeout_seconds: 0,
+      notify_on_fallback: false,
+    },
+    options: undefined,
+    pluginConfig: undefined,
+    sessionStates: new Map(),
+    sessionLastAccess: new Map(),
+    sessionRetryInFlight: new Set(),
+    sessionAwaitingFallbackResult: new Set(),
+    sessionFallbackTimeouts: new Map(),
+    sessionStatusRetryKeys: new Map(),
+    internallyAbortedSessions: new Set(),
+  }
+}
+
+describe("createAutoRetryHelpers", () => {
+  test("#given an existing fallback result is pending #when a new fallback retry is skipped by the prompt gate #then the previous pending state is preserved", async () => {
+    // given
+    const promptCalls = { count: 0 }
+    const deps = createDeps(promptCalls)
+    const helpers = createAutoRetryHelpers(deps)
+    const sessionID = "session-auto-retry"
+    const state = createFallbackState("anthropic/claude-opus-4-7")
+    state.pendingFallbackModel = "openai/gpt-5.4"
+    deps.sessionStates.set(sessionID, state)
+    deps.sessionAwaitingFallbackResult.add(sessionID)
+
+    // when
+    await helpers.autoRetryWithFallback(sessionID, "google/gemini-2.5-pro", undefined, "session.status")
+
+    // then
+    expect(promptCalls.count).toBe(0)
+    expect(deps.sessionAwaitingFallbackResult.has(sessionID)).toBe(true)
+    expect(state.pendingFallbackModel).toBe("openai/gpt-5.4")
+  })
+})
diff --git a/src/hooks/runtime-fallback/auto-retry.ts b/src/hooks/runtime-fallback/auto-retry.ts
index 5e907b270..d67d027a0 100644
--- a/src/hooks/runtime-fallback/auto-retry.ts
+++ b/src/hooks/runtime-fallback/auto-retry.ts
@@ -137,6 +137,8 @@ export function createAutoRetryHelpers(deps: HookDeps) {
       return
     }
 
+    const hadAwaitingFallbackResult = sessionAwaitingFallbackResult.has(sessionID)
+    const previousPendingFallbackModel = sessionStates.get(sessionID)?.pendingFallbackModel
     sessionRetryInFlight.add(sessionID)
     let retryDispatched = false
     try {
@@ -154,8 +156,10 @@ export function createAutoRetryHelpers(deps: HookDeps) {
 
         const retryAgent = resolvedAgent ?? getSessionAgent(sessionID)
         const launchAgent = resolveRegisteredAgentName(retryAgent)
-        sessionAwaitingFallbackResult.add(sessionID)
-        scheduleSessionFallbackTimeout(sessionID, retryAgent)
+        if (!hadAwaitingFallbackResult) {
+          sessionAwaitingFallbackResult.add(sessionID)
+          scheduleSessionFallbackTimeout(sessionID, retryAgent)
+        }
 
         const promptResult = await dispatchInternalPrompt({
           mode: "async",
@@ -185,6 +189,10 @@ export function createAutoRetryHelpers(deps: HookDeps) {
           })
           return
         }
+        sessionAwaitingFallbackResult.add(sessionID)
+        if (hadAwaitingFallbackResult) {
+          scheduleSessionFallbackTimeout(sessionID, retryAgent)
+        }
         retryDispatched = true
       } else {
         log(`[${HOOK_NAME}] No user message found for auto-retry (${source})`, { sessionID })
@@ -194,11 +202,19 @@ export function createAutoRetryHelpers(deps: HookDeps) {
     } finally {
       sessionRetryInFlight.delete(sessionID)
       if (!retryDispatched) {
-        sessionAwaitingFallbackResult.delete(sessionID)
-        clearSessionFallbackTimeout(sessionID)
+        if (hadAwaitingFallbackResult) {
+          sessionAwaitingFallbackResult.add(sessionID)
+        } else {
+          sessionAwaitingFallbackResult.delete(sessionID)
+          clearSessionFallbackTimeout(sessionID)
+        }
         const state = sessionStates.get(sessionID)
-        if (state?.pendingFallbackModel) {
-          state.pendingFallbackModel = undefined
+        if (state) {
+          if (hadAwaitingFallbackResult) {
+            state.pendingFallbackModel = previousPendingFallbackModel
+          } else if (state.pendingFallbackModel) {
+            state.pendingFallbackModel = undefined
+          }
         }
       }
     }