plugin-ai-api 1.0.25 → 1.1.0

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

@@ -10,60 +10,61 @@
 import { Context } from '@nocobase/actions';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { PluginAiApiServer } from '../plugin';
-import { invalidateUserPermissionCache, resolveUserAccessScope } from '../utils/user-permissions';
+import { invalidateGroupAccessCache, resolveUserAccessScope } from '../utils/user-permissions';
 
 /**
- * Covers the cross-node half of permission revocation.
+ * Covers the cross-node half of group-access invalidation.
  *
  * 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);
+function row(values: Record<string, unknown>) {
+  return { get: (key?: string) => (key === undefined ? values : values[key]) };
+}
+
+function mockContext(userId: number, groupValues: Record<string, unknown>) {
+  const group = row({ allowedLlmServices: [], allowAllModels: true, allowedModels: [], ...groupValues });
+  const findOne = vi.fn(async () => row({ group }));
   const ctx = {
     state: { currentUser: { id: userId } },
-    db: { getRepository: () => ({ findOne }) },
+    db: {
+      getRepository: (name: string) =>
+        name === 'aiApiGroupMembers' ? { findOne } : { findOne: async () => null, create: vi.fn() },
+    },
     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();
+  invalidateGroupAccessCache();
 });
 
-describe('cross-node permission invalidation', () => {
+describe('cross-node group access 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'] }));
+    const { ctx } = mockContext(1, { id: 5, allowedLlmServices: ['openai'] });
 
-    await resolveUserAccessScope(ctx);
-    await resolveUserAccessScope(ctx);
-    expect(findOne).toHaveBeenCalledTimes(1);
+    const before = await resolveUserAccessScope(ctx);
+    expect(await resolveUserAccessScope(ctx)).toBe(before);
 
-    await plugin.handleSyncMessage({ type: 'invalidateUserPermissions', userId: 1 });
+    await plugin.handleSyncMessage({ type: 'invalidateGroupAccess', groupId: 5 });
 
-    await resolveUserAccessScope(ctx);
-    expect(findOne).toHaveBeenCalledTimes(2);
+    expect(await resolveUserAccessScope(ctx)).not.toBe(before);
   });
 
-  it('ignores unrelated message types and other users', async () => {
+  it('ignores unrelated message types and other groups', async () => {
     const plugin = Object.create(PluginAiApiServer.prototype) as PluginAiApiServer;
-    const { ctx, findOne } = mockContext(1, row({ allowedLlmServices: ['openai'] }));
-    await resolveUserAccessScope(ctx);
+    const { ctx } = mockContext(1, { id: 5, allowedLlmServices: ['openai'] });
+    const before = await resolveUserAccessScope(ctx);
 
-    await plugin.handleSyncMessage({ type: 'somethingElse', userId: 1 });
-    await plugin.handleSyncMessage({ type: 'invalidateUserPermissions', userId: 2 });
+    await plugin.handleSyncMessage({ type: 'somethingElse', groupId: 5 });
+    await plugin.handleSyncMessage({ type: 'invalidateGroupAccess', groupId: 6 });
 
-    await resolveUserAccessScope(ctx);
-    expect(findOne).toHaveBeenCalledTimes(1);
+    expect(await resolveUserAccessScope(ctx)).toBe(before);
   });
 
   it('tolerates a malformed message instead of throwing into the subscriber', async () => {
@@ -77,18 +78,16 @@ describe('cross-node permission invalidation', () => {
     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);
+    const { ctx } = mockContext(1, { id: 5, allowedLlmServices: ['openai'] });
+    const before = await resolveUserAccessScope(ctx);
 
-    // revokeUserPermissions is private; reach it the way the db hook does.
-    (plugin as unknown as { revokeUserPermissions: (id: unknown, tx?: unknown) => void }).revokeUserPermissions(1);
+    // invalidateGroupAccess is private; reach it the way the db hook does.
+    (plugin as unknown as { invalidateGroupAccess: (id: unknown, tx?: unknown) => void }).invalidateGroupAccess(5);
 
     // Local cache cleared without any message coming back to us.
-    await resolveUserAccessScope(ctx);
-    expect(findOne).toHaveBeenCalledTimes(2);
+    expect(await resolveUserAccessScope(ctx)).not.toBe(before);
     expect(sendSyncMessage).toHaveBeenCalledWith(
-      { type: 'invalidateUserPermissions', userId: 1 },
+      { type: 'invalidateGroupAccess', groupId: 5 },
       { transaction: undefined },
     );
   });
@@ -99,11 +98,11 @@ describe('cross-node permission invalidation', () => {
     Object.assign(plugin, { sendSyncMessage });
     const transaction = { id: 'tx-1' };
 
-    (plugin as unknown as { revokeUserPermissions: (id: unknown, tx?: unknown) => void }).revokeUserPermissions(
+    (plugin as unknown as { invalidateGroupAccess: (id: unknown, tx?: unknown) => void }).invalidateGroupAccess(
       7,
       transaction,
     );
 
-    expect(sendSyncMessage).toHaveBeenCalledWith({ type: 'invalidateUserPermissions', userId: 7 }, { transaction });
+    expect(sendSyncMessage).toHaveBeenCalledWith({ type: 'invalidateGroupAccess', groupId: 7 }, { transaction });
   });
 });

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

@@ -0,0 +1,160 @@
+import type { Context } from '@nocobase/actions';
+import { createMockDatabase, type Database } from '@nocobase/database';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { resolveUserGroup, getDefaultGroup } from '../quota-groups';
+
+describe('AI API usage groups', () => {
+  let db: Database;
+
+  beforeEach(async () => {
+    db = await createMockDatabase();
+    db.collection({
+      name: 'aiApiUsageGroups',
+      fields: [
+        { name: 'name', type: 'string' },
+        { name: 'isDefault', type: 'boolean' },
+        { name: 'quotaMode', type: 'string' },
+        { name: 'rateLimitPerMinute', type: 'integer' },
+        { name: 'enabled', type: 'boolean' },
+        { name: 'periodType', type: 'string' },
+        { name: 'timezone', type: 'string' },
+        { name: 'requestLimit', type: 'bigInt' },
+        { name: 'totalTokenLimit', type: 'bigInt' },
+        { name: 'costLimit', type: 'decimal', precision: 20, scale: 8 },
+        { name: 'currency', type: 'string' },
+        { name: 'rejectUnpricedModel', type: 'boolean' },
+        { name: 'missingUsageBehavior', type: 'string' },
+        { name: 'contextOverflowBehavior', type: 'string' },
+        { name: 'allowedLlmServices', type: 'json' },
+        { name: 'allowAllModels', type: 'boolean' },
+        { name: 'allowedModels', type: 'json' },
+      ],
+    });
+    db.collection({
+      name: 'aiApiGroupMembers',
+      fields: [
+        { name: 'groupId', type: 'bigInt' },
+        {
+          name: 'group',
+          type: 'belongsTo',
+          target: 'aiApiUsageGroups',
+          targetKey: 'id',
+          foreignKey: 'groupId',
+        },
+        { name: 'userId', type: 'bigInt' },
+      ],
+      indexes: [{ fields: ['userId'], unique: true }],
+    });
+    await db.sync({ force: true });
+  });
+
+  afterEach(async () => {
+    await db.close();
+  });
+
+  function context(): Context {
+    return { db, request: {}, state: {} } as unknown as Context;
+  }
+
+  it('creates the default group lazily with open model access', async () => {
+    const group = await getDefaultGroup(context());
+    expect(group.name).toBe('Default');
+    expect(group.isDefault).toBe(true);
+    expect(group.quotaMode).toBe('per_user');
+    // The default group must never lock everyone out: empty lists mean "no narrowing".
+    expect(group.allowedLlmServices).toEqual([]);
+    expect(group.allowAllModels).toBe(true);
+    expect(group.allowedModels).toEqual([]);
+
+    const second = await getDefaultGroup(context());
+    expect(second.id).toBe(group.id);
+  });
+
+  it('resolves an unassigned user to the default group', async () => {
+    const group = await resolveUserGroup(context(), 99);
+    expect(group.name).toBe('Default');
+    expect(group.isDefault).toBe(true);
+  });
+
+  it('resolves an assigned user to their explicit group', async () => {
+    const custom = await db.getRepository('aiApiUsageGroups').create({
+      values: {
+        name: 'Pro',
+        isDefault: false,
+        quotaMode: 'share',
+        rateLimitPerMinute: 120,
+        enabled: true,
+        periodType: 'monthly',
+        timezone: 'UTC',
+        currency: 'USD',
+        rejectUnpricedModel: true,
+        missingUsageBehavior: 'use_reserved',
+        contextOverflowBehavior: 'reject',
+      },
+    });
+    await db.getRepository('aiApiGroupMembers').create({
+      values: { groupId: custom.get('id'), userId: 42 },
+    });
+
+    const group = await resolveUserGroup(context(), 42);
+    expect(group.name).toBe('Pro');
+    expect(group.quotaMode).toBe('share');
+  });
+
+  it('carries the model access fields through group resolution', async () => {
+    const custom = await db.getRepository('aiApiUsageGroups').create({
+      values: {
+        name: 'Restricted',
+        isDefault: false,
+        quotaMode: 'per_user',
+        rateLimitPerMinute: 60,
+        enabled: false,
+        periodType: 'monthly',
+        timezone: 'UTC',
+        currency: 'USD',
+        rejectUnpricedModel: true,
+        missingUsageBehavior: 'use_reserved',
+        contextOverflowBehavior: 'reject',
+        allowedLlmServices: ['svc'],
+        allowAllModels: false,
+        allowedModels: ['svc/model-a'],
+      },
+    });
+    await db.getRepository('aiApiGroupMembers').create({
+      values: { groupId: custom.get('id'), userId: 43 },
+    });
+
+    const group = await resolveUserGroup(context(), 43);
+    expect(group.allowedLlmServices).toEqual(['svc']);
+    expect(group.allowAllModels).toBe(false);
+    expect(group.allowedModels).toEqual(['svc/model-a']);
+  });
+
+  it('drops non-string entries from the access lists', async () => {
+    const custom = await db.getRepository('aiApiUsageGroups').create({
+      values: {
+        name: 'Messy',
+        isDefault: false,
+        quotaMode: 'per_user',
+        rateLimitPerMinute: 60,
+        enabled: false,
+        periodType: 'monthly',
+        timezone: 'UTC',
+        currency: 'USD',
+        rejectUnpricedModel: true,
+        missingUsageBehavior: 'use_reserved',
+        contextOverflowBehavior: 'reject',
+        allowedLlmServices: ['svc', null, 42],
+        allowAllModels: false,
+        allowedModels: [{ k: 1 }, 'svc/model-a'],
+      },
+    });
+    await db.getRepository('aiApiGroupMembers').create({
+      values: { groupId: custom.get('id'), userId: 44 },
+    });
+
+    const group = await resolveUserGroup(context(), 44);
+    expect(group.allowedLlmServices).toEqual(['svc']);
+    expect(group.allowedModels).toEqual(['svc/model-a']);
+  });
+});

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

@@ -10,6 +10,7 @@ describe('AI API usage monitor summary', () => {
       inputTokens: '100',
       outputTokens: '25',
       totalTokens: '125',
+      promptCacheTokens: '40',
     });
     const findAll = vi.fn().mockResolvedValue([
       { currency: 'USD', totalCost: '0.12500000' },
@@ -53,6 +54,7 @@ describe('AI API usage monitor summary', () => {
       inputTokens: 100,
       outputTokens: 25,
       totalTokens: 125,
+      promptCacheTokens: 40,
       costsByCurrency: [
         { currency: 'USD', totalCost: '0.12500000' },
         { currency: 'EUR', totalCost: '0.05000000' },

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

@@ -15,7 +15,15 @@ interface ModelResult {
   usage_metadata?: Record<string, unknown>;
 }
 
-function createContext(result: ModelResult) {
+interface RepositoryOverride {
+  findOne: () => unknown;
+}
+
+function createContext(
+  result: ModelResult,
+  metadataGet?: (key: string) => unknown,
+  repositories?: Record<string, RepositoryOverride>,
+) {
   const model = {
     invoke: vi.fn().mockResolvedValue(result),
     modelKwargs: {},
@@ -28,6 +36,10 @@ function createContext(result: ModelResult) {
     }
   }
 
+  const getMetadataValue =
+    metadataGet ??
+    ((key: string) => (key === 'contextWindow' ? 128_000 : key === 'maxCompletionTokens' ? 16_384 : true));
+
   const ctx = {
     app: {
       pm: {
@@ -40,13 +52,21 @@ function createContext(result: ModelResult) {
     },
     db: {
       getRepository: vi.fn((name: string) => {
+        if (repositories?.[name]) {
+          return repositories[name];
+        }
         if (name === 'aiApiModelMetadata') {
           return {
             findOne: vi.fn().mockResolvedValue({
-              get: (key: string) => (key === 'contextWindow' ? 128_000 : key === 'maxCompletionTokens' ? 16_384 : true),
+              get: getMetadataValue,
             }),
           };
         }
+        if (name === 'aiApiUsageGroups') {
+          return {
+            findOne: vi.fn().mockResolvedValue({ id: 1, name: 'Default', isDefault: true, allowAllModels: true }),
+          };
+        }
         return { findOne: vi.fn().mockResolvedValue(null) };
       }),
     },
@@ -160,6 +180,11 @@ function createStreamingContext(
             }),
           };
         }
+        if (name === 'aiApiUsageGroups') {
+          return {
+            findOne: vi.fn().mockResolvedValue({ id: 1, name: 'Default', isDefault: true, allowAllModels: true }),
+          };
+        }
         return { findOne: vi.fn().mockResolvedValue(null) };
       }),
     },
@@ -182,6 +207,15 @@ function createStreamingContext(
   return { ctx, model, writes };
 }
 
+const metadataWithSystemPrompt = (key: string) =>
+  key === 'contextWindow'
+    ? 128_000
+    : key === 'maxCompletionTokens'
+      ? 16_384
+      : key === 'systemPrompt'
+        ? 'You are the initial prompt.'
+        : true;
+
 describe('AI API chat usage collection', () => {
   beforeEach(() => {
     vi.mocked(resolveModelString).mockResolvedValue({
@@ -343,6 +377,91 @@ describe('AI API chat usage collection', () => {
     });
   });
 
+  it('persists cached prompt tokens from streaming chat completions', async () => {
+    const { ctx, writes } = createStreamingContext(
+      {
+        content: 'Hi',
+        usage_metadata: {
+          input_tokens: 100,
+          output_tokens: 50,
+          total_tokens: 150,
+          input_token_details: { cache_read: 80 },
+        },
+      },
+      { include_usage: true },
+    );
+
+    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(usageChunk.usage).toEqual({
+      prompt_tokens: 100,
+      completion_tokens: 50,
+      total_tokens: 150,
+      prompt_tokens_details: { cached_tokens: 80 },
+    });
+    expect(ctx.state.aiApiUsageResult).toMatchObject({
+      source: 'provider',
+      usage: { prompt_tokens: 100, completion_tokens: 50, total_tokens: 150, prompt_cache_tokens: 80 },
+    });
+  });
+
+  it('falls back to the usage chunk response_metadata in streaming chat completions', async () => {
+    const { ctx } = createStreamingContext(
+      {
+        content: 'Hi',
+        usage_metadata: { input_tokens: 12, output_tokens: 6, total_tokens: 18 },
+        response_metadata: { usage: { prompt_tokens_details: { cached_tokens: 9 } } },
+      },
+      { include_usage: true },
+    );
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.state.aiApiUsageResult).toMatchObject({
+      source: 'provider',
+      usage: { prompt_tokens: 12, completion_tokens: 6, total_tokens: 18, prompt_cache_tokens: 9 },
+    });
+  });
+
+  it('persists cached prompt tokens from streaming legacy completions', async () => {
+    const { ctx, writes } = createStreamingContext(
+      {
+        content: 'Hi',
+        usage_metadata: {
+          input_tokens: 30,
+          output_tokens: 10,
+          total_tokens: 40,
+          input_token_details: { cache_read: 25 },
+        },
+      },
+      { include_usage: true },
+      { 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(usageChunk.usage).toEqual({
+      prompt_tokens: 30,
+      completion_tokens: 10,
+      total_tokens: 40,
+      prompt_tokens_details: { cached_tokens: 25 },
+    });
+    expect(ctx.state.aiApiUsageResult).toMatchObject({
+      source: 'provider',
+      usage: { prompt_cache_tokens: 25 },
+    });
+  });
+
   it('always emits chat usage and forces provider collection when include_usage is false', async () => {
     const { ctx, model, writes } = createStreamingContext(
       {
@@ -434,6 +553,82 @@ describe('AI API chat usage collection', () => {
     expect(model.invoke).not.toHaveBeenCalled();
   });
 
+  it('prepends the initial system prompt from model metadata in chat completions', async () => {
+    const { ctx, model } = createContext({ content: 'Hello back' }, metadataWithSystemPrompt);
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(200);
+    const invokedMessages = model.invoke.mock.calls[0][0] as [string, unknown][];
+    expect(invokedMessages).toEqual([
+      ['system', 'You are the initial prompt.'],
+      ['user', 'Hello'],
+    ]);
+  });
+
+  it('keeps the client system prompt after the initial system prompt in chat completions', async () => {
+    const { ctx, model } = createContext({ content: 'Hello back' }, metadataWithSystemPrompt);
+    (ctx.request.body as Record<string, unknown>).messages = [
+      { role: 'system', content: 'Client system prompt' },
+      { role: 'user', content: 'Hello' },
+    ];
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(200);
+    const invokedMessages = model.invoke.mock.calls[0][0] as [string, unknown][];
+    expect(invokedMessages).toEqual([
+      ['system', 'You are the initial prompt.'],
+      ['system', 'Client system prompt'],
+      ['user', 'Hello'],
+    ]);
+  });
+
+  it('prepends the initial system prompt in legacy completions', async () => {
+    const { ctx, model } = createContext({ content: 'Hello back' }, metadataWithSystemPrompt);
+    (ctx.request.body as Record<string, unknown>).prompt = 'Hello';
+
+    await handleCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(200);
+    const invokedMessages = model.invoke.mock.calls[0][0] as [string, unknown][];
+    expect(invokedMessages).toEqual([
+      ['system', 'You are the initial prompt.'],
+      ['human', 'Hello'],
+    ]);
+  });
+
+  it('ignores the default AI Employee prompt in direct LLM chat completions', async () => {
+    const employeeFindOne = vi.fn().mockResolvedValue({ about: 'Employee persona prompt' });
+    const { ctx, model } = createContext({ content: 'Hello back' }, undefined, {
+      aiApiConfig: { findOne: vi.fn().mockResolvedValue({ defaultAiEmployee: 'alice' }) },
+      aiEmployees: { findOne: employeeFindOne },
+    });
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(200);
+    expect(employeeFindOne).not.toHaveBeenCalled();
+    const invokedMessages = model.invoke.mock.calls[0][0] as [string, unknown][];
+    expect(invokedMessages).toEqual([['user', 'Hello']]);
+  });
+
+  it('ignores the default AI Employee prompt in legacy completions', async () => {
+    const employeeFindOne = vi.fn().mockResolvedValue({ about: 'Employee persona prompt' });
+    const { ctx, model } = createContext({ content: 'Hello back' }, undefined, {
+      aiApiConfig: { findOne: vi.fn().mockResolvedValue({ defaultAiEmployee: 'alice' }) },
+      aiEmployees: { findOne: employeeFindOne },
+    });
+    (ctx.request.body as Record<string, unknown>).prompt = 'Hello';
+
+    await handleCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(200);
+    expect(employeeFindOne).not.toHaveBeenCalled();
+    const invokedMessages = model.invoke.mock.calls[0][0] as [string, unknown][];
+    expect(invokedMessages).toEqual([['human', 'Hello']]);
+  });
+
   it('does not emit a usage-only chunk when the provider omits usage metadata', async () => {
     const { ctx, writes } = createStreamingContext({ content: 'Silent' }, { include_usage: true });
 
@@ -483,4 +678,69 @@ describe('AI API chat usage collection', () => {
       expect.objectContaining({ seed: 42, reasoning_effort: 'medium' }),
     );
   });
+
+  it('propagates the provider finish_reason in non-streaming chat completions', async () => {
+    const { ctx } = createContext({
+      content: 'Hello back',
+      response_metadata: { finish_reason: 'length' },
+      usage_metadata: { input_tokens: 8, output_tokens: 3, total_tokens: 11 },
+    });
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(200);
+    expect((ctx.body as { choices: Array<{ finish_reason: string }> }).choices[0].finish_reason).toBe('length');
+  });
+
+  it('propagates the provider finish_reason in streaming chat completions', async () => {
+    const { ctx, writes } = createStreamingContext(
+      {
+        content: 'Hi',
+        response_metadata: { finish_reason: 'length' },
+        usage_metadata: { input_tokens: 5, output_tokens: 4, total_tokens: 9 },
+      },
+      { include_usage: true },
+    );
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    const frames = writes
+      .filter((line) => line.startsWith('data: ') && !line.includes('[DONE]'))
+      .map((line) => JSON.parse(line.slice(6)));
+    const finishChunk = frames[frames.length - 2];
+    expect(finishChunk.choices[0].finish_reason).toBe('length');
+  });
+
+  it('propagates the provider finish_reason in legacy completions', async () => {
+    const { ctx } = createContext({
+      content: 'Hello back',
+      response_metadata: { finish_reason: 'length' },
+    });
+    (ctx.request.body as Record<string, unknown>).prompt = 'Hello';
+
+    await handleCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(200);
+    expect((ctx.body as { choices: Array<{ finish_reason: string }> }).choices[0].finish_reason).toBe('length');
+  });
+
+  it('propagates the provider finish_reason in streaming legacy completions', async () => {
+    const { ctx, writes } = createStreamingContext(
+      {
+        content: 'Hi',
+        response_metadata: { finish_reason: 'length' },
+        usage_metadata: { input_tokens: 2, output_tokens: 5, total_tokens: 7 },
+      },
+      { include_usage: true },
+      { 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 finishChunk = frames[frames.length - 2];
+    expect(finishChunk.choices[0].finish_reason).toBe('length');
+  });
 });

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

@@ -54,6 +54,44 @@ describe('AI API usage normalization', () => {
     });
   });
 
+  it('is idempotent so streaming double-normalization keeps prompt_cache_tokens', () => {
+    const streamChunkUsage = {
+      input_tokens: 100,
+      output_tokens: 50,
+      total_tokens: 150,
+      input_token_details: { cache_read: 80 },
+    };
+    const firstPass = normalizeUsage(streamChunkUsage);
+    expect(firstPass?.prompt_cache_tokens).toBe(80);
+
+    // Streaming routes normalize the chunk once, then setAiApiUsageResult
+    // normalizes the result again — the extracted value must survive.
+    expect(normalizeUsage(firstPass)).toEqual(firstPass);
+  });
+
+  it('keeps prompt_cache_tokens when setAiApiUsageResult receives pre-normalized usage', () => {
+    const ctx = createContext();
+    const preNormalized = normalizeUsage({
+      input_tokens: 100,
+      output_tokens: 50,
+      total_tokens: 150,
+      input_token_details: { cache_read: 80 },
+    });
+
+    const usage = setAiApiUsageResult(ctx, preNormalized, { gatewayResponseId: 'gateway-stream-id' });
+
+    expect(usage).toEqual({
+      prompt_tokens: 100,
+      completion_tokens: 50,
+      total_tokens: 150,
+      prompt_cache_tokens: 80,
+    });
+    expect(ctx.state.aiApiUsageResult).toMatchObject({
+      source: 'provider',
+      usage: { prompt_cache_tokens: 80 },
+    });
+  });
+
   it('rejects synthetic or invalid usage values', () => {
     expect(normalizeUsage(undefined)).toBeUndefined();
     expect(normalizeUsage({ prompt_tokens: '12', completion_tokens: -1, total_tokens: Number.NaN })).toBeUndefined();