plugin-ai-api 1.0.25 → 1.0.28

package/src/server/services/file-processor.ts

@@ -1,5 +1,8 @@
 import type { Context } from '@nocobase/actions';
+import dns from 'dns';
+import { isIP } from 'net';
 import { basename } from 'path';
+import { getAiApiConfig } from '../utils/request-cache';
 
 export interface FileContentBlock {
   type: string;
@@ -29,9 +32,11 @@ export interface PdfToImageRenderer {
 export type FileProcessorErrorCode =
   | 'invalid_url'
   | 'unsupported_protocol'
+  | 'blocked_host'
   | 'fetch_failed'
   | 'file_too_large'
   | 'content_type_not_allowed'
+  | 'too_many_redirects'
   | 'missing_url';
 
 export class FileProcessorError extends Error {
@@ -47,13 +52,16 @@ export class FileProcessorError extends Error {
 export interface FetchFileOptions {
   maxSizeBytes?: number;
   timeoutMs?: number;
+  maxRedirects?: number;
   allowedProtocols?: string[];
   allowedContentTypes?: string[];
 }
 
 const DEFAULT_MAX_FILE_SIZE = 50 * 1024 * 1024; // 50 MB
 const DEFAULT_TIMEOUT_MS = 30_000;
+const DEFAULT_MAX_REDIRECTS = 5;
 const ALLOWED_PROTOCOLS = new Set(['http:', 'https:']);
+const REDIRECT_STATUSES = new Set([301, 302, 303, 307, 308]);
 
 export class FileProcessorService {
   private processors: FileProcessor[] = [];
@@ -106,12 +114,6 @@ function isRecord(value: unknown): value is Record<string, unknown> {
   return typeof value === 'object' && value !== null && !Array.isArray(value);
 }
 
-function getUrlString(value: unknown): string | undefined {
-  if (typeof value === 'string') return value;
-  if (isRecord(value) && typeof value.url === 'string') return value.url;
-  return undefined;
-}
-
 function extractFilename(url: string, contentDisposition: string | null): string | undefined {
   if (contentDisposition) {
     const match = contentDisposition.match(/filename="?([^"]+)"?/);
@@ -126,44 +128,232 @@ function extractFilename(url: string, contentDisposition: string | null): string
   return undefined;
 }
 
+function parseIpv4Octets(ip: string): number[] | undefined {
+  const parts = ip.split('.');
+  if (parts.length !== 4) return undefined;
+  const octets: number[] = [];
+  for (const part of parts) {
+    if (!/^\d{1,3}$/.test(part)) return undefined;
+    const value = Number(part);
+    if (value > 255) return undefined;
+    octets.push(value);
+  }
+  return octets;
+}
+
+function isBlockedIpv4(ip: string): boolean {
+  const octets = parseIpv4Octets(ip);
+  if (!octets) return true;
+  const [a, b] = octets;
+  if (a === 0) return true; // 0.0.0.0/8 "this" network
+  if (a === 10) return true; // 10.0.0.0/8 private
+  if (a === 100 && b >= 64 && b <= 127) return true; // 100.64.0.0/10 carrier-grade NAT
+  if (a === 127) return true; // 127.0.0.0/8 loopback
+  if (a === 169 && b === 254) return true; // 169.254.0.0/16 link-local (cloud metadata)
+  if (a === 172 && b >= 16 && b <= 31) return true; // 172.16.0.0/12 private
+  if (a === 192 && b === 0) return true; // 192.0.0.0/24 IETF protocol assignments
+  if (a === 192 && b === 168) return true; // 192.168.0.0/16 private
+  if (a === 198 && (b === 18 || b === 19)) return true; // 198.18.0.0/15 benchmarking
+  if (a >= 224) return true; // multicast, reserved, broadcast
+  return false;
+}
+
+function parseHexGroup(group: string): number | undefined {
+  if (group.length < 1 || group.length > 4 || !/^[0-9a-fA-F]+$/.test(group)) return undefined;
+  return parseInt(group, 16);
+}
+
+function expandIpv6(input: string): number[] | undefined {
+  let address = input;
+  const zoneIndex = address.indexOf('%');
+  if (zoneIndex !== -1) address = address.slice(0, zoneIndex);
+  if (!address) return undefined;
+
+  // Rewrite an embedded IPv4 suffix (e.g. ::ffff:127.0.0.1) as two hex groups.
+  const lastColon = address.lastIndexOf(':');
+  if (lastColon !== -1 && address.includes('.', lastColon)) {
+    const octets = parseIpv4Octets(address.slice(lastColon + 1));
+    if (!octets) return undefined;
+    const high = ((octets[0] << 8) | octets[1]).toString(16);
+    const low = ((octets[2] << 8) | octets[3]).toString(16);
+    address = `${address.slice(0, lastColon + 1)}${high}:${low}`;
+  }
+
+  const groups: number[] = [];
+  const doubleColonIndex = address.indexOf('::');
+  if (doubleColonIndex !== -1) {
+    if (address.indexOf('::', doubleColonIndex + 1) !== -1) return undefined;
+    const head = address.slice(0, doubleColonIndex);
+    const tail = address.slice(doubleColonIndex + 2);
+    const headGroups = head === '' ? [] : head.split(':');
+    const tailGroups = tail === '' ? [] : tail.split(':');
+    const fillCount = 8 - headGroups.length - tailGroups.length;
+    if (fillCount < 1) return undefined;
+    const allGroups = [...headGroups];
+    for (let i = 0; i < fillCount; i += 1) allGroups.push('0');
+    allGroups.push(...tailGroups);
+    for (const group of allGroups) {
+      const parsed = parseHexGroup(group);
+      if (parsed === undefined) return undefined;
+      groups.push(parsed);
+    }
+  } else {
+    for (const group of address.split(':')) {
+      const parsed = parseHexGroup(group);
+      if (parsed === undefined) return undefined;
+      groups.push(parsed);
+    }
+  }
+
+  return groups.length === 8 ? groups : undefined;
+}
+
+function isBlockedIpv6(ip: string): boolean {
+  const groups = expandIpv6(ip);
+  if (!groups) return true;
+  const [g0, g1, g2, g3, g4, g5, g6, g7] = groups;
+
+  if (groups.every((group) => group === 0)) return true; // :: unspecified
+  if (g0 === 0 && g1 === 0 && g2 === 0 && g3 === 0 && g4 === 0 && g5 === 0 && g6 === 0 && g7 === 1) {
+    return true; // ::1 loopback
+  }
+
+  // Embedded IPv4 forms: IPv4-mapped (::ffff:a.b.c.d), IPv4-compatible (::a.b.c.d),
+  // and NAT64 (64:ff9b::a.b.c.d) — judge by the embedded IPv4 address.
+  const isIpv4Mapped = g0 === 0 && g1 === 0 && g2 === 0 && g3 === 0 && g4 === 0 && g5 === 0xffff;
+  const isIpv4Compatible = g0 === 0 && g1 === 0 && g2 === 0 && g3 === 0 && g4 === 0 && g5 === 0;
+  const isNat64 = g0 === 0x64 && g1 === 0xff9b && g2 === 0 && g3 === 0 && g4 === 0 && g5 === 0;
+  if (isIpv4Mapped || isIpv4Compatible || isNat64) {
+    const embedded = `${(g6 >> 8) & 0xff}.${g6 & 0xff}.${(g7 >> 8) & 0xff}.${g7 & 0xff}`;
+    return isBlockedIpv4(embedded);
+  }
+
+  if ((g0 & 0xfe00) === 0xfc00) return true; // fc00::/7 unique local
+  if ((g0 & 0xffc0) === 0xfe80) return true; // fe80::/10 link-local
+  if ((g0 & 0xff00) === 0xff00) return true; // ff00::/8 multicast
+  return false;
+}
+
+/**
+ * Returns true when an IP literal must not be fetched: private, loopback,
+ * link-local, and other reserved ranges for both IPv4 and IPv6 (including IPv6
+ * forms that embed an IPv4 address). Unparseable input is blocked defensively.
+ */
+export function isBlockedAddress(ip: string): boolean {
+  const version = isIP(ip);
+  if (version === 4) return isBlockedIpv4(ip);
+  if (version === 6) return isBlockedIpv6(ip);
+  return true;
+}
+
+async function assertHostAllowed(hostname: string): Promise<void> {
+  if (isIP(hostname) !== 0) {
+    if (isBlockedAddress(hostname)) {
+      throw new FileProcessorError('blocked_host', `Host '${hostname}' is a blocked address.`);
+    }
+    return;
+  }
+
+  let addresses: Array<{ address: string }>;
+  try {
+    addresses = await dns.promises.lookup(hostname, { all: true });
+  } catch (error) {
+    const message = error instanceof Error && error.message ? error.message : String(error);
+    throw new FileProcessorError('fetch_failed', `Could not resolve host '${hostname}': ${message}`);
+  }
+
+  if (!addresses.length) {
+    throw new FileProcessorError('fetch_failed', `Could not resolve host '${hostname}'.`);
+  }
+
+  for (const { address } of addresses) {
+    if (isBlockedAddress(address)) {
+      throw new FileProcessorError('blocked_host', `Host '${hostname}' resolves to blocked address '${address}'.`);
+    }
+  }
+}
+
+async function fetchWithSsrfGuard(
+  initialUrl: string,
+  options: { allowedProtocols: Set<string>; maxRedirects: number; signal: AbortSignal },
+): Promise<{ response: Response; finalUrl: string }> {
+  let currentUrl = initialUrl;
+  for (let hop = 0; ; hop += 1) {
+    let parsed: URL;
+    try {
+      parsed = new URL(currentUrl);
+    } catch {
+      throw new FileProcessorError('invalid_url', `File URL '${currentUrl}' is not a valid URL.`);
+    }
+
+    if (!options.allowedProtocols.has(parsed.protocol)) {
+      throw new FileProcessorError('unsupported_protocol', `File URL protocol '${parsed.protocol}' is not allowed.`);
+    }
+
+    // SSRF guard: resolve and validate the host before connecting. A malicious DNS
+    // server could still rebind between this check and the fetch; closing that gap
+    // fully would require pinning the resolved address into the connection.
+    const hostname =
+      parsed.hostname.startsWith('[') && parsed.hostname.endsWith(']') ? parsed.hostname.slice(1, -1) : parsed.hostname;
+    await assertHostAllowed(hostname);
+
+    const response = await fetch(currentUrl, { signal: options.signal, redirect: 'manual' });
+    if (!REDIRECT_STATUSES.has(response.status)) {
+      return { response, finalUrl: currentUrl };
+    }
+
+    if (hop >= options.maxRedirects) {
+      throw new FileProcessorError(
+        'too_many_redirects',
+        `File URL '${initialUrl}' exceeded the limit of ${options.maxRedirects} redirects.`,
+      );
+    }
+
+    const location = response.headers.get('location');
+    if (!location) {
+      throw new FileProcessorError('fetch_failed', `Redirect from '${currentUrl}' is missing a Location header.`);
+    }
+
+    try {
+      currentUrl = new URL(location, currentUrl).toString();
+    } catch {
+      throw new FileProcessorError(
+        'invalid_url',
+        `Redirect Location '${location}' from '${currentUrl}' is not a valid URL.`,
+      );
+    }
+  }
+}
+
 export async function fetchFileAsBase64(
   url: string,
   options: FetchFileOptions = {},
 ): Promise<{ fileData: string; mimeType: string | undefined; filename: string | undefined }> {
   const maxSize = options.maxSizeBytes ?? DEFAULT_MAX_FILE_SIZE;
   const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS;
+  const maxRedirects = options.maxRedirects ?? DEFAULT_MAX_REDIRECTS;
   const allowedProtocols = options.allowedProtocols ? new Set(options.allowedProtocols) : ALLOWED_PROTOCOLS;
 
-  let protocol: string;
-  try {
-    protocol = new URL(url).protocol;
-  } catch {
-    throw new FileProcessorError('invalid_url', `File URL '${url}' is not a valid URL.`);
-  }
-
-  if (!allowedProtocols.has(protocol)) {
-    throw new FileProcessorError('unsupported_protocol', `File URL protocol '${protocol}' is not allowed.`);
-  }
-
   const controller = new AbortController();
   const timeout = setTimeout(() => controller.abort(), timeoutMs);
 
   try {
-    const response = await fetch(url, {
+    const { response, finalUrl } = await fetchWithSsrfGuard(url, {
+      allowedProtocols,
+      maxRedirects,
       signal: controller.signal,
-      redirect: 'follow',
     });
 
     if (!response.ok) {
       throw new FileProcessorError(
         'fetch_failed',
-        `Failed to fetch file from '${url}': ${response.status} ${response.statusText}`,
+        `Failed to fetch file from '${finalUrl}': ${response.status} ${response.statusText}`,
       );
     }
 
     const contentLength = response.headers.get('content-length');
     if (contentLength && Number(contentLength) > maxSize) {
-      throw new FileProcessorError('file_too_large', `File at '${url}' exceeds maximum allowed size.`);
+      throw new FileProcessorError('file_too_large', `File at '${finalUrl}' exceeds maximum allowed size.`);
     }
 
     const contentType = response.headers.get('content-type') || undefined;
@@ -177,12 +367,12 @@ export async function fetchFileAsBase64(
 
     const buffer = Buffer.from(await response.arrayBuffer());
     if (buffer.length > maxSize) {
-      throw new FileProcessorError('file_too_large', `File at '${url}' exceeds maximum allowed size.`);
+      throw new FileProcessorError('file_too_large', `File at '${finalUrl}' exceeds maximum allowed size.`);
     }
 
     const mimeType = contentType?.split(';')[0].trim() ?? 'application/octet-stream';
     const contentDisposition = response.headers.get('content-disposition');
-    const filename = extractFilename(url, contentDisposition) ?? 'file';
+    const filename = extractFilename(finalUrl, contentDisposition) ?? 'file';
 
     return {
       fileData: `data:${mimeType};base64,${buffer.toString('base64')}`,
@@ -292,7 +482,7 @@ export const pdfFileProcessor: FileProcessor = {
     block: FileContentBlock,
     context: FileProcessorContext,
   ): Promise<FileContentBlock | FileContentBlock[]> {
-    const config = await context.ctx.db.getRepository('aiApiConfig').findOne();
+    const config = await getAiApiConfig(context.ctx);
     if (!config?.pdfRenderPagesAsImages) {
       return block;
     }

package/src/server/usage.ts

@@ -93,7 +93,10 @@ export function normalizeUsage(value: unknown): Usage | undefined {
     prompt_tokens: prompt,
     completion_tokens: completion,
     total_tokens: total,
-    prompt_cache_tokens: extractPromptCacheTokens(source),
+    // Keep normalizeUsage idempotent: streaming routes normalize the chunk
+    // usage once and hand the result to setAiApiUsageResult, which normalizes
+    // again — an already-extracted prompt_cache_tokens must survive that pass.
+    prompt_cache_tokens: extractPromptCacheTokens(source) ?? normalizeTokenCount(source.prompt_cache_tokens),
   };
 }
 
@@ -228,6 +231,7 @@ export async function finishUsageRecord(ctx: Context, id: unknown, startedAt: nu
     costStatus: billing.costStatus ?? null,
     modelPriceId: billing.modelPriceId ?? null,
     quotaPolicyId: billing.quotaPolicyId ?? null,
+    groupId: billing.groupId ?? null,
     inputPricePerMillionTokens: billing.inputPricePerMillionTokens ?? null,
     outputPricePerMillionTokens: billing.outputPricePerMillionTokens ?? null,
     fixedCostPerRequest: billing.fixedCostPerRequest ?? null,

package/src/server/utils/direct-llm-context.ts

@@ -1,4 +1,5 @@
 import type { Context } from '@nocobase/actions';
+import { resolveRequestUserGroup } from './request-cache';
 
 export type ContextOverflowBehavior = 'reject' | 'truncate';
 
@@ -13,6 +14,7 @@ export interface OpenAIMessage {
 interface ModelMetadata {
   contextWindow: number;
   maxCompletionTokens: number;
+  systemPrompt?: string;
 }
 
 interface ContextPreparationOptions {
@@ -307,16 +309,16 @@ async function loadModelMetadata(ctx: Context, serviceName: string, modelId: str
       `Context metadata is not configured for '${serviceName}/${modelId}'. Configure context window and max completion tokens.`,
     );
   }
-  return { contextWindow, maxCompletionTokens };
+  const systemPromptValue = getValue<unknown>(row, 'systemPrompt');
+  const systemPrompt = typeof systemPromptValue === 'string' ? systemPromptValue.trim() : '';
+  return { contextWindow, maxCompletionTokens, ...(systemPrompt ? { systemPrompt } : {}) };
 }
 
 async function resolveOverflowBehavior(ctx: Context): Promise<ContextOverflowBehavior> {
   const userId = ctx.state.currentUser?.id;
   if (userId === null || userId === undefined) return 'reject';
-  const policy = await ctx.db.getRepository('aiApiUserQuotaPolicies').findOne({
-    filter: { userId, enabled: true },
-  });
-  return getValue<unknown>(policy, 'contextOverflowBehavior') === 'truncate' ? 'truncate' : 'reject';
+  const group = await resolveRequestUserGroup(ctx, userId);
+  return group.contextOverflowBehavior === 'truncate' ? 'truncate' : 'reject';
 }
 
 function resolveReservedOutputTokens(options: ContextPreparationOptions, metadata: ModelMetadata): number {
@@ -348,10 +350,17 @@ export async function prepareDirectLlmContext(
   }
 
   const fixedOverheadTokens = estimateValueTokens(options.tools) + (options.tools === undefined ? 0 : 4);
-  const originalEstimate = estimateMessagesTokens(options.messages) + fixedOverheadTokens;
+  // The initial system prompt from model metadata is prepended before every
+  // client message — including the client's own system prompt, which is kept.
+  // Being a system message, it is a fixed instruction that truncation never
+  // drops, and it counts toward the input budget.
+  const baseMessages: OpenAIMessage[] = metadata.systemPrompt
+    ? [{ role: 'system', content: metadata.systemPrompt }, ...options.messages]
+    : options.messages;
+  const originalEstimate = estimateMessagesTokens(baseMessages) + fixedOverheadTokens;
   if (originalEstimate <= inputTokenBudget) {
     return {
-      messages: options.messages,
+      messages: baseMessages,
       estimatedInputTokens: originalEstimate,
       inputTokenBudget,
       reservedOutputTokens,
@@ -366,14 +375,14 @@ export async function prepareDirectLlmContext(
     );
   }
 
-  const { turns } = splitTurns(options.messages);
+  const { turns } = splitTurns(baseMessages);
   let remainingTurns = turns;
-  let messages = messagesWithTurns(options.messages, remainingTurns);
+  let messages = messagesWithTurns(baseMessages, remainingTurns);
   let estimatedInputTokens = estimateMessagesTokens(messages) + fixedOverheadTokens;
 
   while (remainingTurns.length > 1 && estimatedInputTokens > inputTokenBudget) {
     remainingTurns = remainingTurns.slice(1);
-    messages = messagesWithTurns(options.messages, remainingTurns);
+    messages = messagesWithTurns(baseMessages, remainingTurns);
     estimatedInputTokens = estimateMessagesTokens(messages) + fixedOverheadTokens;
   }
 
@@ -389,6 +398,6 @@ export async function prepareDirectLlmContext(
     estimatedInputTokens,
     inputTokenBudget,
     reservedOutputTokens,
-    truncated: messages.length !== options.messages.length,
+    truncated: messages.length !== baseMessages.length,
   };
 }

package/src/server/utils/rate-limiter.ts

@@ -26,7 +26,7 @@ export class RateLimiter {
    * Check and record a request for a user.
    *
    * @param userId  The user ID (string or numeric)
-   * @param limit   Max allowed requests per window (from aiApiConfig.rateLimitPerMinute)
+   * @param limit   Max allowed requests per window (from the user's usage group)
    * @returns { allowed: true } or { allowed: false, retryAfterMs: number }
    */
   check(userId: string | number, limit: number): { allowed: true } | { allowed: false; retryAfterMs: number } {

package/src/server/utils/request-cache.ts

@@ -0,0 +1,59 @@
+/**
+ * This file is part of the NocoBase (R) project.
+ * Copyright (c) 2020-2024 NocoBase Co., Ltd.
+ * Authors: NocoBase Team.
+ *
+ * This project is dual-licensed under AGPL-3.0 and NocoBase Commercial License.
+ * For more information, please refer to: https://www.nocobase.com/agreement.
+ */
+
+import type { Context } from '@nocobase/actions';
+import type { Model } from '@nocobase/database';
+import { resolveUserGroup, type AiApiUsageGroup } from '../quota-groups';
+
+interface AiApiRequestCache {
+  configLoaded: boolean;
+  config: Model | null;
+  groupKey?: string;
+  group?: AiApiUsageGroup;
+}
+
+function getCache(ctx: Context): AiApiRequestCache {
+  if (!ctx.state.aiApiRequestCache) {
+    ctx.state.aiApiRequestCache = { configLoaded: false, config: null };
+  }
+  return ctx.state.aiApiRequestCache;
+}
+
+/**
+ * Returns the aiApiConfig row, reading it at most once per request. The same row
+ * is consulted by the body limit, mode resolution, model whitelist, quota and
+ * billing steps of a single gateway request, so caching it on ctx.state removes
+ * several duplicate queries per request.
+ */
+export async function getAiApiConfig(ctx: Context): Promise<Model | null> {
+  const cache = getCache(ctx);
+  if (!cache.configLoaded) {
+    cache.config = (await ctx.db.getRepository('aiApiConfig').findOne()) ?? null;
+    cache.configLoaded = true;
+  }
+  return cache.config;
+}
+
+/**
+ * Returns the caller's usage group, resolving it at most once per request per
+ * user id. Wraps resolveUserGroup, which otherwise runs membership queries on
+ * every call site (rate limiting, permissions, billing, context overflow).
+ */
+export async function resolveRequestUserGroup(
+  ctx: Context,
+  userId: string | number | bigint | undefined | null,
+): Promise<AiApiUsageGroup> {
+  const cache = getCache(ctx);
+  const groupKey = userId === undefined || userId === null ? '' : String(userId);
+  if (!cache.group || cache.groupKey !== groupKey) {
+    cache.group = await resolveUserGroup(ctx, userId);
+    cache.groupKey = groupKey;
+  }
+  return cache.group;
+}

package/src/server/utils/resolve-service.ts

@@ -8,6 +8,7 @@
  */
 
 import { Context } from '@nocobase/actions';
+import { getAiApiConfig } from './request-cache';
 
 /**
  * Resolve an LLM service by name or title.
@@ -64,7 +65,7 @@ export async function resolveModelString(
   }
 
   // ─── Strategy 2: Use default LLM service from config ───
-  const config = await ctx.db.getRepository('aiApiConfig').findOne();
+  const config = await getAiApiConfig(ctx);
   if (config?.defaultLlmService) {
     const service = await repo.findOne({ filter: { name: config.defaultLlmService } });
     if (service) {