ai-chatbot/lib/ai/runtime/agent-model-failover.ts

import { is3DModelDescriptor, isImageGenerationModelDescriptor, isVideoModelDescriptor } from "../model-classification";
import { isMediaInputFailure } from "./request-failure-classification";
import type { getEnabledAiModelsWithProvider } from "@/lib/db/queries";

export type AgentModelEntry = Awaited<ReturnType<typeof getEnabledAiModelsWithProvider>>[number];
export const AGENT_MODEL_COOLDOWN_MS = 30 * 60_000;
export const AGENT_MODEL_ATTEMPT_MS = 45_000;
export type AgentModelFault = { reason: string; status: number | null };
export interface AgentCircuitStore {
  isPaused(entry: AgentModelEntry): Promise<boolean>;
  fail(entry: AgentModelEntry, fault: AgentModelFault): Promise<void>;
}

export class AgentModelProtocolError extends Error {
  constructor() { super("Agent model returned an invalid response"); this.name = "AgentModelProtocolError"; }
}

const COMPATIBLE = new Set(["vip98k", "openai", "openai-compatible", "alibaba", "qwen", "local-qwen", "local", "zhipu", "moonshot", "deepseek", "xai", "siliconflow", "custom"]);
export function isAgentMultimodalCandidate(entry: AgentModelEntry) {
  const { model, provider } = entry;
  const descriptor = { ...model, providerType: provider.type };
  const caps = model.capabilities as Record<string, unknown> | null;
  return model.isEnabled && provider.isEnabled && model.supportsVision &&
    model.healthStatus !== "auto_disabled" && COMPATIBLE.has(provider.type) &&
    !["image", "image-tool", "image-edit", "audio", "speech", "embedding", "rerank"].includes(String(caps?.kind || "")) &&
    caps?.supportsTools !== false && caps?.supportsToolCalling !== false && caps?.supportsToolCalls !== false && caps?.supportsFunctionCalling !== false &&
    !isImageGenerationModelDescriptor(descriptor) && !isVideoModelDescriptor(descriptor) && !is3DModelDescriptor(descriptor);
}

export function agentModelCandidates(primary: AgentModelEntry | undefined, entries: AgentModelEntry[]) {
  const eligible = entries.filter(isAgentMultimodalCandidate);
  const ordered = [primary?.model.id, primary?.model.fallbackModelId1, primary?.model.fallbackModelId2];
  return [...eligible].filter(entry => entry.model.id === primary?.model.id || !entry.model.excludeFromFailover)
    .sort((a, b) => {
      const rank = (entry: AgentModelEntry) => { const index = ordered.indexOf(entry.model.id); return index < 0 ? 3 : index; };
      return rank(a) - rank(b) || Number(b.model.isDefault) - Number(a.model.isDefault) ||
        (a.model.sortOrder ?? 0) - (b.model.sortOrder ?? 0) || a.model.id.localeCompare(b.model.id);
    });
}

export function agentModelFault(error: unknown, signal?: AbortSignal): AgentModelFault | null {
  if (signal?.aborted) return null;
  const e = error as { status?: number; statusCode?: number; message?: string; name?: string; cause?: unknown } | null;
  const status = Number(e?.statusCode || e?.status) || null;
  const message = String(e?.message || "").toLowerCase();
  if (isMediaInputFailure(message) || /content.?policy|content.?filter|moderation|safety|违规|敏感|user cancel|user abort/.test(message)) return null;
  if (/tools?.*(?:not supported|unsupported)|does not support.*(?:tools?|function.call)/.test(message)) return { reason: "tools_unsupported", status };
  if (status && [400, 413, 415, 422].includes(status)) return null;
  if (status === 401 || /invalid authentication|invalid.api.key/.test(message)) return { reason: "authentication", status };
  if (status === 403) return { reason: "access_denied", status };
  if (status === 429 || status === 402) return { reason: "rate_or_quota", status };
  if (status === 404 || /model.not.found|invalid.model/.test(message)) return { reason: "model_unavailable", status };
  if (status === 408) return { reason: "connection_or_timeout", status };
  if (status && status >= 500) return { reason: "upstream_unavailable", status };
  if (["AgentModelProtocolError", "AI_NoOutputGeneratedError", "AI_InvalidResponseDataError", "AI_TypeValidationError", "AI_JSONParseError"].includes(String(e?.name)) ||
    /未返回可用的回答|未返回所要求的工具|empty.response|no.output.generated/.test(message)) return { reason: "invalid_response", status };
  if (/timeout|timed out|fetch failed|econnreset|econnrefused|network|socket|terminated/.test(message) || e?.name === "TimeoutError") return { reason: "connection_or_timeout", status };
  return e?.cause ? agentModelFault(e.cause, signal) : null;
}

// Retries only the model computation. Callers must never include business tools
// in execute(): native tool calls are exposed only after a complete valid reply.
export function createAgentModelFailover(entries: AgentModelEntry[], circuit: AgentCircuitStore) {
  let preferred: string | undefined;
  return async function run<T>(execute: (entry: AgentModelEntry, signal: AbortSignal) => Promise<T>, signal?: AbortSignal): Promise<T> {
    const deadline = AbortSignal.timeout(150_000);
    const totalSignal = signal ? AbortSignal.any([signal, deadline]) : deadline;
    const ordered = [...entries].sort((a, b) => Number(b.model.id === preferred) - Number(a.model.id === preferred));
    let last: unknown;
    for (const entry of ordered) {
      totalSignal.throwIfAborted();
      if (await circuit.isPaused(entry)) continue;
      const attemptSignal = AbortSignal.any([totalSignal, AbortSignal.timeout(AGENT_MODEL_ATTEMPT_MS)]);
      try {
        const result = await execute(entry, attemptSignal);
        totalSignal.throwIfAborted();
        preferred = entry.model.id;
        // A late success must never erase a concurrent request's newer fault.
        return result;
      } catch (error) {
        const fault = agentModelFault(attemptSignal.aborted && !totalSignal.aborted
          ? Object.assign(new Error("Agent model timed out"), { name: "TimeoutError" }) : error, totalSignal);
        if (!fault) throw error;
        await circuit.fail(entry, fault);
        last = new Error(`Agent 模型通道暂不可用（${fault.reason}），故障通道已暂停 30 分钟`);
      }
    }
    throw last || new Error("Agent 多模态通道暂不可用，故障通道正在冷却，请稍后重试");
  };
}