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;
}