fix: make model cooldown recovery query safe
Some checks failed
Lint / build (20) (push) Has been cancelled
Playwright Tests / test (push) Has been cancelled

This commit is contained in:
Your Name 2026-05-14 16:51:44 +08:00
parent 8fd5eae96a
commit 132af95df9

View File

@ -1,6 +1,6 @@
import "server-only";
import { and, desc, eq, inArray, sql } from "drizzle-orm";
import { and, desc, eq, inArray } from "drizzle-orm";
import { db } from "@/lib/db";
import {
aiModel,
@ -23,6 +23,14 @@ const AUTO_DISABLE_FAILURE_THRESHOLD = 3;
const AUTO_DISABLE_COOLDOWN_MS = 2 * 60 * 60 * 1000;
const OBSERVATION_RECOVERY_SUCCESS_THRESHOLD = 3;
const RECENT_EVENT_LIMIT = 80;
const RECOVERABLE_FAILURE_EVENT_TYPES = new Set<AiModelHealthEventType>([
"timeout",
"quota",
"unavailable",
"invalid_model",
"download_error",
"error",
]);
function normalizeText(value: unknown) {
return String(value || "").trim();
@ -190,41 +198,48 @@ export async function reviveAiModelsFromCooldown(params?: { modelId?: string })
const conditions = [
eq(aiModel.isEnabled, false),
eq(aiModel.excludeFromFailover, false),
sql`COALESCE(${aiModel.healthAutoDisabledAt}, ${aiModel.healthLastFailureAt}, ${aiModel.healthUpdatedAt}) <= ${cooldownCutoff}`,
sql`(
${aiModel.healthStatus} = 'auto_disabled'
OR (
${aiModel.healthStatus} = 'degraded'
AND ${aiModel.healthLastFailureAt} IS NOT NULL
AND ${aiModel.healthConsecutiveFailures} > 0
AND ${aiModel.healthLastEventType} IN (
'timeout',
'quota',
'unavailable',
'invalid_model',
'download_error',
'error'
)
)
)`,
];
if (params?.modelId) {
conditions.push(eq(aiModel.id, normalizeText(params.modelId)));
}
const dueModels = await db
const candidates = await db
.select({
id: aiModel.id,
modelId: aiModel.modelId,
displayName: aiModel.displayName,
healthStatus: aiModel.healthStatus,
healthConsecutiveFailures: aiModel.healthConsecutiveFailures,
healthLastEventType: aiModel.healthLastEventType,
healthAutoDisabledAt: aiModel.healthAutoDisabledAt,
healthLastFailureAt: aiModel.healthLastFailureAt,
healthUpdatedAt: aiModel.healthUpdatedAt,
})
.from(aiModel)
.where(and(...conditions));
const dueModels = candidates.filter((model) => {
const cooldownStartedAt =
model.healthAutoDisabledAt ||
model.healthLastFailureAt ||
model.healthUpdatedAt;
if (!cooldownStartedAt || cooldownStartedAt > cooldownCutoff) {
return false;
}
if (model.healthStatus === "auto_disabled") {
return true;
}
return (
model.healthStatus === "degraded" &&
Boolean(model.healthLastFailureAt) &&
(model.healthConsecutiveFailures || 0) > 0 &&
RECOVERABLE_FAILURE_EVENT_TYPES.has(model.healthLastEventType as AiModelHealthEventType)
);
});
if (dueModels.length === 0) {
return [];
}