feat: implement model lock functionality for connection management

src/sse/services/auth.js

@@ -1,67 +1,11 @@
 import { getProviderConnections, validateApiKey, updateProviderConnection, getSettings } from "@/lib/localDb";
-import { isAccountUnavailable, getUnavailableUntil, getEarliestRateLimitedUntil, formatRetryAfter, checkFallbackError } from "open-sse/services/accountFallback.js";
+import { formatRetryAfter, checkFallbackError, isModelLockActive, buildModelLockUpdate, buildClearModelLocksUpdate, getEarliestModelLockUntil } from "open-sse/services/accountFallback.js";
 import { resolveProviderId } from "@/shared/constants/providers.js";
 import * as log from "../utils/logger.js";
 
 // Mutex to prevent race conditions during account selection
 let selectionMutex = Promise.resolve();
 
-// ---------------------------------------------------------------------------
-// Model-level rate-limit locking (in-memory)
-//
-// Providers like Antigravity maintain *separate* quota buckets per model
-// family (e.g. Claude vs Gemini). When a 429 arrives for one model, the
-// standard account-level DB lock would block ALL models on that account,
-// wasting quota that is still available for other model families.
-//
-// This module tracks model-specific locks in memory so that only the
-// affected model is skipped during account selection while the rest of the
-// account's quota remains accessible.
-//
-// Locks are intentionally in-memory: they clear on restart and require no
-// database schema migration.
-// ---------------------------------------------------------------------------
-
-/** Providers known to have independent per-model quota buckets */
-const MULTI_BUCKET_PROVIDERS = new Set(["antigravity"]);
-
-/** Map<"connectionId:model", expiryTimestamp> */
-const modelLocks = new Map();
-
-/** Default lock duration for model-level rate limits (5 minutes) */
-const DEFAULT_MODEL_LOCK_MS = 5 * 60 * 1000;
-
-/**
- * Check whether a specific model is temporarily locked on a connection.
- * Expired locks are cleaned up lazily.
- */
-function isModelLocked(connectionId, model) {
-  if (!connectionId || !model) return false;
-  const key = `${connectionId}:${model}`;
-  const expiry = modelLocks.get(key);
-  if (!expiry) return false;
-  if (expiry > Date.now()) return true;
-  modelLocks.delete(key); // clean up expired
-  return false;
-}
-
-/**
- * Lock a model on a specific connection for `durationMs` milliseconds.
- */
-function lockModel(connectionId, model, durationMs) {
-  if (!connectionId || !model) return;
-  const key = `${connectionId}:${model}`;
-  modelLocks.set(key, Date.now() + durationMs);
-  log.warn("AUTH", `Model lock: ${model} on ${connectionId.slice(0, 8)} for ${Math.round(durationMs / 1000)}s`);
-}
-
-/**
- * Check whether a provider uses per-model quota buckets.
- */
-function isMultiBucketProvider(provider) {
-  return MULTI_BUCKET_PROVIDERS.has(provider);
-}
-
 /**
  * Get provider credentials from localDb
  * Filters out unavailable accounts and returns the selected account based on strategy
@@ -85,50 +29,35 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul
     log.debug("AUTH", `${provider} | total connections: ${connections.length}, excludeId: ${excludeConnectionId || "none"}, model: ${model || "any"}`);
 
     if (connections.length === 0) {
-      // Check all connections (including inactive) to see if rate limited
-      const allConnections = await getProviderConnections({ provider: providerId });
-      log.debug("AUTH", `${provider} | all connections (incl inactive): ${allConnections.length}`);
-      if (allConnections.length > 0) {
-        const earliest = getEarliestRateLimitedUntil(allConnections);
-        if (earliest) {
-          log.warn("AUTH", `${provider} | all ${allConnections.length} accounts rate limited (${formatRetryAfter(earliest)})`);
-          return { allRateLimited: true, retryAfter: earliest, retryAfterHuman: formatRetryAfter(earliest) };
-        }
-        log.warn("AUTH", `${provider} | ${allConnections.length} accounts found but none active`);
-        allConnections.forEach(c => {
-          log.debug("AUTH", `  → ${c.id?.slice(0, 8)} | isActive=${c.isActive} | rateLimitedUntil=${c.rateLimitedUntil || "none"} | testStatus=${c.testStatus}`);
-        });
-      }
       log.warn("AUTH", `No credentials for ${provider}`);
       return null;
     }
 
-    // Filter out unavailable accounts and excluded connection
-    const multiBucket = isMultiBucketProvider(provider);
+    // Filter out model-locked and excluded connections
     const availableConnections = connections.filter(c => {
       if (excludeConnectionId && c.id === excludeConnectionId) return false;
-      if (isAccountUnavailable(c.rateLimitedUntil)) return false;
-      if (multiBucket && model && isModelLocked(c.id, model)) return false;
+      if (isModelLockActive(c, model)) return false;
       return true;
     });
 
     log.debug("AUTH", `${provider} | available: ${availableConnections.length}/${connections.length}`);
     connections.forEach(c => {
       const excluded = excludeConnectionId && c.id === excludeConnectionId;
-      const rateLimited = isAccountUnavailable(c.rateLimitedUntil);
-      const modelLocked = multiBucket && model && isModelLocked(c.id, model);
-      if (excluded || rateLimited || modelLocked) {
-        log.debug("AUTH", `  → ${c.id?.slice(0, 8)} | ${excluded ? "excluded" : ""} ${rateLimited ? `rateLimited until ${c.rateLimitedUntil}` : ""} ${modelLocked ? `modelLocked(${model})` : ""}`);
+      const locked = isModelLockActive(c, model);
+      if (excluded || locked) {
+        const lockUntil = getEarliestModelLockUntil(c);
+        log.debug("AUTH", `  → ${c.id?.slice(0, 8)} | ${excluded ? "excluded" : ""} ${locked ? `modelLocked(${model}) until ${lockUntil}` : ""}`);
       }
     });
 
     if (availableConnections.length === 0) {
-      const earliest = getEarliestRateLimitedUntil(connections);
+      // Find earliest lock expiry across all connections for retry timing
+      const lockedConns = connections.filter(c => isModelLockActive(c, model));
+      const expiries = lockedConns.map(c => getEarliestModelLockUntil(c)).filter(Boolean);
+      const earliest = expiries.sort()[0] || null;
       if (earliest) {
-        // Find the connection with the earliest rateLimitedUntil to get its error info
-        const rateLimitedConns = connections.filter(c => c.rateLimitedUntil && new Date(c.rateLimitedUntil).getTime() > Date.now());
-        const earliestConn = rateLimitedConns.sort((a, b) => new Date(a.rateLimitedUntil) - new Date(b.rateLimitedUntil))[0];
-        log.warn("AUTH", `${provider} | all ${connections.length} active accounts rate limited (${formatRetryAfter(earliest)}) | lastErrorCode=${earliestConn?.errorCode}, lastError=${earliestConn?.lastError?.slice(0, 50)}`);
+        const earliestConn = lockedConns[0];
+        log.warn("AUTH", `${provider} | all ${connections.length} accounts locked for ${model || "all"} (${formatRetryAfter(earliest)}) | lastError=${earliestConn?.lastError?.slice(0, 50)}`);
         return {
           allRateLimited: true,
           retryAfter: earliest,
@@ -137,15 +66,6 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul
           lastErrorCode: earliestConn?.errorCode || null
         };
       }
-      if (multiBucket && model) {
-        log.warn("AUTH", `${provider} | all accounts model-locked for ${model}`);
-        return {
-          allRateLimited: true,
-          retryAfter: new Date(Date.now() + 60000).toISOString(),
-          retryAfterHuman: "reset after 1m",
-          lastError: `All accounts rate limited for model ${model}`
-        };
-      }
       log.warn("AUTH", `${provider} | all ${connections.length} accounts unavailable`);
       return null;
     }
@@ -209,7 +129,8 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul
       // Include current status for optimization check
       testStatus: connection.testStatus,
       lastError: connection.lastError,
-      rateLimitedUntil: connection.rateLimitedUntil
+      // Pass full connection for clearAccountError to read modelLock_* keys
+      _connection: connection
     };
   } finally {
     if (resolveMutex) resolveMutex();
@@ -217,9 +138,8 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul
 }
 
 /**
- * Mark account as unavailable — reads backoffLevel from DB, calculates cooldown with exponential backoff, saves new level.
- * For multi-bucket providers (e.g. Antigravity), 429 errors lock only the specific model in memory
- * rather than the entire account in the database.
+ * Mark account+model as unavailable — locks modelLock_${model} in DB.
+ * All errors (429, 401, 5xx, etc.) lock per model, not per account.
  * @param {string} connectionId
  * @param {number} status - HTTP status code from upstream
  * @param {string} errorText
@@ -228,7 +148,6 @@ export async function getProviderCredentials(provider, excludeConnectionId = nul
  * @returns {{ shouldFallback: boolean, cooldownMs: number }}
  */
 export async function markAccountUnavailable(connectionId, status, errorText, provider = null, model = null) {
-  // Read current connection to get backoffLevel
   const connections = await getProviderConnections({ provider });
   const conn = connections.find(c => c.id === connectionId);
   const backoffLevel = conn?.backoffLevel || 0;
@@ -236,17 +155,11 @@ export async function markAccountUnavailable(connectionId, status, errorText, pr
   const { shouldFallback, cooldownMs, newBackoffLevel } = checkFallbackError(status, errorText, backoffLevel);
   if (!shouldFallback) return { shouldFallback: false, cooldownMs: 0 };
 
-  if (isMultiBucketProvider(provider) && status === 429 && model) {
-    const lockDuration = cooldownMs > 0 ? cooldownMs : DEFAULT_MODEL_LOCK_MS;
-    lockModel(connectionId, model, lockDuration);
-    return { shouldFallback: true, cooldownMs: 0 };
-  }
-
-  const rateLimitedUntil = getUnavailableUntil(cooldownMs);
   const reason = typeof errorText === "string" ? errorText.slice(0, 100) : "Provider error";
+  const lockUpdate = buildModelLockUpdate(model, cooldownMs);
 
   await updateProviderConnection(connectionId, {
-    rateLimitedUntil,
+    ...lockUpdate,
     testStatus: "unavailable",
     lastError: reason,
     errorCode: status,
@@ -254,6 +167,9 @@ export async function markAccountUnavailable(connectionId, status, errorText, pr
     backoffLevel: newBackoffLevel ?? backoffLevel
   });
 
+  const lockKey = Object.keys(lockUpdate)[0];
+  log.warn("AUTH", `${connectionId.slice(0, 8)} locked ${lockKey} for ${Math.round(cooldownMs / 1000)}s [${status}]`);
+
   if (provider && status && reason) {
     console.error(`❌ ${provider} [${status}]: ${reason}`);
   }
@@ -263,21 +179,26 @@ export async function markAccountUnavailable(connectionId, status, errorText, pr
 
 /**
  * Clear account error status (only if currently has error)
- * Optimized to avoid unnecessary DB updates
+ * Clears all modelLock_* fields and resets error state.
  */
 export async function clearAccountError(connectionId, currentConnection) {
-  // Only update if currently has error status
-  const hasError = currentConnection.testStatus === "unavailable" ||
-    currentConnection.lastError ||
-    currentConnection.rateLimitedUntil;
+  // Support both direct connection object and credentials wrapper
+  const conn = currentConnection._connection || currentConnection;
+  const now = Date.now();
+
+  // Collect all modelLock_* keys (both active and expired)
+  const allLockKeys = Object.keys(conn).filter(k => k.startsWith("modelLock_"));
+  const hasError = conn.testStatus === "unavailable" || conn.lastError || allLockKeys.length > 0;
 
   if (!hasError) return; // Skip if already clean
 
+  // Clear all modelLock_* keys (lazy cleanup of expired ones included)
+  const clearLocks = Object.fromEntries(allLockKeys.map(k => [k, null]));
   await updateProviderConnection(connectionId, {
+    ...clearLocks,
     testStatus: "active",
     lastError: null,
     lastErrorAt: null,
-    rateLimitedUntil: null,
     backoffLevel: 0
   });
   log.info("AUTH", `Account ${connectionId.slice(0, 8)} error cleared`);