plugin-ai-api 1.0.21 → 1.0.24

package/src/server/__tests__/models.test.ts

@@ -7,11 +7,33 @@
  * For more information, please refer to: https://www.nocobase.com/agreement.
  */
 
-import { describe, expect, it } from 'vitest';
-import { buildModelObject } from '../routes/models';
+import { Context } from '@nocobase/actions';
+import { describe, expect, it, vi } from 'vitest';
+import { buildModelObject, handleGetModel, handleListModels } from '../routes/models';
 
 const CREATED = 1_700_000_000;
 
+function permissionLookupFailureContext() {
+  const ctx = {
+    app: { name: 'main', pm: { get: () => ({}) } },
+    state: { currentUser: { id: 1 } },
+    db: {
+      getRepository: (name: string) => {
+        if (name === 'aiApiConfig') return { findOne: vi.fn(async () => null) };
+        if (name === 'llmServices') return { find: vi.fn(async () => []) };
+        if (name === 'aiApiUserPermissions') {
+          return { findOne: vi.fn(async () => Promise.reject(new Error('permission database unavailable'))) };
+        }
+        return { find: vi.fn(async () => []) };
+      },
+    },
+    log: { error: vi.fn(), warn: vi.fn() },
+    status: 0,
+    body: undefined,
+  } as unknown as Context;
+  return ctx;
+}
+
 describe('buildModelObject', () => {
   it('returns the base OpenAI model shape with no override', () => {
     const model = buildModelObject('svc/gpt-4o', CREATED, 'My Service');
@@ -72,3 +94,23 @@ describe('buildModelObject', () => {
     expect(buildModelObject('svc/m', CREATED, 'Svc', { contextWindow: 10 }).active).toBe(true);
   });
 });
+
+describe('model catalog permission lookup failures', () => {
+  it('returns a retryable 503 when listing models', async () => {
+    const ctx = permissionLookupFailureContext();
+
+    await handleListModels(ctx, undefined as never);
+
+    expect(ctx.status).toBe(503);
+    expect(ctx.body).toMatchObject({ error: { code: 'permission_check_failed' } });
+  });
+
+  it('returns a retryable 503 when retrieving a model', async () => {
+    const ctx = permissionLookupFailureContext();
+
+    await handleGetModel(ctx, 'openai/gpt-4o', undefined as never);
+
+    expect(ctx.status).toBe(503);
+    expect(ctx.body).toMatchObject({ error: { code: 'permission_check_failed' } });
+  });
+});

package/src/server/__tests__/openai-format.test.ts

@@ -1,4 +1,4 @@
-import { toOpenAIResponse, toOpenAIStreamChunk } from '../utils/openai-format';
+import { toOpenAIResponse, toOpenAIStreamChunk, toOpenAIUsageChunk } from '../utils/openai-format';
 import { isStreamingRequested } from '../utils/streaming';
 import { applyProviderRequestParameters, getProviderRequestParameters } from '../routes/chat-completions';
 
@@ -52,6 +52,57 @@ describe('AI API OpenAI tool-call formatting', () => {
   });
 });
 
+describe('AI API OpenAI usage-only streaming chunks', () => {
+  it('includes a null usage field for normal chunks', () => {
+    const chunk = toOpenAIStreamChunk({
+      id: 'chatcmpl-1',
+      model: 'service/model',
+      delta: { content: 'Hello' },
+    });
+
+    expect(chunk.choices).toHaveLength(1);
+    expect(chunk).toHaveProperty('usage', null);
+  });
+
+  it('includes a null usage field for finish chunks', () => {
+    const chunk = toOpenAIStreamChunk({
+      id: 'chatcmpl-1',
+      model: 'service/model',
+      delta: {},
+      finishReason: 'stop',
+    });
+
+    expect(chunk.choices[0].finish_reason).toBe('stop');
+    expect(chunk).toHaveProperty('usage', null);
+  });
+
+  it('formats a usage-only chat completion chunk with an empty choices array', () => {
+    const chunk = toOpenAIUsageChunk({
+      id: 'chatcmpl-1',
+      model: 'service/model',
+      usage: { prompt_tokens: 8, completion_tokens: 3, total_tokens: 11 },
+      object: 'chat.completion.chunk',
+    });
+
+    expect(chunk.object).toBe('chat.completion.chunk');
+    expect(chunk.choices).toEqual([]);
+    expect(chunk.usage).toEqual({ prompt_tokens: 8, completion_tokens: 3, total_tokens: 11 });
+  });
+
+  it('formats a usage-only legacy text completion chunk', () => {
+    const chunk = toOpenAIUsageChunk({
+      id: 'cmpl-1',
+      model: 'service/model',
+      usage: { prompt_tokens: 2, completion_tokens: 5, total_tokens: 7 },
+      object: 'text_completion',
+    });
+
+    expect(chunk.object).toBe('text_completion');
+    expect(chunk.choices).toEqual([]);
+    expect(chunk.usage).toEqual({ prompt_tokens: 2, completion_tokens: 5, total_tokens: 7 });
+  });
+});
+
 describe('AI API provider parameter forwarding', () => {
   it('forwards model and tool-call parameters not managed by the gateway', () => {
     const parameters = getProviderRequestParameters({

package/src/server/__tests__/permission-sync.test.ts

@@ -0,0 +1,109 @@
+/**
+ * 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 { beforeEach, describe, expect, it, vi } from 'vitest';
+import { PluginAiApiServer } from '../plugin';
+import { invalidateUserPermissionCache, resolveUserAccessScope } from '../utils/user-permissions';
+
+/**
+ * Covers the cross-node half of permission revocation.
+ *
+ * syncMessageManager hardcodes `skipSelf: true` (sync-message-manager.ts:59,73), so the node
+ * that writes the change never receives its own broadcast. That makes two things load-bearing
+ * and easy to regress: the writer must invalidate its own cache locally, and every other node
+ * must invalidate on receipt. Neither is observable from a single-node test of the cache alone.
+ */
+function mockContext(userId: number, row: unknown) {
+  const findOne = vi.fn(async () => row);
+  const ctx = {
+    state: { currentUser: { id: userId } },
+    db: { getRepository: () => ({ findOne }) },
+    app: { name: 'main' },
+    log: { warn: vi.fn(), error: vi.fn() },
+  } as unknown as Context;
+  return { ctx, findOne };
+}
+
+function row(values: Record<string, unknown>) {
+  return { get: (key: string) => values[key] };
+}
+
+beforeEach(() => {
+  invalidateUserPermissionCache();
+});
+
+describe('cross-node permission invalidation', () => {
+  it("drops the receiving node's cached scope", async () => {
+    const plugin = Object.create(PluginAiApiServer.prototype) as PluginAiApiServer;
+    const { ctx, findOne } = mockContext(1, row({ allowedLlmServices: ['openai'] }));
+
+    await resolveUserAccessScope(ctx);
+    await resolveUserAccessScope(ctx);
+    expect(findOne).toHaveBeenCalledTimes(1);
+
+    await plugin.handleSyncMessage({ type: 'invalidateUserPermissions', userId: 1 });
+
+    await resolveUserAccessScope(ctx);
+    expect(findOne).toHaveBeenCalledTimes(2);
+  });
+
+  it('ignores unrelated message types and other users', async () => {
+    const plugin = Object.create(PluginAiApiServer.prototype) as PluginAiApiServer;
+    const { ctx, findOne } = mockContext(1, row({ allowedLlmServices: ['openai'] }));
+    await resolveUserAccessScope(ctx);
+
+    await plugin.handleSyncMessage({ type: 'somethingElse', userId: 1 });
+    await plugin.handleSyncMessage({ type: 'invalidateUserPermissions', userId: 2 });
+
+    await resolveUserAccessScope(ctx);
+    expect(findOne).toHaveBeenCalledTimes(1);
+  });
+
+  it('tolerates a malformed message instead of throwing into the subscriber', async () => {
+    const plugin = Object.create(PluginAiApiServer.prototype) as PluginAiApiServer;
+    await expect(plugin.handleSyncMessage(undefined as never)).resolves.toBeUndefined();
+    await expect(plugin.handleSyncMessage({} as never)).resolves.toBeUndefined();
+  });
+
+  it('invalidates locally as well as broadcasting, since the publisher is skipped', async () => {
+    const plugin = Object.create(PluginAiApiServer.prototype) as PluginAiApiServer;
+    const sendSyncMessage = vi.fn(async () => undefined);
+    Object.assign(plugin, { sendSyncMessage });
+
+    const { ctx, findOne } = mockContext(1, row({ allowedLlmServices: ['openai'] }));
+    await resolveUserAccessScope(ctx);
+    expect(findOne).toHaveBeenCalledTimes(1);
+
+    // revokeUserPermissions is private; reach it the way the db hook does.
+    (plugin as unknown as { revokeUserPermissions: (id: unknown, tx?: unknown) => void }).revokeUserPermissions(1);
+
+    // Local cache cleared without any message coming back to us.
+    await resolveUserAccessScope(ctx);
+    expect(findOne).toHaveBeenCalledTimes(2);
+    expect(sendSyncMessage).toHaveBeenCalledWith(
+      { type: 'invalidateUserPermissions', userId: 1 },
+      { transaction: undefined },
+    );
+  });
+
+  it('defers the broadcast to the transaction so other nodes cannot re-cache the old row', async () => {
+    const plugin = Object.create(PluginAiApiServer.prototype) as PluginAiApiServer;
+    const sendSyncMessage = vi.fn(async () => undefined);
+    Object.assign(plugin, { sendSyncMessage });
+    const transaction = { id: 'tx-1' };
+
+    (plugin as unknown as { revokeUserPermissions: (id: unknown, tx?: unknown) => void }).revokeUserPermissions(
+      7,
+      transaction,
+    );
+
+    expect(sendSyncMessage).toHaveBeenCalledWith({ type: 'invalidateUserPermissions', userId: 7 }, { transaction });
+  });
+});

package/src/server/__tests__/usage-route.test.ts

@@ -2,6 +2,7 @@ import type { Context } from '@nocobase/actions';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import type PluginAiApiServer from '../plugin';
 import { handleChatCompletions } from '../routes/chat-completions';
+import { handleCompletions } from '../routes/completions';
 import { resolveModelString } from '../utils/resolve-service';
 
 vi.mock('../utils/resolve-service', () => ({
@@ -14,18 +15,20 @@ interface ModelResult {
   usage_metadata?: Record<string, unknown>;
 }
 
-function createContext(result: ModelResult): Context {
+function createContext(result: ModelResult) {
   const model = {
     invoke: vi.fn().mockResolvedValue(result),
     modelKwargs: {},
   };
+  let providerCreateCount = 0;
   class TestProvider {
     createModel() {
+      providerCreateCount += 1;
       return model;
     }
   }
 
-  return {
+  const ctx = {
     app: {
       pm: {
         get: vi.fn().mockReturnValue({
@@ -36,7 +39,16 @@ function createContext(result: ModelResult): Context {
       },
     },
     db: {
-      getRepository: vi.fn().mockReturnValue({ findOne: vi.fn().mockResolvedValue(null) }),
+      getRepository: vi.fn((name: string) => {
+        if (name === 'aiApiModelMetadata') {
+          return {
+            findOne: vi.fn().mockResolvedValue({
+              get: (key: string) => (key === 'contextWindow' ? 128_000 : key === 'maxCompletionTokens' ? 16_384 : true),
+            }),
+          };
+        }
+        return { findOne: vi.fn().mockResolvedValue(null) };
+      }),
     },
     log: { error: vi.fn() },
     request: {
@@ -49,6 +61,125 @@ function createContext(result: ModelResult): Context {
     state: {},
     set: vi.fn(),
   } as unknown as Context;
+
+  return { ctx, model, getProviderCreateCount: () => providerCreateCount };
+}
+
+class ListenerTarget {
+  private listeners = new Map<string, Set<() => void>>();
+  aborted = false;
+  writableEnded = false;
+
+  once(event: string, listener: () => void) {
+    const wrapped = () => {
+      this.off(event, wrapped);
+      listener();
+    };
+    const group = this.listeners.get(event) ?? new Set();
+    group.add(wrapped);
+    this.listeners.set(event, group);
+  }
+
+  off(event: string, listener: () => void) {
+    this.listeners.get(event)?.delete(listener);
+  }
+
+  emit(event: string) {
+    for (const listener of [...(this.listeners.get(event) ?? [])]) listener();
+  }
+}
+
+function createStreamingContext(
+  result: ModelResult,
+  streamOptions?: Record<string, unknown>,
+  requestBody?: Record<string, unknown>,
+) {
+  const req = new ListenerTarget();
+  const res = new ListenerTarget();
+  const writes: string[] = [];
+
+  res.write = vi.fn((data: unknown) => {
+    writes.push(String(data));
+    return true;
+  });
+  res.end = vi.fn(() => {
+    res.writableEnded = true;
+  });
+
+  const contentChunks = (typeof result.content === 'string' ? [result.content] : []).filter(Boolean);
+  const chunks = [
+    ...contentChunks.map((content) => ({ content })),
+    { content: '', usage_metadata: result.usage_metadata, response_metadata: result.response_metadata },
+  ];
+
+  const model = {
+    invoke: vi.fn().mockResolvedValue({
+      content: result.content,
+      usage_metadata: result.usage_metadata,
+    }),
+    stream: vi.fn().mockResolvedValue({
+      [Symbol.asyncIterator]() {
+        let index = 0;
+        return {
+          async next() {
+            if (index >= chunks.length) return { done: true, value: undefined };
+            return { done: false, value: chunks[index++] };
+          },
+          async return() {
+            index = chunks.length;
+            return { done: true, value: undefined };
+          },
+        };
+      },
+    }),
+    modelKwargs: {},
+  };
+
+  class TestProvider {
+    createModel() {
+      return model;
+    }
+  }
+
+  const ctx = {
+    app: {
+      pm: {
+        get: vi.fn().mockReturnValue({
+          aiManager: {
+            llmProviders: new Map([['test-provider', { provider: TestProvider }]]),
+          },
+        }),
+      },
+    },
+    db: {
+      getRepository: vi.fn((name: string) => {
+        if (name === 'aiApiModelMetadata') {
+          return {
+            findOne: vi.fn().mockResolvedValue({
+              get: (key: string) => (key === 'contextWindow' ? 128_000 : key === 'maxCompletionTokens' ? 16_384 : true),
+            }),
+          };
+        }
+        return { findOne: vi.fn().mockResolvedValue(null) };
+      }),
+    },
+    log: { error: vi.fn() },
+    req,
+    res,
+    request: {
+      body: {
+        model: 'test-service/test-model',
+        messages: [{ role: 'user', content: 'Hello' }],
+        stream: true,
+        stream_options: streamOptions,
+        ...requestBody,
+      },
+    },
+    state: {} as Record<string, unknown>,
+    set: vi.fn(),
+  } as unknown as Context;
+
+  return { ctx, model, writes };
 }
 
 describe('AI API chat usage collection', () => {
@@ -65,7 +196,7 @@ describe('AI API chat usage collection', () => {
   });
 
   it('keeps the public zero fallback but marks missing provider usage as unavailable internally', async () => {
-    const ctx = createContext({ content: 'Hello back' });
+    const { ctx } = createContext({ content: 'Hello back' });
 
     await handleChatCompletions(ctx, {} as PluginAiApiServer);
 
@@ -86,7 +217,7 @@ describe('AI API chat usage collection', () => {
   });
 
   it('stores provider usage and provider request ID separately from the gateway response ID', async () => {
-    const ctx = createContext({
+    const { ctx } = createContext({
       content: 'Hello back',
       response_metadata: { request_id: 'provider-request-1' },
       usage_metadata: { input_tokens: 8, output_tokens: 3, total_tokens: 11 },
@@ -106,4 +237,133 @@ describe('AI API chat usage collection', () => {
       usage: { prompt_tokens: 8, completion_tokens: 3, total_tokens: 11 },
     });
   });
+
+  it('emits a usage-only chunk immediately before [DONE] for streaming chat completions', async () => {
+    const { ctx, writes } = createStreamingContext(
+      {
+        content: 'Hi',
+        response_metadata: { request_id: 'provider-request-stream-1' },
+        usage_metadata: { input_tokens: 5, output_tokens: 4, total_tokens: 9 },
+      },
+      { include_usage: true },
+    );
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    const dataLines = writes.filter((line) => line.startsWith('data: '));
+    const doneFrame = dataLines.find((line) => line.includes('[DONE]'));
+    expect(doneFrame).toBeDefined();
+    const frames = dataLines.filter((line) => !line.includes('[DONE]')).map((line) => JSON.parse(line.slice(6)));
+    const doneIndex = frames.length;
+
+    const finishChunk = frames[doneIndex - 2];
+    const usageChunk = frames[doneIndex - 1];
+
+    expect(finishChunk.choices[0].finish_reason).toBe('stop');
+    expect(usageChunk.choices).toEqual([]);
+    expect(usageChunk.usage).toEqual({ prompt_tokens: 5, completion_tokens: 4, total_tokens: 9 });
+    expect(usageChunk).toHaveProperty('usage.prompt_tokens', 5);
+    expect(ctx.state.aiApiUsageResult).toMatchObject({
+      source: 'provider',
+      providerRequestId: 'provider-request-stream-1',
+    });
+  });
+
+  it('always emits chat usage and forces provider collection when include_usage is false', async () => {
+    const { ctx, model, writes } = createStreamingContext(
+      {
+        content: 'Hi',
+        usage_metadata: { input_tokens: 5, output_tokens: 4, total_tokens: 9 },
+      },
+      { include_usage: false, include_obfuscation: false },
+    );
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    const frames = writes
+      .filter((line) => line.startsWith('data: ') && !line.includes('[DONE]'))
+      .map((line) => JSON.parse(line.slice(6)));
+    const usageChunk = frames[frames.length - 1];
+
+    expect(frames.slice(0, -1).every((frame) => frame.usage === null)).toBe(true);
+    expect(usageChunk.choices).toEqual([]);
+    expect(usageChunk.usage).toEqual({ prompt_tokens: 5, completion_tokens: 4, total_tokens: 9 });
+    expect(model.stream).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({
+        stream_options: { include_usage: true, include_obfuscation: false },
+      }),
+    );
+  });
+
+  it('always emits legacy completion usage and forwards all stream options', async () => {
+    const { ctx, model, writes } = createStreamingContext(
+      {
+        content: 'Hi',
+        usage_metadata: { input_tokens: 2, output_tokens: 5, total_tokens: 7 },
+      },
+      { include_usage: false, include_obfuscation: false },
+      { prompt: 'Hello' },
+    );
+
+    await handleCompletions(ctx, {} as PluginAiApiServer);
+
+    const frames = writes
+      .filter((line) => line.startsWith('data: ') && !line.includes('[DONE]'))
+      .map((line) => JSON.parse(line.slice(6)));
+    const usageChunk = frames[frames.length - 1];
+
+    expect(frames.slice(0, -1).every((frame) => frame.usage === null)).toBe(true);
+    expect(usageChunk.object).toBe('text_completion');
+    expect(usageChunk.choices).toEqual([]);
+    expect(usageChunk.usage).toEqual({ prompt_tokens: 2, completion_tokens: 5, total_tokens: 7 });
+    expect(model.stream).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({
+        stream_options: { include_usage: true, include_obfuscation: false },
+      }),
+    );
+  });
+
+  it('rejects context overflow before creating a provider or reserving quota', async () => {
+    const { ctx, getProviderCreateCount } = createContext({ content: 'unused' });
+    (ctx.request.body as Record<string, unknown>).messages = [{ role: 'user', content: 'x'.repeat(600_000) }];
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(400);
+    expect(ctx.body).toMatchObject({ error: { code: 'context_length_exceeded' } });
+    expect(ctx.state.aiApiLlmBilling).toBeUndefined();
+    expect(getProviderCreateCount()).toBe(0);
+  });
+
+  it('rejects oversized legacy prompts before creating a provider or reserving quota', async () => {
+    const { ctx, model, getProviderCreateCount } = createContext({ content: 'unused' });
+    (ctx.request.body as Record<string, unknown>).prompt = 'x'.repeat(600_000);
+
+    await handleCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(400);
+    expect(ctx.body).toMatchObject({ error: { code: 'context_length_exceeded' } });
+    expect(ctx.state.aiApiLlmBilling).toBeUndefined();
+    expect(getProviderCreateCount()).toBe(0);
+    expect(model.invoke).not.toHaveBeenCalled();
+  });
+
+  it('does not emit a usage-only chunk when the provider omits usage metadata', async () => {
+    const { ctx, writes } = createStreamingContext({ content: 'Silent' }, { include_usage: true });
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    const dataLines = writes.filter((line) => line.startsWith('data: '));
+    const doneFrame = dataLines.find((line) => line.includes('[DONE]'));
+    expect(doneFrame).toBeDefined();
+    const frames = dataLines.filter((line) => !line.includes('[DONE]')).map((line) => JSON.parse(line.slice(6)));
+    const doneIndex = frames.length;
+
+    const precedingChunk = frames[doneIndex - 1];
+    expect(precedingChunk.choices[0].finish_reason).toBe('stop');
+    expect(precedingChunk.usage).toBeNull();
+    expect(ctx.state.aiApiUsageResult).toMatchObject({ source: 'unavailable' });
+  });
 });

package/src/server/__tests__/user-permissions-resource.test.ts

@@ -0,0 +1,66 @@
+/**
+ * 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, Next } from '@nocobase/actions';
+import { describe, expect, it, vi } from 'vitest';
+import aiApiUserPermissionsResource from '../resource/ai-api-user-permissions';
+
+type ActionHandler = (ctx: Context, next: Next) => Promise<void>;
+
+const listUsers = aiApiUserPermissionsResource.actions?.listUsers as ActionHandler;
+
+function model(values: Record<string, unknown>) {
+  return { get: (key: string) => values[key] };
+}
+
+describe('aiApiUserPermissions:listUsers', () => {
+  it('returns the canonical rows shape for NocoBase data wrapping', async () => {
+    const user = model({
+      id: 7,
+      nickname: 'Ada',
+      username: 'ada',
+      email: 'ada@example.com',
+      password: 'must-not-leak',
+    });
+    const findAndCount = vi.fn(async () => [[user], 1] as const);
+    const ctx = {
+      action: { params: {} },
+      db: { getRepository: () => ({ findAndCount }) },
+    } as unknown as Context;
+
+    await listUsers(ctx, async () => undefined);
+
+    expect(ctx.body).toEqual({
+      rows: [{ id: 7, nickname: 'Ada', username: 'ada', email: 'ada@example.com' }],
+      count: 1,
+      page: 1,
+      pageSize: 50,
+    });
+    expect(findAndCount).toHaveBeenCalledWith(
+      expect.objectContaining({ fields: ['id', 'nickname', 'username', 'email'] }),
+    );
+  });
+
+  it('excludes users that already hold a grant when requested', async () => {
+    const findAndCount = vi.fn(async () => [[], 0] as const);
+    const ctx = {
+      action: { params: { excludeGranted: true } },
+      db: {
+        getRepository: (name: string) =>
+          name === 'aiApiUserPermissions'
+            ? { find: vi.fn(async () => [model({ userId: 3 }), model({ userId: 8 })]) }
+            : { findAndCount },
+      },
+    } as unknown as Context;
+
+    await listUsers(ctx, async () => undefined);
+
+    expect(findAndCount).toHaveBeenCalledWith(expect.objectContaining({ filter: { id: { $notIn: [3, 8] } } }));
+  });
+});