import { getProviderConnections } from "@/lib/db/providers"; import { getSettings } from "@/lib/db/settings"; import { resolveProxyForConnection } from "@/lib/db/settings"; import { extractResolvedProxyConfig } from "@/lib/tokenHealthCheck"; import { refreshAndUpdateCredentials } from "@/lib/usage/providerLimits"; import { runWithProxyContext } from "@omniroute/open-sse/utils/proxyFetch"; import { logger } from "@omniroute/open-sse/utils/logger"; import { matchesCron } from "@/lib/jobs/cronMatch"; import { getCircuitBreakerStore } from "./warmupScheduler/circuitBreakerFactory"; import { TERMINAL_CONNECTION_STATUSES } from "@/lib/quota/connectionRecovery"; import type { WarmupResult, WarmupFailureKind, WarmupTarget } from "./warmupScheduler/core"; export type { WarmupResult, WarmupFailureKind } from "./warmupScheduler/core"; interface WarmupConnection { id: string; provider?: string; authType?: string; email?: string | null; name?: string | null; testStatus?: string | null; accessToken?: string | null; refreshToken?: string | null; tokenExpiresAt?: string | null; providerSpecificData?: unknown; } const log = logger("WarmupScheduler"); const WARMUP_MESSAGES = ["hi", "hello", "ping", "ready"]; let messageCounter = 0; function getWarmupMessage(): string { const msg = WARMUP_MESSAGES[messageCounter % WARMUP_MESSAGES.length]; messageCounter++; return msg; } declare global { var __omnirouteWarmupScheduler: { timer: NodeJS.Timeout | null; executing: boolean; lastFireMinute: number; }; } const STATE = (globalThis.__omnirouteWarmupScheduler ??= { timer: null, executing: false, lastFireMinute: -1, }); const TRUE_ENV_VALUES = new Set(["1", "true", "yes", "on"]); function isEnabled(): boolean { const raw = process.env.OMNIROUTE_WARMUP_ENABLED; return raw ? TRUE_ENV_VALUES.has(raw.trim().toLowerCase()) : false; } function getCron(): string { return process.env.OMNIROUTE_WARMUP_CRON || "0 7 * * *"; } function getConcurrency(): number { const raw = process.env.OMNIROUTE_WARMUP_CONCURRENCY; const parsed = raw ? parseInt(raw, 10) : NaN; return Math.min(10, Math.max(1, Number.isFinite(parsed) ? parsed : 3)); } function toPacificTime(date: Date): Date { const parts = new Intl.DateTimeFormat("en-US", { timeZone: "America/Los_Angeles", hour12: false, hourCycle: "h23", year: "numeric", month: "numeric", day: "numeric", hour: "numeric", minute: "numeric", second: "numeric", }).formatToParts(date); const get = (t: string) => parseInt(parts.find((p) => p.type === t)?.value || "0", 10); return new Date( get("year"), get("month") - 1, get("day"), get("hour"), get("minute"), get("second") ); } export function startWarmupScheduler(): NodeJS.Timeout | null { if (STATE.timer) return STATE.timer; if (!isEnabled()) { log.info("disabled (OMNIROUTE_WARMUP_ENABLED not set)"); return null; } const cron = getCron(); log.info("scheduler started", { cron, concurrency: getConcurrency() }); STATE.timer = setInterval(tick, 60_000); STATE.timer.unref(); tick(); return STATE.timer; } export function stopWarmupScheduler(): void { if (STATE.timer) { clearInterval(STATE.timer); STATE.timer = null; } STATE.executing = false; STATE.lastFireMinute = -1; } /** Test-only: reset the globalThis singleton so each test starts fresh. */ export function __resetWarmupState(): void { if (STATE.timer) { clearInterval(STATE.timer); } STATE.timer = null; STATE.executing = false; STATE.lastFireMinute = -1; } async function tick(): Promise { if (STATE.executing) return; const now = new Date(); const ptNow = toPacificTime(now); if (!matchesCron(getCron(), ptNow)) { STATE.lastFireMinute = -1; return; } const minuteKey = Math.floor(ptNow.getTime() / 60_000); if (minuteKey === STATE.lastFireMinute) return; STATE.lastFireMinute = minuteKey; STATE.executing = true; try { await executeWarmup(); } catch (err) { log.error("tick failed", { err }); } finally { STATE.executing = false; } } async function executeWarmup(): Promise { const settings = await getSettings(); const enabledMap = (settings?.claudeWarmup as Record | undefined)?.connections; const connections = (await getProviderConnections({ provider: "claude", isActive: true, })) as unknown as WarmupConnection[]; const concurrency = getConcurrency(); const cbStore = await getCircuitBreakerStore(); const targets: WarmupTarget[] = []; const headers = await getWarmupHeaders(); for (const conn of connections) { if (enabledMap?.[conn.id] !== true) { log.debug("warmup skip", { connectionId: conn.id, reason: "not opted-in" }); continue; } if (classifyForWarmup(conn) !== "subscription") { log.debug("warmup skip", { connectionId: conn.id, reason: "not subscription" }); continue; } if (conn.testStatus && TERMINAL_CONNECTION_STATUSES.has(conn.testStatus.toLowerCase())) { log.debug("warmup skip", { connectionId: conn.id, reason: "terminal", status: conn.testStatus, }); continue; } if (await cbStore.isInBackoff(conn.id)) { log.debug("warmup skip", { connectionId: conn.id, reason: "backoff" }); continue; } const cbState = await cbStore.get(conn.id); if (cbState?.lastResult === "forbidden") { log.debug("warmup skip", { connectionId: conn.id, reason: "forbidden" }); continue; } const proxyResolution = await resolveProxyForConnection(conn.id).catch((err) => { log.warn("proxy resolution failed, falling back to direct", { connectionId: conn.id, err }); return null; }); const proxyConfig = ( proxyResolution ? extractResolvedProxyConfig(proxyResolution) : null ) as WarmupTarget["proxyConfig"]; targets.push({ connectionId: conn.id, label: conn.email || conn.name || conn.id, accessToken: conn.accessToken, refreshToken: conn.refreshToken, tokenExpiresAt: conn.tokenExpiresAt, authType: conn.authType, providerSpecificData: (conn.providerSpecificData as Record | undefined) ?? undefined, baseUrl: "https://api.anthropic.com/v1/messages", urlSuffix: "?beta=true", headers, proxyConfig, model: process.env.OMNIROUTE_WARMUP_MODEL || "claude-3-5-haiku-20241022", }); } if (targets.length === 0) { log.info("no subscription connections to warm up"); return; } for (let i = 0; i < targets.length; i += concurrency) { const chunk = targets.slice(i, i + concurrency); const results = await Promise.allSettled(chunk.map((t) => executeWarmupTarget(t))); for (let j = 0; j < chunk.length; j++) { const target = chunk[j]; const settled = results[j]; const result: WarmupResult = settled.status === "fulfilled" ? settled.value : { success: false, tokensUsed: 0, durationMs: 0, failureKind: "unknown", error: String(settled.reason), }; try { await cbStore.recordResult(target.connectionId, result); } catch (err) { log.error("persist failed", { err, connectionId: target.connectionId }); } } } } async function getWarmupHeaders(): Promise> { const { getClaudeCliHeaders } = await import("@omniroute/open-sse/config/providers/shared"); return getClaudeCliHeaders(); } type WarmupPath = "subscription" | "skip"; function classifyForWarmup(conn: { provider?: string; authType?: string; accessToken?: string; providerSpecificData?: unknown; }): WarmupPath { if (conn.provider !== "claude") return "skip"; if (conn.authType === "api_key" || conn.authType === "apikey") return "skip"; if (conn.authType !== "oauth") return "skip"; if (!conn.accessToken) return "skip"; const psd = (conn.providerSpecificData as Record | undefined | null) || {}; const orgType = psd.organizationType as string | undefined; const subStatus = psd.subscriptionStatus as string | undefined; if (["claude_pro", "claude_max", "claude_team", "claude_enterprise"].includes(orgType)) { return "subscription"; } if (orgType === "free") return "skip"; if (subStatus === "active") return "subscription"; return "skip"; } async function executeWarmupTarget(target: WarmupTarget): Promise { const start = Date.now(); const message = getWarmupMessage(); const doFetch = async (accessToken: string): Promise => { const fetchFn = () => fetch(`${target.baseUrl}${target.urlSuffix}`, { method: "POST", headers: { ...target.headers, Authorization: `Bearer ${accessToken}`, "content-type": "application/json", }, body: JSON.stringify({ model: target.model, max_tokens: 1, messages: [{ role: "user", content: message }], }), signal: AbortSignal.timeout(10_000), }); return target.proxyConfig ? runWithProxyContext(target.proxyConfig, fetchFn) : fetchFn(); }; const cleanupBody = (resp: Response) => { resp.body?.cancel?.().catch(() => {}); }; try { let resp = await doFetch(target.accessToken); if (resp.status === 401) { const refreshed = await refreshAndUpdateCredentials( { id: target.connectionId, provider: "claude", authType: "oauth", accessToken: target.accessToken, refreshToken: target.refreshToken, tokenExpiresAt: target.tokenExpiresAt, providerSpecificData: target.providerSpecificData, } as any, { allowRotatingRefresh: true, force: true } ).catch(() => null); if (refreshed?.refreshed === true && refreshed.connection?.accessToken) { cleanupBody(resp); resp = await doFetch(refreshed.connection.accessToken); if (resp.ok) { const tokensUsed = await extractTokens(resp); cleanupBody(resp); return { success: true, tokensUsed, durationMs: Date.now() - start, retryAttempted: true, }; } return classifyResponse(resp, start, true); } cleanupBody(resp); return { success: false, tokensUsed: 0, durationMs: Date.now() - start, failureKind: "auth", error: "Token expired (401 after retry)", retryAttempted: true, }; } return classifyResponse(resp, start, false); } catch (err) { const isTimeout = err instanceof Error && err.name === "TimeoutError"; return { success: false, tokensUsed: 0, durationMs: Date.now() - start, failureKind: isTimeout ? "network" : "unknown", error: err instanceof Error ? err.message : "Unknown", }; } } async function classifyResponse( resp: Response, start: number, retryAttempted: boolean ): Promise { const cleanup = () => { resp.body?.cancel?.().catch(() => {}); }; if (resp.ok) { const tokensUsed = await extractTokens(resp); cleanup(); return { success: true, tokensUsed, durationMs: Date.now() - start, retryAttempted }; } if (resp.status === 401) { cleanup(); return { success: false, tokensUsed: 0, durationMs: Date.now() - start, failureKind: "auth", error: "Auth error (401)", retryAttempted, }; } if (resp.status === 403) { cleanup(); return { success: false, tokensUsed: 0, durationMs: Date.now() - start, failureKind: "forbidden", error: "Forbidden (403)", retryAttempted, }; } if (resp.status === 429) { const retryAfter = resp.headers.get("retry-after"); const retryAfterSec = retryAfter ? parseInt(retryAfter, 10) : NaN; cleanup(); return { success: false, tokensUsed: 0, durationMs: Date.now() - start, failureKind: "rate_limit", error: "Rate limited (429)", retryAfterSeconds: Number.isFinite(retryAfterSec) ? retryAfterSec : undefined, retryAttempted, }; } cleanup(); return { success: false, tokensUsed: 0, durationMs: Date.now() - start, failureKind: "unknown", error: `HTTP ${resp.status}`, retryAttempted, }; } async function extractTokens(resp: Response): Promise { try { const body = (await resp.json()) as { usage?: { input_tokens?: number; output_tokens?: number }; }; return (body?.usage?.input_tokens ?? 0) + (body?.usage?.output_tokens ?? 0); } catch { return 5; } }