plugin-ai-api 1.0.24 → 1.0.28

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({
@@ -205,6 +239,7 @@ describe('AI API chat usage collection', () => {
       prompt_tokens: 0,
       completion_tokens: 0,
       total_tokens: 0,
+      prompt_tokens_details: { cached_tokens: null },
     });
     expect(ctx.state.aiApiUsageResult).toMatchObject({
       source: 'unavailable',
@@ -229,6 +264,7 @@ describe('AI API chat usage collection', () => {
       prompt_tokens: 8,
       completion_tokens: 3,
       total_tokens: 11,
+      prompt_tokens_details: { cached_tokens: null },
     });
     expect(ctx.state.aiApiUsageResult).toMatchObject({
       source: 'provider',
@@ -238,6 +274,73 @@ describe('AI API chat usage collection', () => {
     });
   });
 
+  it('extracts cached prompt tokens from OpenAI-style usage metadata', async () => {
+    const { ctx } = createContext({
+      content: 'Hello back',
+      usage_metadata: {
+        input_tokens: 8,
+        output_tokens: 3,
+        total_tokens: 11,
+        prompt_tokens_details: { cached_tokens: 7 },
+      },
+    });
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect((ctx.body as { usage: object }).usage).toEqual({
+      prompt_tokens: 8,
+      completion_tokens: 3,
+      total_tokens: 11,
+      prompt_tokens_details: { cached_tokens: 7 },
+    });
+    expect(ctx.state.aiApiUsageResult).toMatchObject({
+      source: 'provider',
+      usage: { prompt_tokens: 8, completion_tokens: 3, total_tokens: 11, prompt_cache_tokens: 7 },
+    });
+  });
+
+  it('falls back to response_metadata when usage metadata omits cached tokens', async () => {
+    const { ctx } = createContext({
+      content: 'Hello back',
+      usage_metadata: { input_tokens: 8, output_tokens: 3, total_tokens: 11 },
+      response_metadata: { usage: { prompt_tokens_details: { cached_tokens: 5 } } },
+    });
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect((ctx.body as { usage: object }).usage).toEqual({
+      prompt_tokens: 8,
+      completion_tokens: 3,
+      total_tokens: 11,
+      prompt_tokens_details: { cached_tokens: 5 },
+    });
+    expect(ctx.state.aiApiUsageResult).toMatchObject({
+      source: 'provider',
+      usage: { prompt_cache_tokens: 5 },
+    });
+  });
+
+  it('extracts cached prompt tokens from LangChain-style input_token_details', async () => {
+    const { ctx } = createContext({
+      content: 'Hello back',
+      usage_metadata: {
+        input_tokens: 8,
+        output_tokens: 3,
+        total_tokens: 11,
+        input_token_details: { cache_read: 4 },
+      },
+    });
+
+    await handleChatCompletions(ctx, {} as PluginAiApiServer);
+
+    expect((ctx.body as { usage: object }).usage).toEqual({
+      prompt_tokens: 8,
+      completion_tokens: 3,
+      total_tokens: 11,
+      prompt_tokens_details: { cached_tokens: 4 },
+    });
+  });
+
   it('emits a usage-only chunk immediately before [DONE] for streaming chat completions', async () => {
     const { ctx, writes } = createStreamingContext(
       {
@@ -261,7 +364,12 @@ describe('AI API chat usage collection', () => {
 
     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.usage).toEqual({
+      prompt_tokens: 5,
+      completion_tokens: 4,
+      total_tokens: 9,
+      prompt_tokens_details: { cached_tokens: null },
+    });
     expect(usageChunk).toHaveProperty('usage.prompt_tokens', 5);
     expect(ctx.state.aiApiUsageResult).toMatchObject({
       source: 'provider',
@@ -269,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(
       {
@@ -287,7 +480,12 @@ describe('AI API chat usage collection', () => {
 
     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(usageChunk.usage).toEqual({
+      prompt_tokens: 5,
+      completion_tokens: 4,
+      total_tokens: 9,
+      prompt_tokens_details: { cached_tokens: null },
+    });
     expect(model.stream).toHaveBeenCalledWith(
       expect.anything(),
       expect.objectContaining({
@@ -316,7 +514,12 @@ describe('AI API chat usage collection', () => {
     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(usageChunk.usage).toEqual({
+      prompt_tokens: 2,
+      completion_tokens: 5,
+      total_tokens: 7,
+      prompt_tokens_details: { cached_tokens: null },
+    });
     expect(model.stream).toHaveBeenCalledWith(
       expect.anything(),
       expect.objectContaining({
@@ -350,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 });
 
@@ -366,4 +645,102 @@ describe('AI API chat usage collection', () => {
     expect(precedingChunk.usage).toBeNull();
     expect(ctx.state.aiApiUsageResult).toMatchObject({ source: 'unavailable' });
   });
+
+  it('forwards passthrough provider parameters in legacy non-stream completions', async () => {
+    const { ctx, model } = createContext({
+      content: 'Hello back',
+      usage_metadata: { input_tokens: 2, output_tokens: 5, total_tokens: 7 },
+    });
+    (ctx.request.body as Record<string, unknown>).prompt = 'Hello';
+    (ctx.request.body as Record<string, unknown>).seed = 42;
+    (ctx.request.body as Record<string, unknown>).reasoning_effort = 'medium';
+
+    await handleCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(ctx.status).toBe(200);
+    expect(model.invoke).toHaveBeenCalledWith(
+      expect.anything(),
+      expect.objectContaining({ seed: 42, reasoning_effort: 'medium' }),
+    );
+  });
+
+  it('forwards passthrough provider parameters in legacy streaming completions', async () => {
+    const { ctx, model } = createStreamingContext(
+      { content: 'Hi', usage_metadata: { input_tokens: 2, output_tokens: 5, total_tokens: 7 } },
+      { include_usage: false, include_obfuscation: false },
+      { prompt: 'Hello', seed: 42, reasoning_effort: 'medium' },
+    );
+
+    await handleCompletions(ctx, {} as PluginAiApiServer);
+
+    expect(model.stream).toHaveBeenCalledWith(
+      expect.anything(),
+      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

@@ -24,6 +24,7 @@ describe('AI API usage normalization', () => {
       prompt_tokens: 12,
       completion_tokens: 5,
       total_tokens: 17,
+      prompt_cache_tokens: null,
     });
   });
 
@@ -32,6 +33,62 @@ describe('AI API usage normalization', () => {
       prompt_tokens: 0,
       completion_tokens: 0,
       total_tokens: 0,
+      prompt_cache_tokens: null,
+    });
+  });
+
+  it('extracts prompt_cache_tokens when present in various provider formats', () => {
+    expect(
+      normalizeUsage({ prompt_tokens: 10, completion_tokens: 5, prompt_tokens_details: { cached_tokens: 8 } }),
+    ).toEqual({
+      prompt_tokens: 10,
+      completion_tokens: 5,
+      total_tokens: 15,
+      prompt_cache_tokens: 8,
+    });
+    expect(normalizeUsage({ input_tokens: 20, output_tokens: 10, input_token_details: { cache_read: 15 } })).toEqual({
+      prompt_tokens: 20,
+      completion_tokens: 10,
+      total_tokens: 30,
+      prompt_cache_tokens: 15,
+    });
+  });
+
+  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 },
     });
   });