plugin-ai-api 1.0.20 → 1.0.23

package/src/server/routes/completions.ts

@@ -8,12 +8,25 @@
  */
 
 import { Context } from '@nocobase/actions';
-import { generateCompletionId, toOpenAIError, formatSSE, formatSSEDone } from '../utils/openai-format';
+import {
+  generateCompletionId,
+  toOpenAIError,
+  formatSSE,
+  formatSSEDone,
+  toOpenAIUsageChunk,
+} from '../utils/openai-format';
 import { resolveModelString } from '../utils/resolve-service';
-import { createRequestAbortController, isStreamingRequested, writeResponse } from '../utils/streaming';
+import { enforceModelAccess } from '../utils/user-permissions';
+import {
+  createRequestAbortController,
+  isClientDisconnected,
+  isStreamingRequested,
+  writeResponse,
+} from '../utils/streaming';
 import { extractProviderRequestId, normalizeUsage, setAiApiUsageResult, type Usage } from '../usage';
 import type PluginAiApiServer from '../plugin';
 import { AiApiQuotaError, markLlmProviderAttempted, prepareLlmBilling } from '../billing';
+import { markAiApiFirstProviderOutput } from '../utils/app-observability';
 
 /**
  * POST /api/ai-llm/v1/completions
@@ -87,22 +100,10 @@ export async function handleCompletions(ctx: Context, plugin: PluginAiApiServer)
       return;
     }
 
-    // ─── Check whitelist ───
+    // ─── Check whitelist (global config ∩ per-user grant) ───
     const config = await ctx.db.getRepository('aiApiConfig').findOne();
-    if (config?.enabledLlmServices?.length) {
-      const serviceName = service.name;
-      const serviceTitle = service.title;
-      const isAllowed = config.enabledLlmServices.some((s: string) => s === serviceName || s === serviceTitle);
-      if (!isAllowed) {
-        ctx.status = 403;
-        ctx.body = toOpenAIError(
-          403,
-          `LLM service '${service.title || service.name}' is not enabled for API access`,
-          'invalid_request_error',
-          'model_not_available',
-        );
-        return;
-      }
+    if (!(await enforceModelAccess(ctx, config?.enabledLlmServices, service, modelId))) {
+      return;
     }
 
     // ─── Create LLM provider instance ───
@@ -160,7 +161,14 @@ export async function handleCompletions(ctx: Context, plugin: PluginAiApiServer)
     markLlmProviderAttempted(ctx);
 
     if (stream) {
-      await handleStreamingTextCompletion(ctx, chatModel, langchainMessages, completionId, body.model);
+      await handleStreamingTextCompletion(
+        ctx,
+        chatModel,
+        langchainMessages,
+        completionId,
+        body.model,
+        body.stream_options,
+      );
     } else {
       await handleNonStreamingTextCompletion(ctx, chatModel, langchainMessages, completionId, body.model);
     }
@@ -231,6 +239,7 @@ async function handleStreamingTextCompletion(
   messages: [string, string][],
   completionId: string,
   modelName: string,
+  streamOptions: Record<string, unknown> | undefined,
 ) {
   ctx.set({
     'Content-Type': 'text/event-stream',
@@ -244,7 +253,10 @@ async function handleStreamingTextCompletion(
   let usage: Usage | undefined;
   let providerRequestId: string | undefined;
   try {
-    const stream = await chatModel.stream(messages, { signal: requestAbort.signal });
+    const stream = await chatModel.stream(messages, {
+      stream_options: { ...streamOptions, include_usage: true },
+      signal: requestAbort.signal,
+    });
 
     for await (const chunk of stream) {
       if (requestAbort.signal.aborted) throw requestAbort.signal.reason;
@@ -257,6 +269,7 @@ async function handleStreamingTextCompletion(
       }
 
       if (text) {
+        markAiApiFirstProviderOutput(ctx);
         await writeResponse(
           ctx,
           formatSSE({
@@ -273,6 +286,7 @@ async function handleStreamingTextCompletion(
                 finish_reason: null,
               },
             ],
+            usage: null,
           }),
         );
       }
@@ -299,13 +313,29 @@ async function handleStreamingTextCompletion(
             finish_reason: 'stop',
           },
         ],
+        usage: null,
       }),
     );
 
+    if (usage) {
+      await writeResponse(
+        ctx,
+        formatSSE(
+          toOpenAIUsageChunk({
+            id: completionId,
+            model: modelName,
+            object: 'text_completion',
+            usage,
+          }),
+        ),
+      );
+    }
+
     await writeResponse(ctx, formatSSEDone());
     setAiApiUsageResult(ctx, usage, { gatewayResponseId: completionId, providerRequestId });
     ctx.state.aiApiStreamResult = { succeeded: true, id: completionId };
   } catch (err) {
+    const cancelled = isClientDisconnected(ctx, err);
     ctx.log.error('AI API completions streaming error:', err);
     if (!ctx.res.destroyed && !ctx.res.writableEnded) {
       await writeResponse(
@@ -319,7 +349,11 @@ async function handleStreamingTextCompletion(
       );
     }
     setAiApiUsageResult(ctx, usage, { gatewayResponseId: completionId, providerRequestId });
-    ctx.state.aiApiStreamResult = { succeeded: false, id: completionId, errorCode: 'stream_error' };
+    ctx.state.aiApiStreamResult = {
+      succeeded: false,
+      id: completionId,
+      errorCode: cancelled ? 'client_disconnected' : '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 { enforceModelAccess } from '../utils/user-permissions';
 import { setAiApiUsageUnavailable } from '../usage';
 import type PluginAiApiServer from '../plugin';
 
@@ -112,24 +113,18 @@ export async function handleEmbeddings(ctx: Context, plugin: PluginAiApiServer)
     return;
   }
 
-  // ─── Check service whitelist ──────────────────────────────────────────────
+  // ─── Check service whitelist (global config ∩ per-user grant) ─────────────
+  // A config read failure falls open on the global list, but the per-user grant is
+  // still enforced: an explicit deny must never be bypassed by an unreadable config.
+  let globalEnabledServices: unknown = [];
   try {
     const config = await ctx.db.getRepository('aiApiConfig').findOne();
-    if (config?.enabledLlmServices?.length) {
-      const allowed = config.enabledLlmServices.some((s: string) => s === service.name || s === service.title);
-      if (!allowed) {
-        ctx.status = 403;
-        ctx.body = toOpenAIError(
-          403,
-          `LLM service '${service.title || service.name}' is not enabled for API access`,
-          'invalid_request_error',
-          'model_not_available',
-        );
-        return;
-      }
-    }
+    globalEnabledServices = config?.enabledLlmServices ?? [];
   } catch {
-    // Config read failure: fail open
+    // Config read failure: fail open on the global whitelist only.
+  }
+  if (!(await enforceModelAccess(ctx, globalEnabledServices, service, modelId))) {
+    return;
   }
 
   // ─── Get embedding provider ───────────────────────────────────────────────