plugin-ai-api 1.0.15 → 1.0.21

package/src/server/collections/ai-api-user-quota-policies.ts

@@ -0,0 +1,32 @@
+import { defineCollection } from '@nocobase/database';
+
+export default defineCollection({
+  name: 'aiApiUserQuotaPolicies',
+  autoGenId: true,
+  fields: [
+    { name: 'userId', type: 'bigInt', allowNull: false, index: true },
+    {
+      name: 'user',
+      type: 'belongsTo',
+      target: 'users',
+      targetKey: 'id',
+      foreignKey: 'userId',
+      constraints: false,
+    },
+    { name: 'enabled', type: 'boolean', defaultValue: true, index: true },
+    { name: 'periodType', type: 'string', allowNull: false, defaultValue: 'monthly' },
+    { name: 'timezone', type: 'string', allowNull: false, defaultValue: 'UTC' },
+    { name: 'requestLimit', type: 'bigInt', allowNull: true },
+    { name: 'totalTokenLimit', type: 'bigInt', allowNull: true },
+    { name: 'costLimit', type: 'decimal', precision: 20, scale: 8, allowNull: true },
+    { name: 'currency', type: 'string', allowNull: false, defaultValue: 'USD' },
+    { name: 'rejectUnpricedModel', type: 'boolean', defaultValue: true },
+    { name: 'missingUsageBehavior', type: 'string', allowNull: false, defaultValue: 'use_reserved' },
+  ],
+  indexes: [
+    {
+      fields: ['userId'],
+      unique: true,
+    },
+  ],
+});

package/src/server/plugin.ts

@@ -8,10 +8,13 @@
  */
 
 import { Plugin } from '@nocobase/server';
-import { createAiLlmRouter } from './routes/router';
+import { createAiLlmRouter, AI_LLM_PREFIX } from './routes/router';
 import aiApiConfigResource from './resource/ai-api-config';
+import aiApiUsageMonitorResource from './resource/ai-api-usage-monitor';
 import { RateLimiter } from './utils/rate-limiter';
 import { invalidateRolePermissionCache } from './middleware/role-permission';
+import { validateModelPrice, validateModelMetadata, validateQuotaPolicy } from './validation';
+import { AI_API_ACL_SNIPPET } from '../constants';
 
 // Ensure dayjs timezone + utc plugins are loaded.
 // Some Docker builds ship an older @nocobase/utils whose dayjs.js does not
@@ -37,10 +40,35 @@ export class PluginAiApiServer extends Plugin {
 
   async afterAdd() {}
 
-  async beforeLoad() {}
+  async beforeLoad() {
+    this.app.db.on('aiApiModelPrices.beforeSave', async (model) => {
+      await validateModelPrice(this.db, model);
+    });
+    this.app.db.on('aiApiModelMetadata.beforeSave', (model) => {
+      validateModelMetadata(model);
+    });
+    this.app.db.on('aiApiUserQuotaPolicies.beforeSave', (model) => {
+      validateQuotaPolicy(model);
+    });
+  }
 
   async load() {
-    // 1. Register raw Koa middleware for OpenAI-compatible endpoints
+    // 1. Claim body parsing for our own routes before the core bodyParser runs.
+    //    Core registers koa-bodyparser with a global REQUEST_BODY_LIMIT (10mb by
+    //    default) much earlier in the stack, so without this the gateway's own
+    //    configurable limit is unreachable: an oversized vision request would be
+    //    rejected by core with a non-OpenAI error shape. `disableBodyParser` makes
+    //    koa-bodyparser skip the request, leaving ctx.request.body undefined so
+    //    createAiLlmRouter reads and caps the raw stream itself.
+    this.app.use(
+      async (ctx, next) => {
+        if (ctx.path.startsWith(AI_LLM_PREFIX)) ctx.disableBodyParser = true;
+        await next();
+      },
+      { tag: 'aiApiDisableBodyParser', before: 'bodyParser' },
+    );
+
+    // 2. Register raw Koa middleware for OpenAI-compatible endpoints
     //    Must run before 'resourcer' so URL paths match OpenAI convention
     // OIDC access tokens must first pass through plugin-idp-oauth, which validates
     // issuer/audience/scope and rewrites them to a NocoBase internal token.
@@ -48,6 +76,7 @@ export class PluginAiApiServer extends Plugin {
 
     // 2. Register admin config resource
     this.app.resourceManager.define(aiApiConfigResource);
+    this.app.resourceManager.define(aiApiUsageMonitorResource);
 
     this.app.db.on('aiApiRolePermissions.afterSave', (model) => {
       invalidateRolePermissionCache(model.get('roleName'));
@@ -58,8 +87,19 @@ export class PluginAiApiServer extends Plugin {
 
     // 3. Set ACL permissions for admin config + role permissions management
     this.app.acl.registerSnippet({
-      name: `pm.${this.name}.configuration`,
-      actions: ['aiApiConfig:*', 'aiApiRolePermissions:*'],
+      name: AI_API_ACL_SNIPPET,
+      actions: [
+        'aiApiConfig:*',
+        'aiApiRolePermissions:*',
+        'aiApiModelPrices:*',
+        'aiApiModelMetadata:*',
+        'aiApiUserQuotaPolicies:*',
+        'aiApiUserQuotaBuckets:list',
+        'aiApiUserQuotaBuckets:get',
+        'aiApiUsageRecords:list',
+        'aiApiUsageRecords:get',
+        'aiApiUsageMonitor:summary',
+      ],
     });
 
     // 4. GC the rate limiter every 5 minutes to evict stale user entries.
@@ -77,6 +117,8 @@ export class PluginAiApiServer extends Plugin {
           defaultAiEmployee: '',
           enabledLlmServices: [],
           rateLimitPerMinute: 60,
+          quotaEnabled: false,
+          defaultReservationOutputTokens: 4096,
         },
       });
     }

package/src/server/resource/ai-api-config.ts

@@ -1,74 +1,105 @@
-/**
- * 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 { ResourceOptions } from '@nocobase/resourcer';
-
-/**
- * Resource for managing AI API configuration via NocoBase admin UI.
- * Singleton config pattern (same as aiSettings in plugin-ai).
- *
- * Uses custom action name 'save' instead of 'update' because NocoBase's
- * built-in middleware requires filter/filterByTk for the standard 'update' action.
- */
-const aiApiConfigResource: ResourceOptions = {
-  name: 'aiApiConfig',
-  actions: {
-    async get(ctx, next) {
-      let config = await ctx.db.getRepository('aiApiConfig').findOne();
-      if (!config) {
-        config = await ctx.db.getRepository('aiApiConfig').create({
-          values: {
-            mode: 'llm',
-            defaultAiEmployee: '',
-            defaultLlmService: '',
-            enabledLlmServices: [],
-            rateLimitPerMinute: 60,
-            options: {},
-          },
-        });
-      }
-      ctx.body = config;
-      await next();
-    },
-
-    async save(ctx, next) {
-      const values = ctx.action.params.values || (ctx.request.body as any) || {};
-      const repo = ctx.db.getRepository('aiApiConfig');
-      let config = await repo.findOne();
-
-      if (!config) {
-        config = await repo.create({
-          values: {
-            mode: values.mode ?? 'llm',
-            defaultAiEmployee: values.defaultAiEmployee ?? '',
-            defaultLlmService: values.defaultLlmService ?? '',
-            enabledLlmServices: values.enabledLlmServices ?? [],
-            rateLimitPerMinute: values.rateLimitPerMinute ?? 60,
-            options: values.options ?? {},
-          },
-        });
-      } else {
-        const updateData: Record<string, any> = {};
-        if (values.mode !== undefined) updateData.mode = values.mode;
-        if (values.defaultAiEmployee !== undefined) updateData.defaultAiEmployee = values.defaultAiEmployee;
-        if (values.defaultLlmService !== undefined) updateData.defaultLlmService = values.defaultLlmService;
-        if (values.enabledLlmServices !== undefined) updateData.enabledLlmServices = values.enabledLlmServices;
-        if (values.rateLimitPerMinute !== undefined) updateData.rateLimitPerMinute = values.rateLimitPerMinute;
-        if (values.options !== undefined) updateData.options = values.options;
-
-        await config.update(updateData);
-      }
-
-      ctx.body = config;
-      await next();
-    },
-  },
-};
-
-export default aiApiConfigResource;
+/**
+ * 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 { ResourceOptions } from '@nocobase/resourcer';
+import { MAX_REQUEST_BODY_MB_LIMIT } from '../routes/router';
+
+const DEFAULT_MAX_REQUEST_BODY_MB = 10;
+
+/**
+ * The gateway buffers each request body in memory, so an out-of-range value
+ * here is a denial-of-service footgun rather than a harmless setting.
+ */
+function coerceMaxRequestBodyMb(value: unknown): number {
+  const mb = Number(value);
+  if (!Number.isSafeInteger(mb) || mb <= 0) {
+    throw new Error(`maxRequestBodyMb must be a positive integer (1-${MAX_REQUEST_BODY_MB_LIMIT}).`);
+  }
+  if (mb > MAX_REQUEST_BODY_MB_LIMIT) {
+    throw new Error(`maxRequestBodyMb cannot exceed ${MAX_REQUEST_BODY_MB_LIMIT}.`);
+  }
+  return mb;
+}
+
+/**
+ * Resource for managing AI API configuration via NocoBase admin UI.
+ * Singleton config pattern (same as aiSettings in plugin-ai).
+ *
+ * Uses custom action name 'save' instead of 'update' because NocoBase's
+ * built-in middleware requires filter/filterByTk for the standard 'update' action.
+ */
+const aiApiConfigResource: ResourceOptions = {
+  name: 'aiApiConfig',
+  actions: {
+    async get(ctx, next) {
+      let config = await ctx.db.getRepository('aiApiConfig').findOne();
+      if (!config) {
+        config = await ctx.db.getRepository('aiApiConfig').create({
+          values: {
+            mode: 'llm',
+            defaultAiEmployee: '',
+            defaultLlmService: '',
+            enabledLlmServices: [],
+            rateLimitPerMinute: 60,
+            maxRequestBodyMb: 10,
+            quotaEnabled: false,
+            defaultReservationOutputTokens: 4096,
+            options: {},
+          },
+        });
+      }
+      ctx.body = config;
+      await next();
+    },
+
+    async save(ctx, next) {
+      const values = ctx.action.params.values || (ctx.request.body as any) || {};
+      const repo = ctx.db.getRepository('aiApiConfig');
+      let config = await repo.findOne();
+
+      if (!config) {
+        config = await repo.create({
+          values: {
+            mode: values.mode ?? 'llm',
+            defaultAiEmployee: values.defaultAiEmployee ?? '',
+            defaultLlmService: values.defaultLlmService ?? '',
+            enabledLlmServices: values.enabledLlmServices ?? [],
+            rateLimitPerMinute: values.rateLimitPerMinute ?? 60,
+            maxRequestBodyMb: coerceMaxRequestBodyMb(values.maxRequestBodyMb ?? DEFAULT_MAX_REQUEST_BODY_MB),
+            quotaEnabled: values.quotaEnabled ?? false,
+            defaultReservationOutputTokens: values.defaultReservationOutputTokens ?? 4096,
+            options: values.options ?? {},
+          },
+        });
+      } else {
+        const updateData: Record<string, any> = {};
+        if (values.mode !== undefined) updateData.mode = values.mode;
+        if (values.defaultAiEmployee !== undefined) updateData.defaultAiEmployee = values.defaultAiEmployee;
+        if (values.defaultLlmService !== undefined) updateData.defaultLlmService = values.defaultLlmService;
+        if (values.enabledLlmServices !== undefined) updateData.enabledLlmServices = values.enabledLlmServices;
+        if (values.rateLimitPerMinute !== undefined) updateData.rateLimitPerMinute = values.rateLimitPerMinute;
+        if (values.maxRequestBodyMb !== undefined) {
+          updateData.maxRequestBodyMb = coerceMaxRequestBodyMb(values.maxRequestBodyMb);
+        }
+        if (values.quotaEnabled !== undefined) updateData.quotaEnabled = values.quotaEnabled;
+        if (values.defaultReservationOutputTokens !== undefined) {
+          updateData.defaultReservationOutputTokens = values.defaultReservationOutputTokens;
+        }
+        if (values.options !== undefined) updateData.options = values.options;
+
+        await config.update(updateData);
+      }
+
+      ctx.body = config;
+      await next();
+    },
+  },
+};
+
+export default aiApiConfigResource;

package/src/server/resource/ai-api-usage-monitor.ts

@@ -0,0 +1,74 @@
+import type { Context } from '@nocobase/actions';
+import type { ResourceOptions } from '@nocobase/resourcer';
+import { Op, col, fn } from 'sequelize';
+
+interface UsageSummaryRow {
+  requestCount?: string | number;
+  inputTokens?: string | number;
+  outputTokens?: string | number;
+  totalTokens?: string | number;
+}
+
+interface CostSummaryRow {
+  currency?: string;
+  totalCost?: string | number;
+}
+
+function buildWhere(ctx: Context) {
+  const params = ctx.action.params;
+  const where: Record<string, unknown> = {};
+  const start = typeof params.start === 'string' ? new Date(params.start) : undefined;
+  const end = typeof params.end === 'string' ? new Date(params.end) : undefined;
+
+  if ((start && !Number.isNaN(start.getTime())) || (end && !Number.isNaN(end.getTime()))) {
+    const startedAt: Record<symbol, Date> = {};
+    if (start && !Number.isNaN(start.getTime())) startedAt[Op.gte] = start;
+    if (end && !Number.isNaN(end.getTime())) startedAt[Op.lte] = end;
+    where.startedAt = startedAt;
+  }
+  if (params.userId !== undefined && params.userId !== '') where.userId = params.userId;
+  if (params.resolvedService) where.resolvedService = params.resolvedService;
+  if (params.resolvedModel) where.resolvedModel = params.resolvedModel;
+  if (params.status) where.status = params.status;
+  return where;
+}
+
+const aiApiUsageMonitorResource: ResourceOptions = {
+  name: 'aiApiUsageMonitor',
+  actions: {
+    async summary(ctx, next) {
+      const model = ctx.db.getCollection('aiApiUsageRecords').model;
+      const where = buildWhere(ctx);
+      const totals = (await model.findOne({
+        attributes: [
+          [fn('COUNT', col('id')), 'requestCount'],
+          [fn('COALESCE', fn('SUM', col('inputTokens')), 0), 'inputTokens'],
+          [fn('COALESCE', fn('SUM', col('outputTokens')), 0), 'outputTokens'],
+          [fn('COALESCE', fn('SUM', col('totalTokens')), 0), 'totalTokens'],
+        ],
+        where,
+        raw: true,
+      })) as unknown as UsageSummaryRow;
+      const costs = (await model.findAll({
+        attributes: ['currency', [fn('COALESCE', fn('SUM', col('estimatedCost')), 0), 'totalCost']],
+        where: { ...where, estimatedCost: { [Op.ne]: null } },
+        group: ['currency'],
+        raw: true,
+      })) as unknown as CostSummaryRow[];
+
+      ctx.body = {
+        requestCount: Number(totals?.requestCount ?? 0),
+        inputTokens: Number(totals?.inputTokens ?? 0),
+        outputTokens: Number(totals?.outputTokens ?? 0),
+        totalTokens: Number(totals?.totalTokens ?? 0),
+        costsByCurrency: costs.map((item) => ({
+          currency: item.currency || 'USD',
+          totalCost: String(item.totalCost ?? 0),
+        })),
+      };
+      await next();
+    },
+  },
+};
+
+export default aiApiUsageMonitorResource;

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

@@ -27,6 +27,7 @@ import {
 } from '../utils/ai-employee-runtime';
 import { setAiApiUsageUnavailable } from '../usage';
 import type PluginAiApiServer from '../plugin';
+import { markAiApiFirstProviderOutput } from '../utils/app-observability';
 
 /**
  * POST /api/ai-llm/v1/chat/completions (agent mode)
@@ -281,85 +282,97 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
         // Don't actually end — we control this
       };
 
-      // Intercept write() — translate NocoBase SSE → OpenAI SSE
-      (ctx.res as any).write = (data: Buffer | string): boolean => {
-        pendingSse += typeof data === 'string' ? data : data.toString('utf8');
-        const frames = pendingSse.split('\n\n');
-        pendingSse = frames.pop() || '';
-
-        for (const frame of frames) {
-          for (const line of frame.split('\n')) {
-            const trimmed = line.trim();
-            if (!trimmed.startsWith('data: ')) continue;
-
-            const jsonStr = trimmed.substring(6);
-            if (!jsonStr) continue;
-
-            try {
-              const event = JSON.parse(jsonStr);
-
-              if (event.type === 'content' && event.body) {
-                // Content chunk — forward as OpenAI delta
+      // Translate one buffered NocoBase SSE frame → OpenAI SSE.
+      const processSseFrame = (frame: string): void => {
+        for (const line of frame.split('\n')) {
+          const trimmed = line.trim();
+          if (!trimmed.startsWith('data: ')) continue;
+
+          const jsonStr = trimmed.substring(6);
+          if (!jsonStr) continue;
+
+          try {
+            const event = JSON.parse(jsonStr);
+
+            if (event.type === 'content' && event.body) {
+              markAiApiFirstProviderOutput(ctx);
+              // Content chunk — forward as OpenAI delta
+              originalWrite(
+                formatSSE(
+                  toOpenAIStreamChunk({
+                    id: completionId,
+                    model: body.model,
+                    delta: { content: String(event.body) },
+                  }),
+                ),
+              );
+            } else if (event.type === 'tool_call_chunks' && Array.isArray(event.body)) {
+              const chunks = toOpenAIToolCallChunks(event.body);
+              if (chunks.length) {
+                markAiApiFirstProviderOutput(ctx);
+                sawToolCalls = true;
                 originalWrite(
                   formatSSE(
                     toOpenAIStreamChunk({
                       id: completionId,
                       model: body.model,
-                      delta: { content: String(event.body) },
+                      delta: { tool_calls: chunks },
                     }),
                   ),
                 );
-              } else if (event.type === 'tool_call_chunks' && Array.isArray(event.body)) {
-                const chunks = toOpenAIToolCallChunks(event.body);
-                if (chunks.length) {
-                  sawToolCalls = true;
-                  originalWrite(
-                    formatSSE(
-                      toOpenAIStreamChunk({
-                        id: completionId,
-                        model: body.model,
-                        delta: { tool_calls: chunks },
-                      }),
-                    ),
-                  );
-                }
-              } else if (!sawToolCalls && event.type === 'tool_calls' && Array.isArray(event.body?.toolCalls)) {
-                const chunks = toOpenAIToolCallChunks(event.body.toolCalls);
-                if (chunks.length) {
-                  sawToolCalls = true;
-                  originalWrite(
-                    formatSSE(
-                      toOpenAIStreamChunk({
-                        id: completionId,
-                        model: body.model,
-                        delta: { tool_calls: chunks },
-                      }),
-                    ),
-                  );
-                }
-              } else if (event.type === 'error' && event.body) {
-                // Error from the agent — surface as SSE error object
+              }
+            } else if (!sawToolCalls && event.type === 'tool_calls' && Array.isArray(event.body?.toolCalls)) {
+              const chunks = toOpenAIToolCallChunks(event.body.toolCalls);
+              if (chunks.length) {
+                markAiApiFirstProviderOutput(ctx);
+                sawToolCalls = true;
                 originalWrite(
-                  formatSSE({
-                    error: {
-                      message: String(event.body),
-                      type: 'server_error',
-                      code: 'agent_error',
-                    },
-                  }),
+                  formatSSE(
+                    toOpenAIStreamChunk({
+                      id: completionId,
+                      model: body.model,
+                      delta: { tool_calls: chunks },
+                    }),
+                  ),
                 );
               }
-              // stream_start, stream_end, tool_call_status, web_search,
-              // reasoning and new_message are NocoBase-only events and are ignored.
-              // These are NocoBase-internal events not part of the OpenAI protocol.
-            } catch {
-              // Non-JSON SSE line — ignore
+            } else if (event.type === 'error' && event.body) {
+              // Error from the agent — surface as SSE error object
+              originalWrite(
+                formatSSE({
+                  error: {
+                    message: String(event.body),
+                    type: 'server_error',
+                    code: 'agent_error',
+                  },
+                }),
+              );
             }
+            // stream_start, stream_end, tool_call_status, web_search,
+            // reasoning and new_message are NocoBase-only events and are ignored.
+            // These are NocoBase-internal events not part of the OpenAI protocol.
+          } catch {
+            // Non-JSON SSE line — ignore
           }
         }
+      };
+
+      // Intercept write() — buffer input and translate complete frames.
+      (ctx.res as any).write = (data: Buffer | string): boolean => {
+        pendingSse += typeof data === 'string' ? data : data.toString('utf8');
+        const frames = pendingSse.split('\n\n');
+        pendingSse = frames.pop() || '';
+        for (const frame of frames) processSseFrame(frame);
         return true;
       };
 
+      // Flush any trailing frame that was not terminated by '\n\n'.
+      const flushPendingSse = (): void => {
+        if (!pendingSse.trim()) return;
+        processSseFrame(pendingSse);
+        pendingSse = '';
+      };
+
       try {
         const aiEmployee = new AIEmployee(
           createAIEmployeeOptions(ctx, employeeRecord, sessionId, {
@@ -385,6 +398,8 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
         ctx.res.off('close', abortAgent);
 
         if (streamSucceeded && !ctx.res.destroyed) {
+          // Emit any final frame the agent left unterminated before closing.
+          flushPendingSse();
           originalWrite(
             formatSSE(
               toOpenAIStreamChunk({

package/src/server/routes/auth.ts

@@ -47,7 +47,20 @@ export async function authenticateBearer(ctx: Context): Promise<boolean> {
         const rolesRepository = ctx.db.getRepository('users.roles', ctx.state.currentUser.id);
         const roles = await rolesRepository.find({ fields: ['name'] });
         const roleNames = roles.map((role: { name: string }) => role.name);
-        ctx.state.currentRole = roleNames.includes(requestedRole) ? requestedRole : roleNames[0];
+        // An explicit X-Role that the user does not hold must be rejected, not
+        // silently downgraded to the first role — otherwise a caller could probe
+        // for access under a role they were never granted.
+        if (requestedRole && !roleNames.includes(requestedRole)) {
+          ctx.status = 403;
+          ctx.body = toOpenAIError(
+            403,
+            `Requested role '${requestedRole}' is not assigned to this user`,
+            'permission_denied',
+            'role_not_permitted',
+          );
+          return false;
+        }
+        ctx.state.currentRole = requestedRole || roleNames[0];
         ctx.state.currentRoles = ctx.state.currentRole ? [ctx.state.currentRole] : roleNames;
       }
       return true;