plugin-ai-api 1.0.23 → 1.0.24

package/src/client-v2/pages/UserQuotasPage.tsx

@@ -39,6 +39,7 @@ interface QuotaPolicy {
   currency: string;
   rejectUnpricedModel: boolean;
   missingUsageBehavior: 'allow' | 'use_reserved';
+  contextOverflowBehavior: 'reject' | 'truncate';
 }
 
 export default function UserQuotasPage() {
@@ -85,6 +86,7 @@ export default function UserQuotasPage() {
       currency: 'USD',
       rejectUnpricedModel: true,
       missingUsageBehavior: 'use_reserved',
+      contextOverflowBehavior: 'reject',
     } as QuotaPolicy);
     setOpen(true);
   };
@@ -157,6 +159,14 @@ export default function UserQuotasPage() {
       width: 120,
       render: (value, record) => (value == null ? t('Unlimited') : `${value} ${record.currency}`),
     },
+    {
+      title: t('Context overflow behavior'),
+      dataIndex: 'contextOverflowBehavior',
+      key: 'contextOverflowBehavior',
+      width: 150,
+      render: (value: QuotaPolicy['contextOverflowBehavior']) =>
+        value === 'truncate' ? t('Truncate oldest conversation turns') : t('Reject request'),
+    },
     { title: t('Timezone'), dataIndex: 'timezone', key: 'timezone', width: 140 },
     {
       title: t('Status'),
@@ -248,6 +258,14 @@ export default function UserQuotasPage() {
               ]}
             />
           </Form.Item>
+          <Form.Item name="contextOverflowBehavior" label={t('Context overflow behavior')} rules={[{ required: true }]}>
+            <Select
+              options={[
+                { label: t('Reject request'), value: 'reject' },
+                { label: t('Truncate oldest conversation turns'), value: 'truncate' },
+              ]}
+            />
+          </Form.Item>
           <Form.Item name="enabled" label={t('Enabled')} valuePropName="checked">
             <Switch />
           </Form.Item>

package/src/locale/en-US.json

@@ -53,6 +53,9 @@
   "Missing usage behavior": "Missing usage behavior",
   "Use reserved estimate": "Use reserved estimate",
   "Allow without token charge": "Allow without token charge",
+  "Context overflow behavior": "Context overflow behavior",
+  "Reject request": "Reject request",
+  "Truncate oldest conversation turns": "Truncate oldest conversation turns",
   "Started at": "Started at",
   "Requested model": "Requested model",
   "Resolved service": "Resolved service",

package/src/locale/vi-VN.json

@@ -53,6 +53,9 @@
   "Missing usage behavior": "Xử lý khi thiếu token usage",
   "Use reserved estimate": "Dùng số liệu giữ chỗ để ước tính",
   "Allow without token charge": "Cho phép và không tính token",
+  "Context overflow behavior": "Xử lý khi vượt context",
+  "Reject request": "Từ chối request",
+  "Truncate oldest conversation turns": "Cắt các lượt hội thoại cũ nhất",
   "Started at": "Bắt đầu lúc",
   "Requested model": "Model được yêu cầu",
   "Resolved service": "Service đã resolve",

package/src/locale/zh-CN.json

@@ -53,6 +53,9 @@
   "Missing usage behavior": "缺少用量时的行为",
   "Use reserved estimate": "使用预留估算",
   "Allow without token charge": "允许且不计令牌",
+  "Context overflow behavior": "上下文超限处理",
+  "Reject request": "拒绝请求",
+  "Truncate oldest conversation turns": "截断最早的对话轮次",
   "Started at": "开始时间",
   "Requested model": "请求模型",
   "Resolved service": "解析后的服务",

package/src/server/__tests__/direct-llm-context.test.ts

@@ -0,0 +1,125 @@
+import type { Context } from '@nocobase/actions';
+import { describe, expect, it, vi } from 'vitest';
+import { DirectLlmContextError, prepareDirectLlmContext, type OpenAIMessage } from '../utils/direct-llm-context';
+
+function context({ behavior = 'reject', metadata = { contextWindow: 120, maxCompletionTokens: 40 } } = {}): Context {
+  return {
+    state: { currentUser: { id: 1 } },
+    db: {
+      getRepository: vi.fn((name: string) => {
+        if (name === 'aiApiModelMetadata') {
+          return {
+            findOne: vi.fn().mockResolvedValue({ get: (key: string) => metadata[key as keyof typeof metadata] }),
+          };
+        }
+        if (name === 'aiApiUserQuotaPolicies') {
+          return {
+            findOne: vi
+              .fn()
+              .mockResolvedValue({ get: (key: string) => (key === 'contextOverflowBehavior' ? behavior : undefined) }),
+          };
+        }
+        return { findOne: vi.fn() };
+      }),
+    },
+  } as unknown as Context;
+}
+
+function request(messages: OpenAIMessage[], tools?: unknown) {
+  return {
+    serviceName: 'test-service',
+    modelId: 'test-model',
+    messages,
+    tools,
+  };
+}
+
+describe('direct LLM context preparation', () => {
+  it('uses reject when the user has no enabled policy', async () => {
+    const ctx = context();
+    vi.mocked(ctx.db.getRepository).mockImplementation((name: string) => {
+      if (name === 'aiApiModelMetadata') {
+        return {
+          findOne: vi.fn().mockResolvedValue({ get: (key: string) => (key === 'contextWindow' ? 120 : 40) }),
+        } as never;
+      }
+      return { findOne: vi.fn().mockResolvedValue(null) } as never;
+    });
+
+    await expect(
+      prepareDirectLlmContext(ctx, request([{ role: 'user', content: 'x'.repeat(400) }])),
+    ).rejects.toMatchObject<Partial<DirectLlmContextError>>({ code: 'context_length_exceeded' });
+  });
+
+  it('keeps a request that fits the model input budget', async () => {
+    const messages = [{ role: 'user', content: 'hello' }];
+    const prepared = await prepareDirectLlmContext(context(), request(messages));
+
+    expect(prepared.messages).toBe(messages);
+    expect(prepared.truncated).toBe(false);
+    expect(prepared.inputTokenBudget).toBe(80);
+  });
+
+  it('rejects a requested output limit above model metadata', async () => {
+    await expect(
+      prepareDirectLlmContext(context(), { ...request([{ role: 'user', content: 'hello' }]), maxCompletionTokens: 41 }),
+    ).rejects.toMatchObject<Partial<DirectLlmContextError>>({ code: 'max_completion_tokens_exceeds_model_limit' });
+  });
+
+  it('truncates oldest complete turns while preserving instructions and newest turn', async () => {
+    const messages: OpenAIMessage[] = [
+      { role: 'system', content: 'Always answer safely.' },
+      { role: 'user', content: 'first '.repeat(30) },
+      { role: 'assistant', content: 'first answer '.repeat(20) },
+      { role: 'user', content: 'latest question' },
+    ];
+
+    const prepared = await prepareDirectLlmContext(context({ behavior: 'truncate' }), request(messages));
+
+    expect(prepared.truncated).toBe(true);
+    expect(prepared.messages).toEqual([messages[0], messages[3]]);
+    expect(messages).toHaveLength(4);
+  });
+
+  it('keeps assistant tool calls and their tool responses in the same turn', async () => {
+    const messages: OpenAIMessage[] = [
+      { role: 'user', content: 'old question '.repeat(30) },
+      { role: 'assistant', content: '', tool_calls: [{ id: 'call-old', type: 'function' }] },
+      { role: 'tool', tool_call_id: 'call-old', content: 'old result '.repeat(20) },
+      { role: 'user', content: 'latest question' },
+    ];
+
+    const prepared = await prepareDirectLlmContext(context({ behavior: 'truncate' }), request(messages));
+
+    expect(prepared.messages).toEqual([messages[3]]);
+  });
+
+  it('rejects when the newest turn cannot fit without cutting content', async () => {
+    await expect(
+      prepareDirectLlmContext(context({ behavior: 'truncate' }), request([{ role: 'user', content: 'x'.repeat(400) }])),
+    ).rejects.toMatchObject<Partial<DirectLlmContextError>>({ code: 'context_length_exceeded' });
+  });
+
+  it('rejects image content until a model-specific estimator is available', async () => {
+    await expect(
+      prepareDirectLlmContext(
+        context(),
+        request([
+          { role: 'user', content: [{ type: 'image_url', image_url: { url: 'https://example.test/image.png' } }] },
+        ]),
+      ),
+    ).rejects.toMatchObject<Partial<DirectLlmContextError>>({ code: 'context_estimation_unsupported' });
+  });
+
+  it('counts tool definitions as fixed input overhead', async () => {
+    await expect(
+      prepareDirectLlmContext(
+        context(),
+        request(
+          [{ role: 'user', content: 'hello' }],
+          [{ type: 'function', function: { name: 'large', parameters: { text: 'x'.repeat(400) } } }],
+        ),
+      ),
+    ).rejects.toMatchObject<Partial<DirectLlmContextError>>({ code: 'context_length_exceeded' });
+  });
+});

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

@@ -15,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({
@@ -37,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: {
@@ -50,6 +61,8 @@ function createContext(result: ModelResult): Context {
     state: {},
     set: vi.fn(),
   } as unknown as Context;
+
+  return { ctx, model, getProviderCreateCount: () => providerCreateCount };
 }
 
 class ListenerTarget {
@@ -139,7 +152,16 @@ function createStreamingContext(
       },
     },
     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() },
     req,
@@ -174,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);
 
@@ -195,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 },
@@ -303,6 +325,31 @@ describe('AI API chat usage collection', () => {
     );
   });
 
+  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 });
 

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

@@ -0,0 +1,36 @@
+import type { Model } from '@nocobase/database';
+import { describe, expect, it, vi } from 'vitest';
+
+vi.mock('dayjs', () => ({
+  default: () => ({ tz: vi.fn() }),
+}));
+
+import { validateQuotaPolicy } from '../validation';
+
+function quotaPolicy(values: Record<string, unknown>): Model {
+  return {
+    get: (key: string) => values[key],
+  } as Model;
+}
+
+const requiredPolicy = {
+  periodType: 'monthly',
+  missingUsageBehavior: 'use_reserved',
+  timezone: 'UTC',
+};
+
+describe('AI API quota policy validation', () => {
+  it('uses reject when context overflow behavior is absent', () => {
+    expect(() => validateQuotaPolicy(quotaPolicy(requiredPolicy))).not.toThrow();
+  });
+
+  it.each(['reject', 'truncate'] as const)('accepts %s context overflow behavior', (contextOverflowBehavior) => {
+    expect(() => validateQuotaPolicy(quotaPolicy({ ...requiredPolicy, contextOverflowBehavior }))).not.toThrow();
+  });
+
+  it('rejects an unsupported context overflow behavior', () => {
+    expect(() => validateQuotaPolicy(quotaPolicy({ ...requiredPolicy, contextOverflowBehavior: 'compact' }))).toThrow(
+      'contextOverflowBehavior must be reject or truncate.',
+    );
+  });
+});

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

@@ -22,6 +22,7 @@ export default defineCollection({
     { name: 'currency', type: 'string', allowNull: false, defaultValue: 'USD' },
     { name: 'rejectUnpricedModel', type: 'boolean', defaultValue: true },
     { name: 'missingUsageBehavior', type: 'string', allowNull: false, defaultValue: 'use_reserved' },
+    { name: 'contextOverflowBehavior', type: 'string', allowNull: false, defaultValue: 'reject' },
   ],
   indexes: [
     {

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

@@ -31,6 +31,7 @@ import { enforceModelAccess } from '../utils/user-permissions';
 import { extractProviderRequestId, normalizeUsage, setAiApiUsageResult, type Usage } from '../usage';
 import type PluginAiApiServer from '../plugin';
 import { AiApiQuotaError, markLlmProviderAttempted, prepareLlmBilling } from '../billing';
+import { DirectLlmContextError, prepareDirectLlmContext, type OpenAIMessage } from '../utils/direct-llm-context';
 import { markAiApiFirstProviderOutput } from '../utils/app-observability';
 
 /**
@@ -142,8 +143,6 @@ export async function handleChatCompletions(ctx: Context, plugin: PluginAiApiSer
       return;
     }
 
-    await prepareLlmBilling(ctx, resolved);
-
     const providerRequestParameters = getProviderRequestParameters(body);
     if (stream) {
       const streamOptions = isRecord(body.stream_options) ? body.stream_options : {};
@@ -166,13 +165,6 @@ export async function handleChatCompletions(ctx: Context, plugin: PluginAiApiSer
     if (body.presence_penalty !== undefined) modelOptions.presencePenalty = body.presence_penalty;
     if (body.stop !== undefined) modelOptions.stop = body.stop;
 
-    const Provider = providerMeta.provider;
-    const provider = new Provider({
-      app: ctx.app,
-      serviceOptions: service.options,
-      modelOptions,
-    });
-
     // ─── Build system prompt from AI Employee ───
     let systemPrompt = '';
     if (config?.defaultAiEmployee) {
@@ -197,12 +189,31 @@ export async function handleChatCompletions(ctx: Context, plugin: PluginAiApiSer
     }
 
     // ─── Build messages (inject system prompt if not provided by client) ───
-    const messages = [...body.messages];
+    let messages: OpenAIMessage[] = [...body.messages];
     const hasSystemMessage = messages.some((m: any) => m.role === 'system');
     if (systemPrompt && !hasSystemMessage) {
       messages.unshift({ role: 'system', content: systemPrompt });
     }
 
+    const preparedContext = await prepareDirectLlmContext(ctx, {
+      serviceName: service.name,
+      modelId,
+      messages,
+      tools: body.tools,
+      maxCompletionTokens: body.max_completion_tokens,
+      maxTokens: body.max_tokens,
+    });
+    messages = preparedContext.messages;
+
+    await prepareLlmBilling(ctx, resolved);
+
+    const Provider = providerMeta.provider;
+    const provider = new Provider({
+      app: ctx.app,
+      serviceOptions: service.options,
+      modelOptions,
+    });
+
     // ─── Build message tuples for LangChain model ───
     // LangChain chat models accept [role, content] tuples or BaseMessage objects.
     // We use tuples to avoid importing @langchain/core directly.
@@ -257,15 +268,16 @@ export async function handleChatCompletions(ctx: Context, plugin: PluginAiApiSer
     }
   } catch (err) {
     ctx.log.error('AI API chat completions error:', err);
-    if (!ctx.res.headersSent) {
+    if (!ctx.res?.headersSent) {
       const isQuotaError = err instanceof AiApiQuotaError;
-      ctx.status = isQuotaError ? 429 : 500;
+      const isContextError = err instanceof DirectLlmContextError;
+      ctx.status = isQuotaError ? 429 : isContextError ? 400 : 500;
       if (isQuotaError) ctx.set('X-RateLimit-Reason', err.code);
       ctx.body = toOpenAIError(
         ctx.status,
         getErrorMessage(err, 'Internal server error'),
-        isQuotaError ? 'quota_error' : 'server_error',
-        isQuotaError ? err.code : undefined,
+        isQuotaError ? 'quota_error' : isContextError ? 'invalid_request_error' : 'server_error',
+        isQuotaError || isContextError ? err.code : undefined,
       );
     }
   }

package/src/server/routes/completions.ts

@@ -26,6 +26,7 @@ import {
 import { extractProviderRequestId, normalizeUsage, setAiApiUsageResult, type Usage } from '../usage';
 import type PluginAiApiServer from '../plugin';
 import { AiApiQuotaError, markLlmProviderAttempted, prepareLlmBilling } from '../billing';
+import { DirectLlmContextError, prepareDirectLlmContext, type OpenAIMessage } from '../utils/direct-llm-context';
 import { markAiApiFirstProviderOutput } from '../utils/app-observability';
 
 /**
@@ -114,8 +115,6 @@ export async function handleCompletions(ctx: Context, plugin: PluginAiApiServer)
       return;
     }
 
-    await prepareLlmBilling(ctx, resolved);
-
     const modelOptions: Record<string, any> = {
       model: modelId,
       llmService: service.name,
@@ -126,13 +125,6 @@ export async function handleCompletions(ctx: Context, plugin: PluginAiApiServer)
     if (body.max_tokens !== undefined) modelOptions.maxTokens = body.max_tokens;
     if (body.stop !== undefined) modelOptions.stop = body.stop;
 
-    const Provider = providerMeta.provider;
-    const provider = new Provider({
-      app: ctx.app,
-      serviceOptions: service.options,
-      modelOptions,
-    });
-
     // ─── Convert prompt to message tuple ───
     const prompt =
       typeof body.prompt === 'string'
@@ -142,7 +134,7 @@ export async function handleCompletions(ctx: Context, plugin: PluginAiApiServer)
           : String(body.prompt);
 
     // Inject system prompt from AI Employee if configured
-    const langchainMessages: [string, string][] = [];
+    const messages: OpenAIMessage[] = [];
     if (config?.defaultAiEmployee) {
       const employee = await ctx.db.getRepository('aiEmployees').findOne({
         filter: { username: config.defaultAiEmployee },
@@ -150,11 +142,30 @@ export async function handleCompletions(ctx: Context, plugin: PluginAiApiServer)
       if (employee) {
         const systemPrompt = employee.about || employee.defaultPrompt || '';
         if (systemPrompt) {
-          langchainMessages.push(['system', systemPrompt]);
+          messages.push({ role: 'system', content: systemPrompt });
         }
       }
     }
-    langchainMessages.push(['human', prompt]);
+    messages.push({ role: 'user', content: prompt });
+
+    const preparedContext = await prepareDirectLlmContext(ctx, {
+      serviceName: service.name,
+      modelId,
+      messages,
+      maxTokens: body.max_tokens,
+    });
+    await prepareLlmBilling(ctx, resolved);
+
+    const Provider = providerMeta.provider;
+    const provider = new Provider({
+      app: ctx.app,
+      serviceOptions: service.options,
+      modelOptions,
+    });
+    const langchainMessages = preparedContext.messages.map((message): [string, string] => [
+      message.role === 'user' ? 'human' : message.role,
+      String(message.content ?? ''),
+    ]);
 
     const completionId = generateCompletionId().replace('chatcmpl-', 'cmpl-');
     const chatModel = provider.createModel();
@@ -174,15 +185,16 @@ export async function handleCompletions(ctx: Context, plugin: PluginAiApiServer)
     }
   } catch (err) {
     ctx.log.error('AI API completions error:', err);
-    if (!ctx.res.headersSent) {
+    if (!ctx.res?.headersSent) {
       const isQuotaError = err instanceof AiApiQuotaError;
-      ctx.status = isQuotaError ? 429 : 500;
+      const isContextError = err instanceof DirectLlmContextError;
+      ctx.status = isQuotaError ? 429 : isContextError ? 400 : 500;
       if (isQuotaError) ctx.set('X-RateLimit-Reason', err.code);
       ctx.body = toOpenAIError(
         ctx.status,
         getErrorMessage(err, 'Internal server error'),
-        isQuotaError ? 'quota_error' : 'server_error',
-        isQuotaError ? err.code : undefined,
+        isQuotaError ? 'quota_error' : isContextError ? 'invalid_request_error' : 'server_error',
+        isQuotaError || isContextError ? err.code : undefined,
       );
     }
   }