nolo-cli 0.1.26 → 0.1.29

package/server/handlers/chatHandler.ts

@@ -0,0 +1,797 @@
+
+// 文件路径: packages/server/handlers/chatHandler.ts
+
+import { getNoloKey } from "../../ai/llm/getNoloKey";
+import { getUser } from "../../auth/server/getUser";
+import { authenticateRequest } from "../../auth/utils";
+import { getModelPricing, getPrices, getFinalPrice } from "../../ai/llm/getPricing";
+import { checkBalanceForMessages } from "../../auth/billing";
+import { fetchUpstreamWithGentleRetry } from "./chatUpstreamRetry";
+import {
+  createBillingAwareSseRewriter,
+  formatBillingUsageEvent,
+} from "./chatBillingSse";
+import {
+  fetchGoogleNativeImageChat,
+  isGoogleNativeImageModel,
+} from "./googleNativeChat";
+import {
+  GOOGLE_FALLBACK_STATUSES,
+  type ChatProxyPayload,
+  buildChatFallbackPlan,
+  buildChatRoutingDecisionLog,
+  normalizeChatProxyPayload,
+  resolveProviderCredentials,
+  selectUpstreamRequest,
+} from "./chatProxyRouting";
+import { authenticateMachineToken } from "./machines/machineAuth";
+import serverDb from "../../database/server/db";
+const encoder = new TextEncoder();
+
+// 简单的控制台日志封装，带颜色高亮
+const log = {
+  info: (...args: any[]) =>
+    console.log("\x1b[36m%s\x1b[0m", "[INFO]", ...args), // 青色
+  warn: (...args: any[]) =>
+    console.warn("\x1b[33m%s\x1b[0m", "[WARN]", ...args), // 黄色
+  error: (...args: any[]) =>
+    console.error("\x1b[31m%s\x1b[0m", "[ERROR]", ...args), // 红色
+};
+
+const getNowMs = () =>
+  typeof performance !== "undefined" ? performance.now() : Date.now();
+
+const logChatStreamPerf = (
+  startedAtMs: number,
+  stage: string,
+  details: Record<string, unknown> = {}
+) => {
+  log.info("Chat stream perf", {
+    stage,
+    elapsedMs: Math.round((getNowMs() - startedAtMs) * 10) / 10,
+    ...details,
+  });
+};
+
+const CORS = {
+  "Access-Control-Allow-Origin": "*",
+  "Access-Control-Allow-Methods": "POST, OPTIONS",
+  "Access-Control-Allow-Headers": "Content-Type, Authorization",
+};
+
+const formatSseError = (msg: string, code: string) =>
+  encoder.encode(
+    `data: ${JSON.stringify({ error: { msg, code } })} \n\n`
+  );
+
+const CHAT_PROXY_WARMUP_CHUNK = encoder.encode(": nolo-chat-proxy-ready\n\n");
+
+const createIdleTimeoutGuard = (
+  timeoutMs: number
+): {
+  arm: (ctrl: TransformStreamDefaultController<Uint8Array>) => void;
+  clear: () => void;
+  close: () => void;
+} => {
+  let idleTimer: ReturnType<typeof setTimeout> | null = null;
+  let closed = false;
+
+  const clear = () => {
+    if (idleTimer) {
+      clearTimeout(idleTimer);
+      idleTimer = null;
+    }
+  };
+
+  const close = () => {
+    closed = true;
+    clear();
+  };
+
+  const arm = (ctrl: TransformStreamDefaultController<Uint8Array>) => {
+    clear();
+    if (closed) return;
+    idleTimer = setTimeout(() => {
+      if (closed) return;
+      try {
+        ctrl.enqueue(formatSseError(`idle ${timeoutMs / 1000} s`, "IDLE"));
+      } catch {
+        close();
+      }
+    }, timeoutMs);
+  };
+
+  return { arm, clear, close };
+};
+
+// 计算消息预估成本的辅助函数（用于日志记录）
+const calculateMessageCost = (
+  model: string,
+  provider: string,
+  messages: any[],
+  requestedServiceTier?: string
+): number => {
+  try {
+    const serverPrices = getModelPricing(provider, model, requestedServiceTier);
+    if (!serverPrices) {
+      log.info("无法获取模型定价信息");
+      return 0;
+    }
+
+    const userConfig = {
+      inputPrice: 0,
+      outputPrice: 0,
+    };
+
+    const prices = getPrices(userConfig, serverPrices);
+    const finalPrice = getFinalPrice(prices);
+
+    log.info(`模型 ${model} 预估成本计算`, {
+      serverPrices,
+      userConfig,
+      prices,
+      finalPrice,
+      model,
+      provider,
+    });
+
+    return finalPrice;
+  } catch (error) {
+    log.error("计算消息成本时出错:", error);
+    return 0;
+  }
+};
+
+// 预估成本检查：超过余额 50% 时仅做日志警告
+const checkEstimatedCost = (
+  balance: number,
+  estimatedCost: number,
+  imageUrlCount: number
+) => {
+  if (estimatedCost > 0 && estimatedCost > balance * 0.5) {
+    log.warn(
+      `警告：预估成本 $${estimatedCost.toFixed(
+        6
+      )
+      } 超过余额的50 % (${(balance * 0.5).toFixed(2)})`,
+      {
+        estimatedCost,
+        balance,
+        threshold: balance * 0.5,
+        imageUrlCount,
+      }
+    );
+  }
+  return { allowed: true as const };
+};
+
+interface RequestCheckResult {
+  allowed: boolean;
+  error?: {
+    message: string;
+    code: string;
+    status: number;
+    details?: any;
+  };
+  imageUrlCount?: number;
+  estimatedCost?: number;
+}
+
+// 使用 auth/billing 中的通用逻辑，再附加一次成本预估日志
+const performRequestChecks = async (
+  userId: string,
+  balance: number,
+  model: string,
+  provider: string,
+  messages: any[],
+  serviceTier: string | undefined,
+  _auth: any
+): Promise<RequestCheckResult> => {
+  try {
+    const {
+      allowed,
+      error,
+      imageUrlCount,
+    } = checkBalanceForMessages({ balance, messages });
+
+    log.info("图片URL数量统计", { imageUrlCount });
+
+    const estimatedCost = calculateMessageCost(
+      model,
+      provider,
+      messages || [],
+      serviceTier
+    );
+    log.info(`预估消息成本: $${estimatedCost.toFixed(6)} `, {
+      model,
+      provider,
+      estimatedCost,
+      userId,
+    });
+
+    if (!allowed && error) {
+      return {
+        allowed: false,
+        error,
+        imageUrlCount,
+        estimatedCost,
+      };
+    }
+
+    checkEstimatedCost(balance, estimatedCost, imageUrlCount);
+
+    return {
+      allowed: true,
+      imageUrlCount,
+      estimatedCost,
+    };
+  } catch (error) {
+    log.error("请求检查过程中发生错误:", error);
+    return {
+      allowed: false,
+      error: {
+        message: "请求检查失败",
+        code: "CHECK_FAILED",
+        status: 500,
+        details: {
+          error: error instanceof Error ? error.message : String(error),
+        },
+      },
+    };
+  }
+};
+
+const normalizeCustomChatUrl = (rawUrl: unknown): string | undefined => {
+  if (typeof rawUrl !== "string" || !rawUrl.trim()) return undefined;
+  const url = rawUrl.trim().replace(/\/$/, "");
+  if (/\/v\d+(beta\d*)?$/.test(url)) {
+    return `${url}/chat/completions`;
+  }
+  return url;
+};
+
+const hydrateProxyRequestFromAgent = async (
+  request: ChatProxyPayload
+): Promise<ChatProxyPayload> => {
+  const agentKey =
+    typeof request.agentKey === "string" && request.agentKey.trim()
+      ? request.agentKey.trim()
+      : "";
+  if (!agentKey) return request;
+  if (request.url && request.KEY) return request;
+
+  let agent: any = null;
+  try {
+    agent = await serverDb.get(agentKey);
+  } catch {
+    return request;
+  }
+  if (!agent || typeof agent !== "object") return request;
+
+  const payload = {
+    ...request.body,
+    url: request.url || normalizeCustomChatUrl(agent.customProviderUrl),
+    KEY: request.KEY || agent.apiKey,
+    agentHydratedKey: !request.KEY && Boolean(agent.apiKey),
+    agentKey,
+    provider: request.provider ?? agent.provider ?? agent.apiSource,
+    apiSource: request.apiSource ?? agent.apiSource,
+    apiKeyHeader: request.apiKeyHeader ?? agent.apiKeyHeader,
+    model: request.body?.model ?? agent.model,
+  };
+
+  return normalizeChatProxyPayload(payload);
+};
+
+const authenticateChatProxyRequest = async (req: Request) => {
+  const machineAuth = await authenticateMachineToken(req);
+  if (machineAuth instanceof Response) return machineAuth;
+  if (machineAuth?.userId) return machineAuth;
+
+  return authenticateRequest(req);
+};
+
+export async function handleChatRequest(req: Request, extraHeaders = {}) {
+  const TIMEOUT = 300_000; // 全部都用 300s
+  const chatStreamStartedAtMs = getNowMs();
+
+  const auth = await authenticateChatProxyRequest(req);
+  if (auth instanceof Response) return auth;
+
+  const { userId } = auth;
+  const { balance = 0 } = await getUser(userId);
+  log.info("用户余额检查", { userId, balance });
+
+  try {
+    let request = normalizeChatProxyPayload(await req.json());
+    request = await hydrateProxyRequestFromAgent(request);
+    const { body, providerName, model, messages } = request;
+
+    log.info("处理聊天请求", {
+      url: request.url,
+      hasKey: !!request.KEY,
+      apiSource: request.apiSource,
+      userId,
+      bodyKeys: Object.keys(body),
+    });
+
+    log.info("=== Messages Debug Info ===", {
+      type: typeof messages,
+      length: Array.isArray(messages) ? messages.length : "Not an array",
+      content: JSON.stringify(messages, null, 2),
+    });
+    log.info("=== End Messages Debug ===");
+
+    const checkResult = await performRequestChecks(
+      userId,
+      balance,
+      model,
+      request.provider,
+      messages,
+      typeof body.service_tier === "string" ? body.service_tier : undefined,
+      auth
+    );
+
+    if (!checkResult.allowed && checkResult.error) {
+      log.warn("请求检查失败", {
+        userId,
+        balance,
+        errorCode: checkResult.error.code,
+        errorMessage: checkResult.error.message,
+      });
+
+      return new Response(
+        JSON.stringify({
+          error: {
+            message: checkResult.error.message,
+            code: checkResult.error.code,
+            details: checkResult.error.details,
+          },
+        }),
+        {
+          status: checkResult.error.status,
+          headers: {
+            "Content-Type": "application/json",
+            ...CORS,
+            ...extraHeaders,
+          },
+        }
+      );
+    }
+
+    log.info("请求检查通过", {
+      userId,
+      imageUrlCount: checkResult.imageUrlCount,
+      estimatedCost: checkResult.estimatedCost,
+    });
+
+    const credentials = resolveProviderCredentials(request, getNoloKey);
+
+    if (!credentials.apiKey && !credentials.isCustomApi) {
+      log.warn("API密钥缺失", {
+        provider: request.provider,
+        apiSource: request.apiSource,
+        userId,
+      });
+      const providerLabel = request.model || request.provider || "当前模型服务";
+      return new Response(
+        JSON.stringify({
+          error: {
+            message: `缺少模型服务 API Key：${providerLabel} 没有可用密钥。当前还没开始调用任何工具。请在模型设置里选择已配置的自定义 API，或为该 provider 配置平台密钥。`,
+            code: "MISSING_PROVIDER_API_KEY",
+          },
+        }),
+        {
+          status: 400,
+          headers: {
+            "Content-Type": "application/json",
+            ...CORS,
+            ...extraHeaders,
+          },
+        }
+      );
+    }
+
+    const controller = new AbortController();
+    const timer = setTimeout(() => controller.abort("timeout"), TIMEOUT);
+    const fallbackPlan = buildChatFallbackPlan(
+      request,
+      credentials,
+      controller.signal,
+      getNoloKey
+    );
+    const upstreamRequest = selectUpstreamRequest(
+      request,
+      credentials,
+      fallbackPlan,
+      controller.signal
+    );
+
+    log.info("Chat proxy routing decision", {
+      userId,
+      ...buildChatRoutingDecisionLog(
+        request,
+        credentials,
+        fallbackPlan,
+        upstreamRequest
+      ),
+    });
+
+    if (fallbackPlan.missingKimiFallbackKey) {
+      log.warn("Kimi 平台 fallback 未启用：缺少 DEEPINFRA_API_KEY 和 OPENROUTER_API_KEY", {
+        provider: providerName,
+        model: request.originalModel,
+        userId,
+      });
+    }
+
+    if (fallbackPlan.missingGoogleFallbackKey) {
+      log.warn("Google fallback 未启用：缺少 OPENROUTER_API_KEY", {
+        provider: providerName,
+        model,
+        userId,
+      });
+    }
+
+    if (fallbackPlan.routeKimiThroughOpenRouter) {
+      log.warn("Fireworks Kimi 平台请求直接切换到 OpenRouter fallback", {
+        provider: providerName,
+        model,
+        fallbackProvider: upstreamRequest.provider,
+        fallbackModel: upstreamRequest.model,
+        userId,
+      });
+    }
+
+    const shouldStream = body.stream !== false;
+    const streamHeaders = {
+      "Content-Type": "text/event-stream",
+      "Cache-Control": "no-cache, no-transform",
+      Connection: "keep-alive",
+      "X-Accel-Buffering": "no",
+      ...CORS,
+      ...extraHeaders,
+    };
+
+    if (providerName === "google" && isGoogleNativeImageModel(model)) {
+      try {
+        const nativeResponse = await fetchGoogleNativeImageChat({
+          model,
+          body,
+          apiKey: credentials.apiKey,
+          signal: controller.signal,
+          streamHeaders,
+        });
+        clearTimeout(timer);
+        return nativeResponse;
+      } catch (error: any) {
+        const shouldFallbackToOpenRouter =
+          !!fallbackPlan.googleFallbackRequest &&
+          (!("status" in (error ?? {})) ||
+            GOOGLE_FALLBACK_STATUSES.includes(Number(error?.status)));
+
+        if (!shouldFallbackToOpenRouter) {
+          clearTimeout(timer);
+          throw error;
+        }
+
+        log.warn("Google native image 上游失败，切换到 OpenRouter fallback", {
+          provider: providerName,
+          model,
+          status: error?.status,
+          error: error?.message,
+        });
+      }
+    }
+
+    log.info("发送上游请求", {
+      url: upstreamRequest.url,
+      provider: upstreamRequest.provider,
+      model: upstreamRequest.model,
+      timeout: TIMEOUT,
+    });
+    const {
+      response: upstream,
+      provider: actualBillingProvider = upstreamRequest.provider,
+      model: actualBillingModel = upstreamRequest.model,
+      serviceTier: returnedBillingServiceTier,
+    } = await fetchUpstreamWithGentleRetry({
+      url: upstreamRequest.url,
+      provider: upstreamRequest.provider,
+      model: upstreamRequest.model,
+      log,
+      fallbackRequest: upstreamRequest.fallbackRequest,
+      fallbackRequests: upstreamRequest.fallbackRequests,
+      fallbackOnStatuses: upstreamRequest.fallbackOnStatuses,
+      init: upstreamRequest.init,
+    });
+    clearTimeout(timer);
+    const requestedBillingServiceTier =
+      typeof body.service_tier === "string" ? body.service_tier : undefined;
+    const actualBillingServiceTier =
+      actualBillingProvider === providerName && actualBillingModel === model
+        ? requestedBillingServiceTier
+        : returnedBillingServiceTier;
+    const shouldAppendBillingUsage =
+      actualBillingProvider !== providerName ||
+      actualBillingModel !== model ||
+      actualBillingServiceTier !== requestedBillingServiceTier;
+    const billingUsageEvent = shouldAppendBillingUsage
+      ? formatBillingUsageEvent(
+        actualBillingProvider,
+        actualBillingModel,
+        actualBillingServiceTier
+      )
+      : null;
+
+    logChatStreamPerf(chatStreamStartedAtMs, "chat-proxy-upstream-response", {
+      userId,
+      requestedProvider: providerName,
+      requestedModel: model,
+      provider: upstreamRequest.provider,
+      upstreamModel: upstreamRequest.model,
+      actualBillingProvider,
+      actualBillingModel,
+      hasBillingUsageEvent: Boolean(billingUsageEvent),
+      status: upstream.status,
+    });
+
+    if (!upstream.ok) {
+      const t = await upstream.text();
+      const m = (() => {
+        try {
+          return JSON.parse(t)?.error?.message || t;
+        } catch {
+          return t;
+        }
+      })();
+
+      log.error("上游请求失败", {
+        status: upstream.status,
+        statusText: upstream.statusText,
+        response: t,
+      });
+
+      if (!shouldStream) {
+        return new Response(
+          JSON.stringify({
+            error: {
+              message: m,
+              code: `UPSTREAM_${upstream.status}`,
+            },
+          }),
+          {
+            status: upstream.status,
+            headers: {
+              "Content-Type": "application/json",
+              ...CORS,
+              ...extraHeaders,
+            },
+          }
+        );
+      }
+
+      return new Response(
+        new ReadableStream({
+          start(ctrl) {
+            ctrl.enqueue(formatSseError(m, `UPSTREAM_${upstream.status} `));
+            ctrl.close();
+          },
+        }),
+        { status: 200, headers: streamHeaders }
+      );
+    }
+
+    if (!shouldStream) {
+      return new Response(upstream.body, {
+        status: upstream.status,
+        headers: {
+          "Content-Type": upstream.headers.get("Content-Type") || "application/json",
+          ...CORS,
+          ...extraHeaders,
+        },
+      });
+    }
+
+    if (!upstream.body) {
+      log.error("上游响应 body 为空", {
+        url: upstreamRequest.url,
+        status: upstream.status,
+      });
+      return new Response(
+        JSON.stringify({ error: { message: "Empty upstream response", code: "UPSTREAM_EMPTY" } }),
+        { status: 502, headers: { "Content-Type": "application/json", ...CORS, ...extraHeaders } }
+      );
+    }
+
+    const stream = (() => {
+      if (!billingUsageEvent) {
+        const idleGuard = createIdleTimeoutGuard(TIMEOUT);
+        let loggedFirstUpstreamChunk = false;
+        let loggedFirstProxyEnqueue = false;
+        return upstream.body.pipeThrough(
+          new TransformStream({
+            start(ctrl) {
+              ctrl.enqueue(CHAT_PROXY_WARMUP_CHUNK);
+              logChatStreamPerf(
+                chatStreamStartedAtMs,
+                "chat-proxy-warmup-enqueue",
+                {
+                  userId,
+                  provider: actualBillingProvider,
+                  model: actualBillingModel,
+                  byteLength: CHAT_PROXY_WARMUP_CHUNK.byteLength,
+                  hasBillingUsageEvent: false,
+                }
+              );
+              // 初始 idle timer：上游迟迟不发首个 chunk 时超时
+              idleGuard.arm(ctrl);
+            },
+            transform(chunk, ctrl) {
+              idleGuard.clear();
+              if (!loggedFirstUpstreamChunk) {
+                loggedFirstUpstreamChunk = true;
+                logChatStreamPerf(
+                  chatStreamStartedAtMs,
+                  "chat-proxy-upstream-first-chunk",
+                  {
+                    userId,
+                    provider: actualBillingProvider,
+                    model: actualBillingModel,
+                    byteLength: chunk.byteLength,
+                    hasBillingUsageEvent: false,
+                  }
+                );
+              }
+              ctrl.enqueue(chunk);
+              if (!loggedFirstProxyEnqueue) {
+                loggedFirstProxyEnqueue = true;
+                logChatStreamPerf(
+                  chatStreamStartedAtMs,
+                  "chat-proxy-first-enqueue",
+                  {
+                    userId,
+                    provider: actualBillingProvider,
+                    model: actualBillingModel,
+                    byteLength: chunk.byteLength,
+                    hasBillingUsageEvent: false,
+                  }
+                );
+              }
+              idleGuard.arm(ctrl);
+            },
+            flush() {
+              idleGuard.close();
+            },
+            cancel() {
+              idleGuard.close();
+            },
+          })
+        );
+      }
+
+      const rewriter = createBillingAwareSseRewriter(billingUsageEvent);
+      const idleGuard = createIdleTimeoutGuard(TIMEOUT);
+      let loggedFirstUpstreamChunk = false;
+      let loggedFirstProxyEnqueue = false;
+
+      return upstream.body.pipeThrough(
+        new TransformStream({
+          start(ctrl) {
+            ctrl.enqueue(CHAT_PROXY_WARMUP_CHUNK);
+            logChatStreamPerf(
+              chatStreamStartedAtMs,
+              "chat-proxy-warmup-enqueue",
+              {
+                userId,
+                provider: actualBillingProvider,
+                model: actualBillingModel,
+                byteLength: CHAT_PROXY_WARMUP_CHUNK.byteLength,
+                hasBillingUsageEvent: true,
+              }
+            );
+            idleGuard.arm(ctrl);
+          },
+          transform(chunk, ctrl) {
+            idleGuard.clear();
+            if (!loggedFirstUpstreamChunk) {
+              loggedFirstUpstreamChunk = true;
+              logChatStreamPerf(
+                chatStreamStartedAtMs,
+                "chat-proxy-upstream-first-chunk",
+                {
+                  userId,
+                  provider: actualBillingProvider,
+                  model: actualBillingModel,
+                  byteLength: chunk.byteLength,
+                  hasBillingUsageEvent: true,
+                }
+              );
+            }
+            for (const rewrittenChunk of rewriter.transform(chunk)) {
+              ctrl.enqueue(rewrittenChunk);
+              if (!loggedFirstProxyEnqueue) {
+                loggedFirstProxyEnqueue = true;
+                logChatStreamPerf(
+                  chatStreamStartedAtMs,
+                  "chat-proxy-first-enqueue",
+                  {
+                    userId,
+                    provider: actualBillingProvider,
+                    model: actualBillingModel,
+                    upstreamByteLength: chunk.byteLength,
+                    rewrittenByteLength: rewrittenChunk.byteLength,
+                    hasBillingUsageEvent: true,
+                  }
+                );
+              }
+            }
+            idleGuard.arm(ctrl);
+          },
+          flush(ctrl) {
+            for (const rewrittenChunk of rewriter.flush()) {
+              ctrl.enqueue(rewrittenChunk);
+              if (!loggedFirstProxyEnqueue) {
+                loggedFirstProxyEnqueue = true;
+                logChatStreamPerf(
+                  chatStreamStartedAtMs,
+                  "chat-proxy-first-enqueue",
+                  {
+                    userId,
+                    provider: actualBillingProvider,
+                    model: actualBillingModel,
+                    rewrittenByteLength: rewrittenChunk.byteLength,
+                    hasBillingUsageEvent: true,
+                    source: "flush",
+                  }
+                );
+              }
+            }
+            idleGuard.close();
+          },
+          cancel() {
+            idleGuard.close();
+          },
+        })
+      );
+    })();
+
+    log.info("流式响应开始", {
+      userId,
+      model,
+      provider: request.provider,
+      actualBillingProvider,
+      actualBillingModel,
+    });
+    return new Response(stream, { status: 200, headers: streamHeaders });
+  } catch (e: any) {
+    const isAbort = e?.name === "AbortError";
+
+    if (isAbort) {
+      log.error("请求超时", { error: e.message, timeout: TIMEOUT, userId });
+    } else {
+      log.error("代理失败", {
+        error: e.message,
+        stack: e.stack,
+        userId,
+      });
+    }
+
+    return new Response(
+      JSON.stringify({
+        error: {
+          message: isAbort ? e?.message : "Proxy failed",
+          code: "PROXY_FAILED",
+        },
+      }),
+      {
+        status: isAbort ? 504 : 500,
+        headers: {
+          "Content-Type": "application/json",
+          ...CORS,
+          ...extraHeaders,
+        },
+      }
+    );
+  }
+}