ai_api / src /lib /quota /connectionRecovery.ts
Yogesh
initial deploy
cd8bd0a
Raw
History Blame Contribute Delete
11.6 kB
/**
* connectionRecovery.ts β€” Proactive recovery of provider connections whose
* transient cooldown has elapsed.
*
* Today the cooldown released by `markAccountUnavailable()` (testStatus
* 'unavailable' + a future `rateLimitedUntil`) is recovered LAZILY: a connection
* only becomes eligible again when the next real request reads it in
* `getProviderCredentials` (src/sse/services/auth.ts). The first request after
* the cooldown window therefore pays the latency of re-discovering a healthy
* connection.
*
* Modeled on gpt-load's CronChecker, this module identifies the subset of
* connections that are cooling down with an already-elapsed window and clears
* their error state OUTSIDE the request hot path, so they are restored before
* the next real request arrives.
*
* This file holds the PURE selection logic (`selectRecoverableConnections`,
* `isRecoverableCooldownConnection`) β€” no DB, no network, time injected β€” plus a
* thin async tick (`runConnectionRecoveryTick`) that wires the helper to the DB.
* The tick is NOT auto-started on import; the caller (startup bootstrap) decides
* when to schedule it, so importing this module in tests never spawns a timer.
*/
import { cooldownUntilMs } from "@omniroute/open-sse/services/accountFallback.ts";
/**
* The transient-cooldown status written by `markAccountUnavailable()` for a
* recoverable failure. Only connections in this status are candidates for
* proactive recovery.
*/
export const RECOVERABLE_COOLDOWN_STATUS = "unavailable";
/**
* Terminal connection statuses that must NEVER be auto-recovered β€” they stay
* unavailable until credentials/settings change or an operator resets them.
* Mirrors `isTerminalConnectionStatus` (src/sse/services/auth.ts) and
* `TERMINAL_STATUSES` (src/lib/db/providers.ts::clearStaleCrashCooldowns).
*/
export const TERMINAL_CONNECTION_STATUSES = new Set<string>([
"banned",
"expired",
"credits_exhausted",
]);
/** Minimal connection shape needed to decide recoverability. */
export interface RecoverableConnectionInput {
id: string;
testStatus?: string | null;
rateLimitedUntil?: string | null;
}
function normalizeStatus(value: string | null | undefined): string {
return (value || "").trim().toLowerCase();
}
/**
* True when `rateLimitedUntil` is set and its instant is at or before `nowMs`
* (the cooldown window has elapsed). Tolerates ISO strings and numeric-epoch
* strings β€” the `rate_limited_until` TEXT column can hold either (#3954).
*/
function hasElapsedCooldown(rateLimitedUntil: string | null | undefined, nowMs: number): boolean {
if (!rateLimitedUntil) return false;
const ms = cooldownUntilMs(rateLimitedUntil);
return Number.isFinite(ms) && ms <= nowMs;
}
/**
* Decide whether a single connection is a proactive-recovery candidate:
* - has a real id, AND
* - testStatus === 'unavailable' (the transient cooldown status), AND
* - rateLimitedUntil is set and already in the past (< nowMs), AND
* - is NOT in a terminal state (banned / expired / credits_exhausted).
*
* Pure β€” `nowMs` is injected so callers/tests control the clock.
*/
export function isRecoverableCooldownConnection(
connection: RecoverableConnectionInput | null | undefined,
nowMs: number
): boolean {
if (!connection || typeof connection.id !== "string" || connection.id.length === 0) {
return false;
}
const status = normalizeStatus(connection.testStatus);
if (status !== RECOVERABLE_COOLDOWN_STATUS) return false;
if (TERMINAL_CONNECTION_STATUSES.has(status)) return false; // defensive; 'unavailable' is never terminal
return hasElapsedCooldown(connection.rateLimitedUntil, nowMs);
}
/**
* From a list of connections, return only those whose transient cooldown has
* elapsed and are safe to restore. Pure, non-mutating, time injected.
*/
export function selectRecoverableConnections<T extends RecoverableConnectionInput>(
connections: readonly T[] | null | undefined,
nowMs: number
): T[] {
if (!Array.isArray(connections)) return [];
return connections.filter((connection) => isRecoverableCooldownConnection(connection, nowMs));
}
/** Result of one recovery tick (handy for logging / tests of the wiring). */
export interface ConnectionRecoveryTickResult {
scanned: number;
recovered: number;
recoveredIds: string[];
}
/**
* Run one proactive-recovery pass: load active provider connections, select the
* subset whose transient cooldown has elapsed, and clear their error state via
* `clearAccountError` so they are eligible again before the next real request.
*
* Dependencies are injected (default to the real DB / auth modules) so the tick
* can be unit-tested without a live database. Best-effort and never throws β€” a
* failure to recover one connection must not abort the others or the scheduler.
*
* NOTE: not auto-started on import. The startup bootstrap is responsible for
* scheduling it (see runConnectionRecoveryTick usage in the report).
*/
export async function runConnectionRecoveryTick(
deps: {
nowMs?: number;
loadConnections?: () => Promise<RecoverableConnectionInput[]>;
clearConnectionError?: (
connectionId: string,
current: RecoverableConnectionInput
) => Promise<void>;
logger?: { info?: (msg: string) => void; warn?: (msg: string) => void };
} = {}
): Promise<ConnectionRecoveryTickResult> {
const nowMs = deps.nowMs ?? Date.now();
const result: ConnectionRecoveryTickResult = { scanned: 0, recovered: 0, recoveredIds: [] };
let connections: RecoverableConnectionInput[];
try {
const load =
deps.loadConnections ??
(async () => {
// Lazy import keeps this module loadable (and the pure helpers testable)
// without a full DB/auth graph.
const { getProviderConnections } = await import("@/lib/db/providers");
const rows = (await getProviderConnections({ isActive: true })) as Array<{
id?: unknown;
testStatus?: unknown;
rateLimitedUntil?: unknown;
}>;
return (Array.isArray(rows) ? rows : []).map((row) => ({
id: typeof row.id === "string" ? row.id : "",
testStatus: typeof row.testStatus === "string" ? row.testStatus : null,
rateLimitedUntil:
typeof row.rateLimitedUntil === "string" ? row.rateLimitedUntil : null,
}));
});
connections = await load();
} catch (err) {
deps.logger?.warn?.(
`[ConnectionRecovery] failed to load connections: ${err instanceof Error ? err.message : String(err)}`
);
return result;
}
result.scanned = connections.length;
const recoverable = selectRecoverableConnections(connections, nowMs);
if (recoverable.length === 0) return result;
const clear =
deps.clearConnectionError ??
(async (connectionId: string, current: RecoverableConnectionInput) => {
const { clearAccountError } = await import("@/sse/services/auth");
await clearAccountError(connectionId, current);
});
for (const connection of recoverable) {
try {
await clear(connection.id, connection);
result.recovered += 1;
result.recoveredIds.push(connection.id);
} catch (err) {
deps.logger?.warn?.(
`[ConnectionRecovery] failed to recover ${connection.id.slice(0, 8)}: ${err instanceof Error ? err.message : String(err)}`
);
}
}
if (result.recovered > 0) {
deps.logger?.info?.(
`[ConnectionRecovery] proactively restored ${result.recovered} connection(s) with elapsed cooldown`
);
}
return result;
}
// ── Scheduler (opt-out, low frequency) ──────────────────────────────────────
// Mirrors src/lib/tokenHealthCheck.ts: a globalThis-guarded singleton so HMR /
// double-import never stacks timers, an unref'd interval so it never holds the
// process open, and a self-disable in build/test processes. NOT auto-started on
// import β€” startup bootstrap calls initConnectionRecoveryScheduler().
const DEFAULT_TICK_MS = 60 * 1000; // re-validate elapsed cooldowns every 60s
const MIN_TICK_MS = 5 * 1000; // floor to avoid hot-looping if misconfigured
const RECOVERY_LOG_PREFIX = "[ConnectionRecovery]";
const TRUE_ENV_VALUES = new Set(["1", "true", "yes", "on"]);
declare global {
var __omnirouteConnRecovery:
| { initialized: boolean; interval: ReturnType<typeof setInterval> | null }
| undefined;
}
function getRecoveryState() {
if (!globalThis.__omnirouteConnRecovery) {
globalThis.__omnirouteConnRecovery = { initialized: false, interval: null };
}
return globalThis.__omnirouteConnRecovery;
}
function isEnvFlagEnabled(name: string): boolean {
const value = typeof process !== "undefined" ? process.env[name] : undefined;
return !!value && TRUE_ENV_VALUES.has(value.trim().toLowerCase());
}
function isBuildProcess(): boolean {
return typeof process !== "undefined" && process.env.NEXT_PHASE === "phase-production-build";
}
function isAutomatedTestProcess(): boolean {
return (
typeof process !== "undefined" &&
(process.env.NODE_ENV === "test" ||
process.env.VITEST !== undefined ||
process.argv.some((arg) => arg.includes("test")))
);
}
function isRecoverySchedulerDisabled(): boolean {
return (
isEnvFlagEnabled("OMNIROUTE_DISABLE_CONNECTION_RECOVERY") ||
isEnvFlagEnabled("OMNIROUTE_DISABLE_BACKGROUND_SERVICES") ||
isBuildProcess() ||
isAutomatedTestProcess()
);
}
/**
* Resolve the tick interval (ms) from OMNIROUTE_CONNECTION_RECOVERY_INTERVAL_MS,
* falling back to the 60s default and clamping to a small floor.
*/
export function resolveConnectionRecoveryIntervalMs(
rawValue: string | undefined = typeof process !== "undefined"
? process.env.OMNIROUTE_CONNECTION_RECOVERY_INTERVAL_MS
: undefined
): number {
if (!rawValue) return DEFAULT_TICK_MS;
const parsed = Number(rawValue);
if (!Number.isFinite(parsed) || parsed <= 0) return DEFAULT_TICK_MS;
return Math.max(MIN_TICK_MS, Math.floor(parsed));
}
/**
* Start the proactive connection-recovery scheduler (idempotent). No-op in
* build/test processes or when disabled via env. Each tick runs
* runConnectionRecoveryTick() against the real DB.
*/
export function initConnectionRecoveryScheduler(): void {
const state = getRecoveryState();
if (state.initialized || isRecoverySchedulerDisabled()) return;
state.initialized = true;
const tickMs = resolveConnectionRecoveryIntervalMs();
const tickLogger = {
info: (msg: string) => console.log(msg),
warn: (msg: string) => console.warn(msg),
};
const runTick = () => {
runConnectionRecoveryTick({ logger: tickLogger }).catch((err: unknown) => {
const msg = err instanceof Error ? err.message : String(err);
console.warn(`${RECOVERY_LOG_PREFIX} tick error (non-fatal): ${msg}`);
});
};
console.log(
`${RECOVERY_LOG_PREFIX} Starting proactive cooldown recovery (tick every ${Math.round(tickMs / 1000)}s)`
);
// Delay the first tick a little so it never piles onto cold-start work.
const timer = setTimeout(() => {
runTick();
state.interval = setInterval(runTick, tickMs);
(state.interval as { unref?: () => void } | undefined)?.unref?.();
}, 15_000);
(timer as { unref?: () => void } | undefined)?.unref?.();
}
/** Stop the scheduler (tests / hot-reload). */
export function stopConnectionRecoveryScheduler(): void {
const state = getRecoveryState();
if (state.interval) {
clearInterval(state.interval);
state.interval = null;
}
state.initialized = false;
}