From 132af95df9526fa88bb87e7e18a07a68ac24e464 Mon Sep 17 00:00:00 2001
From: Your Name <your.email@example.com>
Date: Thu, 14 May 2026 16:51:44 +0800
Subject: [PATCH] fix: make model cooldown recovery query safe

---
 lib/ai/runtime/model-health.ts | 53 ++++++++++++++++++++++------------
 1 file changed, 34 insertions(+), 19 deletions(-)

diff --git a/lib/ai/runtime/model-health.ts b/lib/ai/runtime/model-health.ts
index 9722e48f8..a06bec3aa 100644
--- a/lib/ai/runtime/model-health.ts
+++ b/lib/ai/runtime/model-health.ts
@@ -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 [];
   }