Spaces:
Runtime error
Runtime error
File size: 11,613 Bytes
cd8bd0a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 | /**
* 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;
}
|