From 095d9f2196181c532b4cebe61f68b3b8d654d5ce Mon Sep 17 00:00:00 2001
From: voodoohop <t.haferlach@gmail.com>
Date: Tue, 9 Jun 2026 20:06:49 +0300
Subject: [PATCH] fix: preserve Perplexity usage cost in non-strict mode

---
 .../perplexity-ai/chatComplete.test.ts        | 108 ++++++++++++++++++
 src/providers/perplexity-ai/chatComplete.ts   |  41 ++++---
 2 files changed, 133 insertions(+), 16 deletions(-)
 create mode 100644 src/providers/perplexity-ai/chatComplete.test.ts

diff --git a/src/providers/perplexity-ai/chatComplete.test.ts b/src/providers/perplexity-ai/chatComplete.test.ts
new file mode 100644
index 000000000..5852cf43d
--- /dev/null
+++ b/src/providers/perplexity-ai/chatComplete.test.ts
@@ -0,0 +1,108 @@
+import {
+  PerplexityAIChatCompleteResponse,
+  PerplexityAIChatCompleteResponseTransform,
+  PerplexityAIChatCompleteStreamChunkTransform,
+} from './chatComplete';
+import { ChatCompletionResponse } from '../types';
+
+const response = {
+  id: 'test-id',
+  model: 'sonar',
+  object: 'chat.completion',
+  created: 1,
+  citations: ['https://example.com'],
+  choices: [
+    {
+      message: { role: 'assistant', content: 'ok' },
+      delta: { role: 'assistant', content: 'ok' },
+      index: 0,
+      finish_reason: 'stop',
+    },
+  ],
+  usage: {
+    prompt_tokens: 10,
+    completion_tokens: 5,
+    total_tokens: 15,
+    search_context_size: 'low',
+    cost: {
+      request_cost: 0.005,
+      total_cost: 0.00501,
+    },
+  },
+} satisfies PerplexityAIChatCompleteResponse;
+
+describe('PerplexityAIChatCompleteResponseTransform', () => {
+  test('preserves provider usage in non-strict mode', () => {
+    const transformed = PerplexityAIChatCompleteResponseTransform(
+      response,
+      200,
+      new Headers(),
+      false
+    );
+
+    expect(transformed).toMatchObject({
+      citations: ['https://example.com'],
+      usage: {
+        prompt_tokens: 10,
+        completion_tokens: 5,
+        total_tokens: 15,
+        search_context_size: 'low',
+        cost: {
+          request_cost: 0.005,
+        },
+      },
+    });
+  });
+
+  test('strips provider usage in strict mode', () => {
+    const transformed = PerplexityAIChatCompleteResponseTransform(
+      response,
+      200,
+      new Headers(),
+      true
+    ) as ChatCompletionResponse;
+
+    expect(transformed).toMatchObject({
+      usage: {
+        prompt_tokens: 10,
+        completion_tokens: 5,
+        total_tokens: 15,
+      },
+    });
+    expect(transformed).not.toHaveProperty('citations');
+    expect(transformed.usage).not.toHaveProperty('cost');
+    expect(transformed.usage).not.toHaveProperty('search_context_size');
+  });
+});
+
+describe('PerplexityAIChatCompleteStreamChunkTransform', () => {
+  test('strips provider usage in strict mode', () => {
+    const chunk = `data: ${JSON.stringify({
+      ...response,
+      object: 'chat.completion.chunk',
+      choices: [
+        {
+          delta: { role: 'assistant', content: '' },
+          index: 0,
+          finish_reason: 'stop',
+        },
+      ],
+    })}`;
+
+    const transformed = PerplexityAIChatCompleteStreamChunkTransform(
+      chunk,
+      'fallback-id',
+      {},
+      true
+    );
+    const [eventLine] = transformed.split('\n');
+    const event = JSON.parse(eventLine.replace(/^data: /, ''));
+
+    expect(event.usage).toEqual({
+      prompt_tokens: 10,
+      completion_tokens: 5,
+      total_tokens: 15,
+    });
+    expect(event).not.toHaveProperty('citations');
+  });
+});
diff --git a/src/providers/perplexity-ai/chatComplete.ts b/src/providers/perplexity-ai/chatComplete.ts
index 11889f7ce..e1c857c9e 100644
--- a/src/providers/perplexity-ai/chatComplete.ts
+++ b/src/providers/perplexity-ai/chatComplete.ts
@@ -101,6 +101,13 @@ interface PerplexityAIChatChoice {
   finish_reason: string | null;
 }
 
+interface PerplexityAIUsage {
+  prompt_tokens: number;
+  completion_tokens: number;
+  total_tokens: number;
+  [key: string]: unknown;
+}
+
 export interface PerplexityAIChatCompleteResponse {
   id: string;
   model: string;
@@ -108,11 +115,7 @@ export interface PerplexityAIChatCompleteResponse {
   created: number;
   citations: string[];
   choices: PerplexityAIChatChoice[];
-  usage: {
-    prompt_tokens: number;
-    completion_tokens: number;
-    total_tokens: number;
-  };
+  usage: PerplexityAIUsage;
 }
 
 export interface PerplexityAIErrorResponse {
@@ -129,14 +132,18 @@ export interface PerplexityAIChatCompletionStreamChunk {
   object: string;
   created: number;
   citations?: string[];
-  usage: {
-    prompt_tokens: number;
-    completion_tokens: number;
-    total_tokens: number;
-  };
+  usage: PerplexityAIUsage;
   choices: PerplexityAIChatChoice[];
 }
 
+function toOpenAIUsage(usage: PerplexityAIUsage) {
+  return {
+    prompt_tokens: usage.prompt_tokens,
+    completion_tokens: usage.completion_tokens,
+    total_tokens: usage.total_tokens,
+  };
+}
+
 export const PerplexityAIChatCompleteResponseTransform: (
   response: PerplexityAIChatCompleteResponse | PerplexityAIErrorResponse,
   responseStatus: number,
@@ -181,11 +188,9 @@ export const PerplexityAIChatCompleteResponseTransform: (
           finish_reason: '',
         },
       ],
-      usage: {
-        prompt_tokens: response.usage.prompt_tokens,
-        completion_tokens: response.usage.completion_tokens,
-        total_tokens: response.usage.total_tokens,
-      },
+      usage: strictOpenAiCompliance
+        ? toOpenAIUsage(response.usage)
+        : response.usage,
     };
   }
 
@@ -229,7 +234,11 @@ export const PerplexityAIChatCompleteStreamChunkTransform: (
         },
       ],
       ...(parsedChunk.usage &&
-        parsedChunk.choices[0]?.finish_reason && { usage: parsedChunk.usage }),
+        parsedChunk.choices[0]?.finish_reason && {
+          usage: strictOpenAiCompliance
+            ? toOpenAIUsage(parsedChunk.usage)
+            : parsedChunk.usage,
+        }),
     })}` + '\n\n';
 
   if (parsedChunk.choices[0]?.finish_reason)