plugin-ai-api 1.0.8 → 1.0.10

package/src/server/middleware/role-permission.ts

@@ -1,66 +1,79 @@
-/**
- * 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 { Context } from '@nocobase/actions';
-import { toOpenAIError } from '../utils/openai-format';
-
-/**
- * Check whether the authenticated role is allowed to use the AI API.
- * Loads the permission record and stores it in ctx.state.aiApiRolePermission.
- *
- * Returns true if access is allowed (caller may proceed).
- * Returns false if access is denied (403 already written to ctx, caller must return).
- *
- * The 'root' and 'admin' roles always bypass the check.
- */
-export async function checkRolePermission(ctx: Context): Promise<boolean> {
-  const roleName = ctx.state.currentRoles?.[0] || 'member';
-
-  // root / admin always allowed
-  if (roleName === 'root' || roleName === 'admin') {
-    return true;
-  }
-
-  const record = await ctx.db.getRepository('aiApiRolePermissions').findOne({
-    filter: { roleName },
-  });
-
-  if (!record?.enabled) {
-    ctx.status = 403;
-    ctx.body = toOpenAIError(
-      403,
-      `Role '${roleName}' is not authorized to use the AI API. ` +
-        `An admin must enable access in Settings → Users & Permissions → [Role] → AI API.`,
-      'permission_denied',
-      'role_not_permitted',
-    );
-    return false;
-  }
-
-  // Store for downstream handlers
-  ctx.state.aiApiRolePermission = record;
-  return true;
-}
-
-/**
- * Check whether the current role is allowed to use a specific AI Employee.
- * Must be called after checkRolePermission (so ctx.state.aiApiRolePermission is set).
- *
- * Returns true when:
- * - Role is admin/root (no permission record stored)
- * - allowAllEmployees is true
- * - The employeeUsername is in the allowedEmployees list
- */
-export function checkEmployeeAccess(ctx: Context, employeeUsername: string): boolean {
-  const perm = ctx.state.aiApiRolePermission;
-  // admin/root paths have no record stored → always allowed
-  if (!perm) return true;
-  if (perm.allowAllEmployees) return true;
-  return ((perm.allowedEmployees as string[]) || []).includes(employeeUsername);
-}
+/**
+ * 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 { Context } from '@nocobase/actions';
+import { toOpenAIError } from '../utils/openai-format';
+
+const PERMISSION_TTL_MS = 15_000;
+const permissionCache = new Map<string, { record: any; expiresAt: number }>();
+
+export function invalidateRolePermissionCache(roleName?: string): void {
+  if (roleName) permissionCache.delete(roleName);
+  else permissionCache.clear();
+}
+
+/**
+ * Check whether the authenticated role is allowed to use the AI API.
+ * Loads the permission record and stores it in ctx.state.aiApiRolePermission.
+ *
+ * Returns true if access is allowed (caller may proceed).
+ * Returns false if access is denied (403 already written to ctx, caller must return).
+ *
+ * The 'root' and 'admin' roles always bypass the check.
+ */
+export async function checkRolePermission(ctx: Context): Promise<boolean> {
+  const roleName = ctx.state.currentRoles?.[0] || 'member';
+
+  // root / admin always allowed
+  if (roleName === 'root' || roleName === 'admin') {
+    return true;
+  }
+
+  const cached = permissionCache.get(roleName);
+  const record =
+    cached && cached.expiresAt > Date.now()
+      ? cached.record
+      : await ctx.db.getRepository('aiApiRolePermissions').findOne({ filter: { roleName } });
+  if (!cached || cached.expiresAt <= Date.now()) {
+    permissionCache.set(roleName, { record, expiresAt: Date.now() + PERMISSION_TTL_MS });
+  }
+
+  if (!record?.enabled) {
+    ctx.status = 403;
+    ctx.body = toOpenAIError(
+      403,
+      `Role '${roleName}' is not authorized to use the AI API. ` +
+        `An admin must enable access in Settings → Users & Permissions → [Role] → AI API.`,
+      'permission_denied',
+      'role_not_permitted',
+    );
+    return false;
+  }
+
+  // Store for downstream handlers
+  ctx.state.aiApiRolePermission = record;
+  return true;
+}
+
+/**
+ * Check whether the current role is allowed to use a specific AI Employee.
+ * Must be called after checkRolePermission (so ctx.state.aiApiRolePermission is set).
+ *
+ * Returns true when:
+ * - Role is admin/root (no permission record stored)
+ * - allowAllEmployees is true
+ * - The employeeUsername is in the allowedEmployees list
+ */
+export function checkEmployeeAccess(ctx: Context, employeeUsername: string): boolean {
+  const perm = ctx.state.aiApiRolePermission;
+  // admin/root paths have no record stored → always allowed
+  if (!perm) return true;
+  if (perm.allowAllEmployees) return true;
+  return ((perm.allowedEmployees as string[]) || []).includes(employeeUsername);
+}

package/src/server/plugin.ts

@@ -1,89 +1,99 @@
-/**
- * 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 { Plugin } from '@nocobase/server';
-import { createAiLlmRouter } from './routes/router';
-import aiApiConfigResource from './resource/ai-api-config';
-import { RateLimiter } from './utils/rate-limiter';
-
-// Ensure dayjs timezone + utc plugins are loaded.
-// Some Docker builds ship an older @nocobase/utils whose dayjs.js does not
-// extend the 'timezone' plugin, causing utcOffset(value) to behave as a
-// getter (returns a number) instead of a setter (returns a dayjs instance).
-// That breaks parse-filter.js's utc2unit() → "m.startOf is not a function".
-// Extending here patches the shared CommonJS dayjs module instance for the
-// entire Node.js process before any AIEmployee call is made.
-import dayjsLib from 'dayjs';
-import utcPlugin from 'dayjs/plugin/utc';
-import timezonePlugin from 'dayjs/plugin/timezone';
-(dayjsLib as any).extend(utcPlugin);
-(dayjsLib as any).extend(timezonePlugin);
-
-export class PluginAiApiServer extends Plugin {
-  /**
-   * Singleton rate limiter — lives for the entire plugin lifetime, shared across all requests.
-   * Uses a 1-minute sliding window to enforce rateLimitPerMinute from aiApiConfig.
-   */
-  rateLimiter = new RateLimiter(60_000);
-
-  private gcInterval: NodeJS.Timeout | null = null;
-
-  async afterAdd() {}
-
-  async beforeLoad() {}
-
-  async load() {
-    // 1. Register raw Koa middleware for OpenAI-compatible endpoints
-    //    Must run before 'resourcer' so URL paths match OpenAI convention
-    this.app.use(createAiLlmRouter(this), { before: 'resourcer' });
-
-    // 2. Register admin config resource
-    this.app.resourceManager.define(aiApiConfigResource);
-
-    // 3. Set ACL permissions for admin config + role permissions management
-    this.app.acl.registerSnippet({
-      name: `pm.${this.name}.configuration`,
-      actions: ['aiApiConfig:*', 'aiApiRolePermissions:*'],
-    });
-
-    // 4. GC the rate limiter every 5 minutes to evict stale user entries.
-    //    .unref() prevents this timer from keeping the process alive on shutdown.
-    this.gcInterval = setInterval(() => this.rateLimiter.gc(), 5 * 60 * 1000);
-    this.gcInterval.unref();
-  }
-
-  async install() {
-    // Create default config record on first install
-    const existing = await this.db.getRepository('aiApiConfig').findOne();
-    if (!existing) {
-      await this.db.getRepository('aiApiConfig').create({
-        values: {
-          defaultAiEmployee: '',
-          enabledLlmServices: [],
-          rateLimitPerMinute: 60,
-        },
-      });
-    }
-  }
-
-  async afterEnable() {}
-
-  async afterDisable() {}
-
-  async remove() {
-    // Clean up the GC timer so we don't leak resources during hot-reload
-    if (this.gcInterval) {
-      clearInterval(this.gcInterval);
-      this.gcInterval = null;
-    }
-    this.rateLimiter.clear();
-  }
-}
-
-export default PluginAiApiServer;
+/**
+ * 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 { Plugin } from '@nocobase/server';
+import { createAiLlmRouter } from './routes/router';
+import aiApiConfigResource from './resource/ai-api-config';
+import { RateLimiter } from './utils/rate-limiter';
+import { invalidateRolePermissionCache } from './middleware/role-permission';
+
+// Ensure dayjs timezone + utc plugins are loaded.
+// Some Docker builds ship an older @nocobase/utils whose dayjs.js does not
+// extend the 'timezone' plugin, causing utcOffset(value) to behave as a
+// getter (returns a number) instead of a setter (returns a dayjs instance).
+// That breaks parse-filter.js's utc2unit() → "m.startOf is not a function".
+// Extending here patches the shared CommonJS dayjs module instance for the
+// entire Node.js process before any AIEmployee call is made.
+import dayjsLib from 'dayjs';
+import utcPlugin from 'dayjs/plugin/utc';
+import timezonePlugin from 'dayjs/plugin/timezone';
+(dayjsLib as any).extend(utcPlugin);
+(dayjsLib as any).extend(timezonePlugin);
+
+export class PluginAiApiServer extends Plugin {
+  /**
+   * Singleton rate limiter — lives for the entire plugin lifetime, shared across all requests.
+   * Uses a 1-minute sliding window to enforce rateLimitPerMinute from aiApiConfig.
+   */
+  rateLimiter = new RateLimiter(60_000);
+
+  private gcInterval: NodeJS.Timeout | null = null;
+
+  async afterAdd() {}
+
+  async beforeLoad() {}
+
+  async load() {
+    // 1. 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.
+    this.app.use(createAiLlmRouter(this), { after: 'idp-oauth-resource-auth', before: 'resourcer' });
+
+    // 2. Register admin config resource
+    this.app.resourceManager.define(aiApiConfigResource);
+
+    this.app.db.on('aiApiRolePermissions.afterSave', (model) => {
+      invalidateRolePermissionCache(model.get('roleName'));
+    });
+    this.app.db.on('aiApiRolePermissions.afterDestroy', (model) => {
+      invalidateRolePermissionCache(model.get('roleName'));
+    });
+
+    // 3. Set ACL permissions for admin config + role permissions management
+    this.app.acl.registerSnippet({
+      name: `pm.${this.name}.configuration`,
+      actions: ['aiApiConfig:*', 'aiApiRolePermissions:*'],
+    });
+
+    // 4. GC the rate limiter every 5 minutes to evict stale user entries.
+    //    .unref() prevents this timer from keeping the process alive on shutdown.
+    this.gcInterval = setInterval(() => this.rateLimiter.gc(), 5 * 60 * 1000);
+    this.gcInterval.unref();
+  }
+
+  async install() {
+    // Create default config record on first install
+    const existing = await this.db.getRepository('aiApiConfig').findOne();
+    if (!existing) {
+      await this.db.getRepository('aiApiConfig').create({
+        values: {
+          defaultAiEmployee: '',
+          enabledLlmServices: [],
+          rateLimitPerMinute: 60,
+        },
+      });
+    }
+  }
+
+  async afterEnable() {}
+
+  async afterDisable() {}
+
+  async remove() {
+    // Clean up the GC timer so we don't leak resources during hot-reload
+    if (this.gcInterval) {
+      clearInterval(this.gcInterval);
+      this.gcInterval = null;
+    }
+    this.rateLimiter.clear();
+  }
+}
+
+export default PluginAiApiServer;

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

@@ -14,9 +14,11 @@ import {
   toOpenAIError,
   formatSSE,
   formatSSEDone,
+  OpenAIToolCallChunk,
 } from '../utils/openai-format';
 import { resolveModelString } from '../utils/resolve-service';
 import { checkEmployeeAccess } from '../middleware/role-permission';
+import { isStreamingRequested } from '../utils/streaming';
 import {
   AgentRuntimeContext,
   createAIEmployeeOptions,
@@ -132,7 +134,7 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
     return;
   }
 
-  const wantStream = body.stream === true;
+  const wantStream = isStreamingRequested(body.stream);
   const lifecycle = getAgentRuntimeLifecycle(ctx);
   let runtimeContext: AgentRuntimeContext | undefined;
   let lifecycleCompleted = false;
@@ -256,6 +258,17 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
 
       const originalWrite = ctx.res.write.bind(ctx.res);
       const originalEnd = ctx.res.end.bind(ctx.res);
+      const aiPlugin = ctx.app.pm.get('ai') as any;
+      const abortAgent = () => {
+        if (!ctx.res.writableEnded) {
+          aiPlugin?.aiEmployeesManager?.conversationController?.get(String(sessionId))?.abort();
+        }
+      };
+      ctx.req.once('aborted', abortAgent);
+      ctx.res.once('close', abortAgent);
+      let streamSucceeded = false;
+      let sawToolCalls = false;
+      let pendingSse = '';
 
       // Intercept end() — prevent AIEmployee from terminating the stream early.
       // We restore and call it ourselves in the finally block.
@@ -269,46 +282,78 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
 
       // Intercept write() — translate NocoBase SSE → OpenAI SSE
       (ctx.res as any).write = (data: Buffer | string): boolean => {
-        const text = typeof data === 'string' ? data : data.toString('utf8');
-
-        for (const line of text.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
-              originalWrite(
-                formatSSE(
-                  toOpenAIStreamChunk({
-                    id: completionId,
-                    model: body.model,
-                    delta: { content: String(event.body) },
+        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
+                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) {
+                  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
+                originalWrite(
+                  formatSSE({
+                    error: {
+                      message: String(event.body),
+                      type: 'server_error',
+                      code: 'agent_error',
+                    },
                   }),
-                ),
-              );
-            } 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
             }
-            // stream_start, stream_end, tool_calls, tool_call_status,
-            // web_search, reasoning, new_message, tool_call_chunks → silently ignored.
-            // These are NocoBase-internal events not part of the OpenAI protocol.
-          } catch {
-            // Non-JSON SSE line — ignore
           }
         }
         return true;
@@ -322,7 +367,10 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
           }),
         );
 
-        await aiEmployee.stream({ userMessages });
+        streamSucceeded = await aiEmployee.stream({ userMessages });
+        if (!streamSucceeded) {
+          throw new Error('AI Employee stream failed');
+        }
         try {
           await lifecycle?.runAfterHooks(runtimeContext, { succeeded: true });
         } finally {
@@ -332,20 +380,26 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
         // Restore original write/end before sending our closing frames
         (ctx.res as any).write = originalWrite;
         (ctx.res as any).end = originalEnd;
-
-        // Send the finish chunk and [DONE] signal
-        originalWrite(
-          formatSSE(
-            toOpenAIStreamChunk({
-              id: completionId,
-              model: body.model,
-              delta: {},
-              finishReason: 'stop',
-            }),
-          ),
-        );
-        originalWrite(formatSSEDone());
-        originalEnd();
+        ctx.req.off('aborted', abortAgent);
+        ctx.res.off('close', abortAgent);
+
+        if (streamSucceeded && !ctx.res.destroyed) {
+          originalWrite(
+            formatSSE(
+              toOpenAIStreamChunk({
+                id: completionId,
+                model: body.model,
+                delta: {},
+                finishReason: 'stop',
+              }),
+            ),
+          );
+          originalWrite(formatSSEDone());
+          ctx.state.aiApiStreamResult = { succeeded: true, id: completionId };
+        } else {
+          ctx.state.aiApiStreamResult = { succeeded: false, id: completionId, errorCode: 'agent_error' };
+        }
+        if (!ctx.res.writableEnded && !ctx.res.destroyed) originalEnd();
       }
     } else {
       // ── NON-STREAMING (invoke) ─────────────────────────────────────────────
@@ -422,6 +476,19 @@ export async function handleAgentCompletions(ctx: Context, plugin: PluginAiApiSe
   }
 }
 
+function toOpenAIToolCallChunks(value: unknown[]): OpenAIToolCallChunk[] {
+  return value.map((call: any, fallbackIndex) => ({
+    index: typeof call.index === 'number' ? call.index : fallbackIndex,
+    ...(call.id ? { id: String(call.id), type: 'function' as const } : {}),
+    function: {
+      ...(call.name ? { name: String(call.name) } : {}),
+      ...(call.args !== undefined
+        ? { arguments: typeof call.args === 'string' ? call.args : JSON.stringify(call.args) }
+        : {}),
+    },
+  }));
+}
+
 /**
  * Extract the last AI message content from a LangGraph invoke() result.
  *