YeonGyu-Kim
1ad5db4e8b
Allow provider cooldown events to override a pending fallback retry so runtime fallback can keep progressing instead of stalling on the same model. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
112 lines
3.5 KiB
TypeScript
112 lines
3.5 KiB
TypeScript
import { describe, expect, it } from "bun:test"
|
|
import type { HookDeps, RuntimeFallbackPluginInput } from "./types"
|
|
import type { AutoRetryHelpers } from "./auto-retry"
|
|
import { createFallbackState } from "./fallback-state"
|
|
import { createSessionStatusHandler } from "./session-status-handler"
|
|
import { SessionCategoryRegistry } from "../../shared/session-category-registry"
|
|
|
|
function createContext(): RuntimeFallbackPluginInput {
|
|
return {
|
|
client: {
|
|
session: {
|
|
abort: async () => ({}),
|
|
messages: async () => ({ data: [] }),
|
|
promptAsync: async () => ({}),
|
|
},
|
|
tui: {
|
|
showToast: async () => ({}),
|
|
},
|
|
},
|
|
directory: "/test/dir",
|
|
}
|
|
}
|
|
|
|
function createDeps(): HookDeps {
|
|
return {
|
|
ctx: createContext(),
|
|
config: {
|
|
enabled: true,
|
|
retry_on_errors: [429, 503, 529],
|
|
max_fallback_attempts: 4,
|
|
cooldown_seconds: 60,
|
|
timeout_seconds: 30,
|
|
notify_on_fallback: false,
|
|
},
|
|
options: undefined,
|
|
pluginConfig: {
|
|
categories: {
|
|
test: {
|
|
fallback_models: ["openai/gpt-5.4", "google/gemini-2.5-pro"],
|
|
},
|
|
},
|
|
},
|
|
sessionStates: new Map(),
|
|
sessionLastAccess: new Map(),
|
|
sessionRetryInFlight: new Set(),
|
|
sessionAwaitingFallbackResult: new Set(),
|
|
sessionFallbackTimeouts: new Map(),
|
|
sessionStatusRetryKeys: new Map(),
|
|
}
|
|
}
|
|
|
|
function createHelpers(abortCalls: string[], retryCalls: Array<{ sessionID: string; model: string; source: string }>): AutoRetryHelpers {
|
|
return {
|
|
abortSessionRequest: async (sessionID: string) => {
|
|
abortCalls.push(sessionID)
|
|
},
|
|
clearSessionFallbackTimeout: () => {},
|
|
scheduleSessionFallbackTimeout: () => {},
|
|
autoRetryWithFallback: async (sessionID: string, model: string, _resolvedAgent: string | undefined, source: string) => {
|
|
retryCalls.push({ sessionID, model, source })
|
|
},
|
|
resolveAgentForSessionFromContext: async () => undefined,
|
|
cleanupStaleSessions: () => {},
|
|
}
|
|
}
|
|
|
|
describe("createSessionStatusHandler", () => {
|
|
it("#given a pending fallback model #when a new provider cooldown retry arrives #then the handler overrides the pending fallback and advances the chain", async () => {
|
|
// given
|
|
SessionCategoryRegistry.clear()
|
|
const sessionID = "session-status-pending-fallback"
|
|
SessionCategoryRegistry.register(sessionID, "test")
|
|
|
|
const deps = createDeps()
|
|
const abortCalls: string[] = []
|
|
const retryCalls: Array<{ sessionID: string; model: string; source: string }> = []
|
|
const state = createFallbackState("anthropic/claude-opus-4-6")
|
|
state.currentModel = "openai/gpt-5.4"
|
|
state.fallbackIndex = 0
|
|
state.attemptCount = 1
|
|
state.pendingFallbackModel = "openai/gpt-5.4"
|
|
state.failedModels.set("anthropic/claude-opus-4-6", Date.now())
|
|
deps.sessionStates.set(sessionID, state)
|
|
|
|
const handler = createSessionStatusHandler(deps, createHelpers(abortCalls, retryCalls), deps.sessionStatusRetryKeys)
|
|
|
|
// when
|
|
await handler({
|
|
sessionID,
|
|
model: "openai/gpt-5.4",
|
|
status: {
|
|
type: "retry",
|
|
attempt: 2,
|
|
message: "All credentials for model gpt-5.4 are cooling down [retrying in 7m 56s attempt #2]",
|
|
},
|
|
})
|
|
|
|
// then
|
|
expect(abortCalls).toEqual([sessionID])
|
|
expect(retryCalls).toEqual([
|
|
{
|
|
sessionID,
|
|
model: "google/gemini-2.5-pro",
|
|
source: "session.status",
|
|
},
|
|
])
|
|
expect(state.currentModel).toBe("google/gemini-2.5-pro")
|
|
expect(state.pendingFallbackModel).toBe("google/gemini-2.5-pro")
|
|
SessionCategoryRegistry.clear()
|
|
})
|
|
})
|