plugin-ai-api 1.0.13 → 1.0.15

package/src/server/migrations/20260727140000-change-usage-user-id-to-bigint.ts

@@ -0,0 +1,222 @@
+import { DataTypes } from '@nocobase/database';
+import { Migration } from '@nocobase/server';
+import { QueryTypes, type TableName } from 'sequelize';
+
+const MAX_SIGNED_BIGINT = 9_223_372_036_854_775_807n;
+const PREFLIGHT_BATCH_SIZE = 1_000;
+
+interface LegacyUsageRow {
+  recordId: string | number | bigint;
+  userId: unknown;
+}
+
+interface InvalidUsageRow {
+  recordId: string;
+  userId: string;
+}
+
+interface InvalidUsageSummary {
+  count: number;
+  samples: InvalidUsageRow[];
+}
+
+function stringifyDatabaseValue(value: unknown): string {
+  if (typeof value === 'bigint') return value.toString();
+  if (value === null) return 'null';
+  if (value === undefined) return 'undefined';
+  return String(value);
+}
+
+export function normalizeLegacyUserId(value: unknown): string | undefined {
+  let normalized: string;
+
+  if (typeof value === 'string') {
+    normalized = value.trim();
+  } else if (typeof value === 'bigint') {
+    normalized = value.toString();
+  } else if (typeof value === 'number' && Number.isSafeInteger(value)) {
+    normalized = String(value);
+  } else {
+    return undefined;
+  }
+
+  if (!/^\d+$/.test(normalized)) return undefined;
+
+  const parsed = BigInt(normalized);
+  if (parsed > MAX_SIGNED_BIGINT) return undefined;
+
+  return normalized;
+}
+
+export function isBigIntColumnType(value: unknown): boolean {
+  return typeof value === 'string' && value.toUpperCase().includes('BIGINT');
+}
+
+export function isLegacyStringColumnType(value: unknown): boolean {
+  if (typeof value !== 'string') return false;
+  const normalized = value.toUpperCase();
+  return normalized.includes('CHAR') || normalized.includes('TEXT');
+}
+
+export default class ChangeUsageUserIdToBigInt extends Migration {
+  on = 'beforeLoad' as const;
+
+  async up() {
+    const { collection, temporary } = this.getUsageCollection();
+    try {
+      if (!(await collection.existsInDb())) return;
+
+      const userIdField = collection.getField('userId');
+      if (!userIdField) {
+        throw new Error('AI API usage migration could not resolve the userId field.');
+      }
+
+      const tableName = collection.getTableNameWithSchema();
+      const columnName = userIdField.columnName();
+      const columns = await this.queryInterface.describeTable(tableName);
+      const column = columns[columnName];
+      if (!column) {
+        throw new Error(`AI API usage migration could not find the physical column ${columnName}.`);
+      }
+
+      if (!isBigIntColumnType(column.type) && !isLegacyStringColumnType(column.type)) {
+        throw new Error(`AI API usage migration does not support converting ${column.type} to BIGINT.`);
+      }
+
+      const invalidRows = await this.findInvalidRows(collection.quotedTableName(), columnName);
+      if (invalidRows.count > 0) {
+        const samples = invalidRows.samples.map((row) => `${row.recordId}:${JSON.stringify(row.userId)}`).join(', ');
+        throw new Error(
+          `AI API usage migration found ${invalidRows.count} invalid userId value(s). ` +
+            `Clean these records before retrying the upgrade. Samples: ${samples}`,
+        );
+      }
+
+      if (isBigIntColumnType(column.type)) {
+        if (column.allowNull === false) return;
+        await this.makeColumnNotNull(tableName, collection.quotedTableName(), columnName);
+        return;
+      }
+
+      await this.changeColumnToBigInt(tableName, collection.quotedTableName(), columnName);
+    } finally {
+      if (temporary) this.db.removeCollection('aiApiUsageRecords');
+    }
+  }
+
+  async down() {
+    const { collection, temporary } = this.getUsageCollection();
+    try {
+      if (!(await collection.existsInDb())) return;
+
+      const userIdField = collection.getField('userId');
+      if (!userIdField) return;
+
+      const tableName = collection.getTableNameWithSchema();
+      const columnName = userIdField.columnName();
+      const columns = await this.queryInterface.describeTable(tableName);
+      if (!isBigIntColumnType(columns[columnName]?.type)) return;
+
+      if (this.db.inDialect('postgres')) {
+        const quotedColumn = this.db.quoteIdentifier(columnName);
+        await this.sequelize.transaction(async (transaction) => {
+          await this.sequelize.query(
+            `ALTER TABLE ${collection.quotedTableName()} ` +
+              `ALTER COLUMN ${quotedColumn} DROP NOT NULL, ` +
+              `ALTER COLUMN ${quotedColumn} TYPE VARCHAR(255) USING ${quotedColumn}::text`,
+            { transaction },
+          );
+        });
+        return;
+      }
+
+      await this.queryInterface.changeColumn(tableName, columnName, {
+        type: DataTypes.STRING,
+        allowNull: true,
+      });
+    } finally {
+      if (temporary) this.db.removeCollection('aiApiUsageRecords');
+    }
+  }
+
+  private getUsageCollection() {
+    const existing = this.db.getCollection('aiApiUsageRecords');
+    if (existing) return { collection: existing, temporary: false };
+
+    const collection = this.db.collection({
+      name: 'aiApiUsageRecords',
+      fields: [{ name: 'userId', type: 'string' }],
+    });
+    return { collection, temporary: true };
+  }
+
+  private async findInvalidRows(quotedTableName: string, columnName: string): Promise<InvalidUsageSummary> {
+    const quotedId = this.db.quoteIdentifier('id');
+    const quotedUserId = this.db.quoteIdentifier(columnName);
+    const invalidRows: InvalidUsageSummary = { count: 0, samples: [] };
+    let offset = 0;
+    let hasMoreRows = true;
+
+    while (hasMoreRows) {
+      const rows = await this.sequelize.query<LegacyUsageRow>(
+        `SELECT ${quotedId} AS ${this.db.quoteIdentifier('recordId')}, ` +
+          `${quotedUserId} AS ${this.db.quoteIdentifier('userId')} ` +
+          `FROM ${quotedTableName} ORDER BY ${quotedId} ASC LIMIT :limit OFFSET :offset`,
+        {
+          replacements: { limit: PREFLIGHT_BATCH_SIZE, offset },
+          type: QueryTypes.SELECT,
+        },
+      );
+
+      for (const row of rows) {
+        if (normalizeLegacyUserId(row.userId) === undefined) {
+          invalidRows.count += 1;
+          if (invalidRows.samples.length < 20) {
+            invalidRows.samples.push({
+              recordId: stringifyDatabaseValue(row.recordId),
+              userId: stringifyDatabaseValue(row.userId),
+            });
+          }
+        }
+      }
+
+      hasMoreRows = rows.length === PREFLIGHT_BATCH_SIZE;
+      offset += rows.length;
+    }
+
+    return invalidRows;
+  }
+
+  private async changeColumnToBigInt(tableName: TableName, quotedTableName: string, columnName: string) {
+    if (this.db.inDialect('postgres')) {
+      const quotedColumn = this.db.quoteIdentifier(columnName);
+      await this.sequelize.transaction(async (transaction) => {
+        await this.sequelize.query(
+          `ALTER TABLE ${quotedTableName} ` +
+            `ALTER COLUMN ${quotedColumn} TYPE BIGINT USING BTRIM(${quotedColumn})::BIGINT, ` +
+            `ALTER COLUMN ${quotedColumn} SET NOT NULL`,
+          { transaction },
+        );
+      });
+      return;
+    }
+
+    await this.queryInterface.changeColumn(tableName, columnName, {
+      type: DataTypes.BIGINT,
+      allowNull: false,
+    });
+  }
+
+  private async makeColumnNotNull(tableName: TableName, quotedTableName: string, columnName: string) {
+    if (this.db.inDialect('postgres')) {
+      const quotedColumn = this.db.quoteIdentifier(columnName);
+      await this.sequelize.query(`ALTER TABLE ${quotedTableName} ALTER COLUMN ${quotedColumn} SET NOT NULL`);
+      return;
+    }
+
+    await this.queryInterface.changeColumn(tableName, columnName, {
+      type: DataTypes.BIGINT,
+      allowNull: false,
+    });
+  }
+}

package/src/server/routes/agent-completions.ts

@@ -25,6 +25,7 @@ import {
   getAgentRuntimeLifecycle,
   loadAIEmployeeConstructor,
 } from '../utils/ai-employee-runtime';
+import { setAiApiUsageUnavailable } from '../usage';
 import type PluginAiApiServer from '../plugin';
 
 /**
@@ -395,8 +396,10 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
             ),
           );
           originalWrite(formatSSEDone());
+          setAiApiUsageUnavailable(ctx, completionId);
           ctx.state.aiApiStreamResult = { succeeded: true, id: completionId };
         } else {
+          setAiApiUsageUnavailable(ctx, completionId);
           ctx.state.aiApiStreamResult = { succeeded: false, id: completionId, errorCode: 'agent_error' };
         }
         if (!ctx.res.writableEnded && !ctx.res.destroyed) originalEnd();
@@ -443,6 +446,7 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
         // Clients can detect agent mode by checking usage.total_tokens === 0.
         usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 },
       };
+      setAiApiUsageUnavailable(ctx, completionId);
     }
 
     // ─── Cleanup: delete the ephemeral conversation (best-effort) ────────────

package/src/server/routes/auth.ts

@@ -41,6 +41,7 @@ export async function authenticateBearer(ctx: Context): Promise<boolean> {
 
   try {
     if (ctx.state.currentUser) {
+      ctx.state.aiApiAuthType = ctx.state.oauthPrincipal ? 'oidc' : 'bearer';
       if (!ctx.state.currentRole) {
         const requestedRole = ctx.get('X-Role');
         const rolesRepository = ctx.db.getRepository('users.roles', ctx.state.currentUser.id);

package/src/server/routes/chat-completions.ts

@@ -21,6 +21,7 @@ import {
 import { resolveModelString } from '../utils/resolve-service';
 import { createRequestAbortController, isStreamingRequested, writeResponse } from '../utils/streaming';
 import { checkEmployeeAccess } from '../middleware/role-permission';
+import { extractProviderRequestId, normalizeUsage, setAiApiUsageResult, type Usage } from '../usage';
 import type PluginAiApiServer from '../plugin';
 
 /**
@@ -249,13 +250,10 @@ async function handleNonStreamingCompletion(
   }
 
   // Extract usage if available
-  const usage = result.usage_metadata
-    ? {
-        prompt_tokens: result.usage_metadata.input_tokens || 0,
-        completion_tokens: result.usage_metadata.output_tokens || 0,
-        total_tokens: result.usage_metadata.total_tokens || 0,
-      }
-    : { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 };
+  const usage = setAiApiUsageResult(ctx, result.usage_metadata, {
+    gatewayResponseId: completionId,
+    providerRequestId: extractProviderRequestId(result),
+  });
 
   ctx.status = 200;
   const toolCalls = normalizeToolCalls(result.tool_calls);
@@ -263,7 +261,7 @@ async function handleNonStreamingCompletion(
     id: completionId,
     model: modelName,
     content,
-    usage,
+    usage: usage ?? { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 },
     toolCalls,
   });
 }
@@ -300,7 +298,8 @@ async function handleStreamingCompletion(
   );
 
   const requestAbort = createRequestAbortController(ctx);
-  let usage: { prompt_tokens: number; completion_tokens: number; total_tokens: number } | undefined;
+  let usage: Usage | undefined;
+  let providerRequestId: string | undefined;
   let finishReason = 'stop';
   try {
     const stream = await chatModel.stream(messages, { ...providerRequestParameters, signal: requestAbort.signal });
@@ -337,12 +336,9 @@ async function handleStreamingCompletion(
         );
       }
       if (chunk.usage_metadata) {
-        usage = {
-          prompt_tokens: chunk.usage_metadata.input_tokens || 0,
-          completion_tokens: chunk.usage_metadata.output_tokens || 0,
-          total_tokens: chunk.usage_metadata.total_tokens || 0,
-        };
+        usage = normalizeUsage(chunk.usage_metadata) ?? usage;
       }
+      providerRequestId = providerRequestId ?? extractProviderRequestId(chunk);
     }
 
     // Send finish chunk
@@ -360,7 +356,8 @@ async function handleStreamingCompletion(
 
     // Send [DONE]
     await writeResponse(ctx, formatSSEDone());
-    ctx.state.aiApiStreamResult = { succeeded: true, id: completionId, usage };
+    setAiApiUsageResult(ctx, usage, { gatewayResponseId: completionId, providerRequestId });
+    ctx.state.aiApiStreamResult = { succeeded: true, id: completionId };
   } catch (err) {
     ctx.log.error('AI API streaming error:', err);
     // Send error as SSE event before closing
@@ -375,7 +372,8 @@ async function handleStreamingCompletion(
         }),
       );
     }
-    ctx.state.aiApiStreamResult = { succeeded: false, id: completionId, usage, errorCode: 'stream_error' };
+    setAiApiUsageResult(ctx, usage, { gatewayResponseId: completionId, providerRequestId });
+    ctx.state.aiApiStreamResult = { succeeded: false, id: completionId, errorCode: 'stream_error' };
   } finally {
     requestAbort.dispose();
     if (!ctx.res.writableEnded && !ctx.res.destroyed) ctx.res.end();

package/src/server/routes/completions.ts

@@ -11,6 +11,7 @@ import { Context } from '@nocobase/actions';
 import { generateCompletionId, toOpenAIError, formatSSE, formatSSEDone } from '../utils/openai-format';
 import { resolveModelString } from '../utils/resolve-service';
 import { createRequestAbortController, isStreamingRequested, writeResponse } from '../utils/streaming';
+import { extractProviderRequestId, normalizeUsage, setAiApiUsageResult, type Usage } from '../usage';
 import type PluginAiApiServer from '../plugin';
 
 /**
@@ -187,13 +188,10 @@ async function handleNonStreamingTextCompletion(
     text = textPart?.text || JSON.stringify(result.content);
   }
 
-  const usage = result.usage_metadata
-    ? {
-        prompt_tokens: result.usage_metadata.input_tokens || 0,
-        completion_tokens: result.usage_metadata.output_tokens || 0,
-        total_tokens: result.usage_metadata.total_tokens || 0,
-      }
-    : { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 };
+  const usage = setAiApiUsageResult(ctx, result.usage_metadata, {
+    gatewayResponseId: completionId,
+    providerRequestId: extractProviderRequestId(result),
+  });
 
   ctx.status = 200;
   ctx.body = {
@@ -210,7 +208,7 @@ async function handleNonStreamingTextCompletion(
         finish_reason: 'stop',
       },
     ],
-    usage,
+    usage: usage ?? { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 },
   };
 }
 
@@ -232,7 +230,8 @@ async function handleStreamingTextCompletion(
   ctx.status = 200;
 
   const requestAbort = createRequestAbortController(ctx);
-  let usage: { prompt_tokens: number; completion_tokens: number; total_tokens: number } | undefined;
+  let usage: Usage | undefined;
+  let providerRequestId: string | undefined;
   try {
     const stream = await chatModel.stream(messages, { signal: requestAbort.signal });
 
@@ -267,12 +266,9 @@ async function handleStreamingTextCompletion(
         );
       }
       if (chunk.usage_metadata) {
-        usage = {
-          prompt_tokens: chunk.usage_metadata.input_tokens || 0,
-          completion_tokens: chunk.usage_metadata.output_tokens || 0,
-          total_tokens: chunk.usage_metadata.total_tokens || 0,
-        };
+        usage = normalizeUsage(chunk.usage_metadata) ?? usage;
       }
+      providerRequestId = providerRequestId ?? extractProviderRequestId(chunk);
     }
 
     // Final chunk
@@ -296,7 +292,8 @@ async function handleStreamingTextCompletion(
     );
 
     await writeResponse(ctx, formatSSEDone());
-    ctx.state.aiApiStreamResult = { succeeded: true, id: completionId, usage };
+    setAiApiUsageResult(ctx, usage, { gatewayResponseId: completionId, providerRequestId });
+    ctx.state.aiApiStreamResult = { succeeded: true, id: completionId };
   } catch (err) {
     ctx.log.error('AI API completions streaming error:', err);
     if (!ctx.res.destroyed && !ctx.res.writableEnded) {
@@ -310,7 +307,8 @@ async function handleStreamingTextCompletion(
         }),
       );
     }
-    ctx.state.aiApiStreamResult = { succeeded: false, id: completionId, usage, errorCode: 'stream_error' };
+    setAiApiUsageResult(ctx, usage, { gatewayResponseId: completionId, providerRequestId });
+    ctx.state.aiApiStreamResult = { succeeded: false, id: completionId, errorCode: 'stream_error' };
   } finally {
     requestAbort.dispose();
     if (!ctx.res.writableEnded && !ctx.res.destroyed) ctx.res.end();

package/src/server/routes/embeddings.ts

@@ -10,6 +10,7 @@
 import { Context } from '@nocobase/actions';
 import { toOpenAIError, toOpenAIEmbeddingsResponse } from '../utils/openai-format';
 import { resolveModelString } from '../utils/resolve-service';
+import { setAiApiUsageUnavailable } from '../usage';
 import type PluginAiApiServer from '../plugin';
 
 /**
@@ -175,6 +176,7 @@ export async function handleEmbeddings(ctx: Context, plugin: PluginAiApiServer)
 
     ctx.status = 200;
     ctx.set('Content-Type', 'application/json');
+    setAiApiUsageUnavailable(ctx);
     ctx.body = toOpenAIEmbeddingsResponse({
       model: body.model,
       embeddings: vectors,

package/src/server/routes/router.ts

@@ -24,6 +24,8 @@ import type PluginAiApiServer from '../plugin';
 
 const API_PREFIX = '/api/ai-llm/v1';
 
+type DataWrappingContext = Context & { withoutDataWrapping?: boolean };
+
 /**
  * Main Koa middleware router for OpenAI-compatible endpoints.
  *
@@ -60,7 +62,7 @@ export function createAiLlmRouter(plugin: PluginAiApiServer) {
 
     // Prevent NocoBase's dataWrapping middleware from wrapping OpenAI-format responses
     // in an extra {"data": ...} envelope, which breaks OpenAI-compatible clients like n8n.
-    (ctx as any).withoutDataWrapping = true;
+    (ctx as DataWrappingContext).withoutDataWrapping = true;
 
     // Parse the sub-path after prefix
     const subPath = path.substring(API_PREFIX.length);
@@ -89,8 +91,9 @@ export function createAiLlmRouter(plugin: PluginAiApiServer) {
       try {
         const rawBody = await getRawBody(ctx);
         ctx.request.body = JSON.parse(rawBody);
-      } catch (bodyErr: any) {
-        const status = bodyErr?.statusCode === 413 ? 413 : 400;
+      } catch (bodyErr: unknown) {
+        const status =
+          bodyErr && typeof bodyErr === 'object' && 'statusCode' in bodyErr && bodyErr.statusCode === 413 ? 413 : 400;
         const message = status === 413 ? 'Request body too large (max 10 MB)' : 'Invalid JSON in request body';
         ctx.status = status;
         ctx.body = toOpenAIError(status, message, 'invalid_request_error');
@@ -120,13 +123,14 @@ export function createAiLlmRouter(plugin: PluginAiApiServer) {
     }
 
     // ─── Route matching ───────────────────────────────────────────────────
-    const model = (ctx.request.body as any)?.model ?? '-';
+    const requestBody = (ctx.request.body || {}) as Record<string, unknown>;
+    const model = requestBody.model === undefined || requestBody.model === null ? '-' : String(requestBody.model);
     const isUsageEndpoint =
       method === 'POST' && (subPath === '/chat/completions' || subPath === '/completions' || subPath === '/embeddings');
+    const isStreamingEndpoint = method === 'POST' && (subPath === '/chat/completions' || subPath === '/completions');
     const resolvedMode = isUsageEndpoint ? await resolveMode(ctx) : 'llm';
-    const requestBody = (ctx.request.body || {}) as Record<string, unknown>;
-    const streaming = isUsageEndpoint && isStreamingRequested(requestBody.stream);
-    if (streaming && (subPath === '/chat/completions' || subPath === '/completions')) {
+    const streaming = isStreamingEndpoint && isStreamingRequested(requestBody.stream);
+    if (streaming) {
       const streamOptions = requestBody.stream_options;
       ctx.request.body = {
         ...requestBody,
@@ -140,7 +144,7 @@ export function createAiLlmRouter(plugin: PluginAiApiServer) {
     let usageId: unknown;
     try {
       usageId = isUsageEndpoint
-        ? await startUsageRecord(ctx, requestId, subPath, String(model), streaming, resolvedMode)
+        ? await startUsageRecord(ctx, requestId, subPath, model, streaming, resolvedMode)
         : undefined;
     } catch (usageError) {
       ctx.log.error('AI API usage record could not be created:', usageError);
@@ -172,8 +176,8 @@ export function createAiLlmRouter(plugin: PluginAiApiServer) {
         const completionsMode = resolvedMode;
         if (completionsMode === 'agent') {
           // Convert legacy prompt → messages format for agent handler
-          const reqBody = ctx.request.body as any;
-          if (reqBody?.prompt !== undefined) {
+          const reqBody = (ctx.request.body || {}) as Record<string, unknown>;
+          if (reqBody.prompt !== undefined) {
             const prompt =
               typeof reqBody.prompt === 'string'
                 ? reqBody.prompt