@lobehub/chat 1.84.26 → 1.84.27

package/CHANGELOG.md

@@ -2,6 +2,31 @@
 
 # Changelog
 
+### [Version 1.84.27](https://github.com/lobehub/lobe-chat/compare/v1.84.26...v1.84.27)
+
+<sup>Released on **2025-05-09**</sup>
+
+#### 💄 Styles
+
+- **misc**: Add reasoning tokens and token usage statistics for Google Gemini.
+
+<br/>
+
+<details>
+<summary><kbd>Improvements and Fixes</kbd></summary>
+
+#### Styles
+
+- **misc**: Add reasoning tokens and token usage statistics for Google Gemini, closes [#7501](https://github.com/lobehub/lobe-chat/issues/7501) ([b466b42](https://github.com/lobehub/lobe-chat/commit/b466b42))
+
+</details>
+
+<div align="right">
+
+[![](https://img.shields.io/badge/-BACK_TO_TOP-151515?style=flat-square)](#readme-top)
+
+</div>
+
 ### [Version 1.84.26](https://github.com/lobehub/lobe-chat/compare/v1.84.25...v1.84.26)
 
 <sup>Released on **2025-05-08**</sup>

package/changelog/v1.json

@@ -1,4 +1,13 @@
 [
+  {
+    "children": {
+      "improvements": [
+        "Add reasoning tokens and token usage statistics for Google Gemini."
+      ]
+    },
+    "date": "2025-05-09",
+    "version": "1.84.27"
+  },
   {
     "children": {
       "improvements": [

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lobehub/chat",
-  "version": "1.84.26",
+  "version": "1.84.27",
   "description": "Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.",
   "keywords": [
     "framework",
@@ -121,7 +121,7 @@
   "dependencies": {
     "@ant-design/icons": "^5.6.1",
     "@ant-design/pro-components": "^2.8.7",
-    "@anthropic-ai/sdk": "^0.40.1",
+    "@anthropic-ai/sdk": "^0.41.0",
     "@auth/core": "^0.38.0",
     "@aws-sdk/client-bedrock-runtime": "^3.779.0",
     "@aws-sdk/client-s3": "^3.779.0",

package/src/config/aiModels/vertexai.ts

@@ -98,9 +98,9 @@ const vertexaiChatModels: AIChatModelCard[] = [
     type: 'chat',
   },
   {
-    abilities: { 
-      functionCall: true, 
-      vision: true 
+    abilities: {
+      functionCall: true,
+      vision: true
     },
     contextWindowTokens: 1_000_000 + 8192,
     description: 'Gemini 1.5 Flash 002 是一款高效的多模态模型，支持广泛应用的扩展。',
@@ -115,9 +115,9 @@ const vertexaiChatModels: AIChatModelCard[] = [
     type: 'chat',
   },
   {
-    abilities: { 
-      functionCall: true, 
-      vision: true 
+    abilities: {
+      functionCall: true,
+      vision: true
     },
     contextWindowTokens: 2_000_000 + 8192,
     description:

package/src/libs/agent-runtime/google/index.ts

@@ -106,6 +106,7 @@ export class LobeGoogleAI implements LobeRuntimeAI {
 
       const contents = await this.buildGoogleMessages(payload.messages);
 
+      const inputStartAt = Date.now();
       const geminiStreamResult = await this.client
         .getGenerativeModel(
           {
@@ -161,7 +162,7 @@ export class LobeGoogleAI implements LobeRuntimeAI {
 
       // Convert the response into a friendly text-stream
       const Stream = this.isVertexAi ? VertexAIStream : GoogleGenerativeAIStream;
-      const stream = Stream(prod, options?.callback);
+      const stream = Stream(prod, { callbacks: options?.callback, inputStartAt });
 
       // Respond with the stream
       return StreamingResponse(stream, { headers: options?.headers });

package/src/libs/agent-runtime/utils/streams/google-ai.test.ts

@@ -34,10 +34,12 @@ describe('GoogleGenerativeAIStream', () => {
     const onCompletionMock = vi.fn();
 
     const protocolStream = GoogleGenerativeAIStream(mockGoogleStream, {
-      onStart: onStartMock,
-      onText: onTextMock,
-      onToolsCalling: onToolCallMock,
-      onCompletion: onCompletionMock,
+      callbacks: {
+        onStart: onStartMock,
+        onText: onTextMock,
+        onToolsCalling: onToolCallMock,
+        onCompletion: onCompletionMock,
+      },
     });
 
     const decoder = new TextDecoder();
@@ -187,7 +189,7 @@ describe('GoogleGenerativeAIStream', () => {
       // usage
       'id: chat_1\n',
       'event: usage\n',
-      `data: {"inputImageTokens":258,"inputTextTokens":8,"totalInputTokens":266,"totalTokens":266}\n\n`,
+      `data: {"inputImageTokens":258,"inputTextTokens":8,"outputTextTokens":0,"totalInputTokens":266,"totalOutputTokens":0,"totalTokens":266}\n\n`,
     ]);
   });
 
@@ -276,7 +278,100 @@ describe('GoogleGenerativeAIStream', () => {
         // usage
         'id: chat_1',
         'event: usage',
-        `data: {"inputTextTokens":19,"totalInputTokens":19,"totalOutputTokens":11,"totalTokens":30}\n`,
+        `data: {"inputTextTokens":19,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":11,"totalTokens":30}\n`,
+      ].map((i) => i + '\n'),
+    );
+  });
+
+  it('should handle stop with content and thought', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = [
+      {
+        candidates: [
+          {
+            content: { parts: [{ text: '234' }], role: 'model' },
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        text: () => '234',
+        usageMetadata: {
+          promptTokenCount: 19,
+          candidatesTokenCount: 3,
+          totalTokenCount: 122,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }],
+          thoughtsTokenCount: 100,
+        },
+        modelVersion: 'gemini-2.0-flash-exp-image-generation',
+      },
+      {
+        text: () => '567890\n',
+        candidates: [
+          {
+            content: { parts: [{ text: '567890\n' }], role: 'model' },
+            finishReason: 'STOP',
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 19,
+          candidatesTokenCount: 11,
+          totalTokenCount: 131,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }],
+          thoughtsTokenCount: 100,
+        },
+        modelVersion: 'gemini-2.0-flash-exp-image-generation',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        data.forEach((item) => {
+          controller.enqueue(item);
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = GoogleGenerativeAIStream(mockGoogleStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual(
+      [
+        'id: chat_1',
+        'event: text',
+        'data: "234"\n',
+
+        'id: chat_1',
+        'event: text',
+        `data: "567890\\n"\n`,
+        // stop
+        'id: chat_1',
+        'event: stop',
+        `data: "STOP"\n`,
+        // usage
+        'id: chat_1',
+        'event: usage',
+        `data: {"inputTextTokens":19,"outputReasoningTokens":100,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":111,"totalTokens":131}\n`,
       ].map((i) => i + '\n'),
     );
   });

package/src/libs/agent-runtime/utils/streams/google-ai.ts

@@ -11,6 +11,7 @@ import {
   StreamToolCallChunkData,
   createCallbacksTransformer,
   createSSEProtocolTransformer,
+  createTokenSpeedCalculator,
   generateToolCallId,
 } from './protocol';
 
@@ -19,31 +20,62 @@ const transformGoogleGenerativeAIStream = (
   context: StreamContext,
 ): StreamProtocolChunk | StreamProtocolChunk[] => {
   // maybe need another structure to add support for multiple choices
+  const candidate = chunk.candidates?.[0];
+  const usage = chunk.usageMetadata;
+  const usageChunks: StreamProtocolChunk[] = [];
+  if (candidate?.finishReason && usage) {
+    const outputReasoningTokens = (usage as any).thoughtsTokenCount || undefined;
+    const totalOutputTokens = (usage.candidatesTokenCount ?? 0) + (outputReasoningTokens ?? 0);
+
+    usageChunks.push(
+      { data: candidate.finishReason, id: context?.id, type: 'stop' },
+      {
+        data: {
+          // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
+          inputImageTokens: (usage as any).promptTokensDetails?.find(
+            (i: any) => i.modality === 'IMAGE',
+          )?.tokenCount,
+          inputTextTokens: (usage as any).promptTokensDetails?.find(
+            (i: any) => i.modality === 'TEXT',
+          )?.tokenCount,
+          outputReasoningTokens,
+          outputTextTokens: totalOutputTokens - (outputReasoningTokens ?? 0),
+          totalInputTokens: usage.promptTokenCount,
+          totalOutputTokens,
+          totalTokens: usage.totalTokenCount,
+        } as ModelTokensUsage,
+        id: context?.id,
+        type: 'usage',
+      },
+    );
+  }
+
   const functionCalls = chunk.functionCalls?.();
 
   if (functionCalls) {
-    return {
-      data: functionCalls.map(
-        (value, index): StreamToolCallChunkData => ({
-          function: {
-            arguments: JSON.stringify(value.args),
-            name: value.name,
-          },
-          id: generateToolCallId(index, value.name),
-          index: index,
-          type: 'function',
-        }),
-      ),
-      id: context.id,
-      type: 'tool_calls',
-    };
+    return [
+      {
+        data: functionCalls.map(
+          (value, index): StreamToolCallChunkData => ({
+            function: {
+              arguments: JSON.stringify(value.args),
+              name: value.name,
+            },
+            id: generateToolCallId(index, value.name),
+            index: index,
+            type: 'function',
+          }),
+        ),
+        id: context.id,
+        type: 'tool_calls',
+      },
+      ...usageChunks,
+    ];
   }
 
   const text = chunk.text?.();
 
-  if (chunk.candidates) {
-    const candidate = chunk.candidates[0];
-
+  if (candidate) {
     // return the grounding
     if (candidate.groundingMetadata) {
       const { webSearchQueries, groundingChunks } = candidate.groundingMetadata;
@@ -64,31 +96,15 @@ const transformGoogleGenerativeAIStream = (
           id: context.id,
           type: 'grounding',
         },
+        ...usageChunks,
       ];
     }
 
     if (candidate.finishReason) {
       if (chunk.usageMetadata) {
-        const usage = chunk.usageMetadata;
         return [
           !!text ? { data: text, id: context?.id, type: 'text' } : undefined,
-          { data: candidate.finishReason, id: context?.id, type: 'stop' },
-          {
-            data: {
-              // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
-              inputImageTokens: (usage as any).promptTokensDetails?.find(
-                (i: any) => i.modality === 'IMAGE',
-              )?.tokenCount,
-              inputTextTokens: (usage as any).promptTokensDetails?.find(
-                (i: any) => i.modality === 'TEXT',
-              )?.tokenCount,
-              totalInputTokens: usage.promptTokenCount,
-              totalOutputTokens: usage.candidatesTokenCount,
-              totalTokens: usage.totalTokenCount,
-            } as ModelTokensUsage,
-            id: context?.id,
-            type: 'usage',
-          },
+          ...usageChunks,
         ].filter(Boolean) as StreamProtocolChunk[];
       }
       return { data: candidate.finishReason, id: context?.id, type: 'stop' };
@@ -117,13 +133,21 @@ const transformGoogleGenerativeAIStream = (
   };
 };
 
+export interface GoogleAIStreamOptions {
+  callbacks?: ChatStreamCallbacks;
+  inputStartAt?: number;
+}
+
 export const GoogleGenerativeAIStream = (
   rawStream: ReadableStream<EnhancedGenerateContentResponse>,
-  callbacks?: ChatStreamCallbacks,
+  { callbacks, inputStartAt }: GoogleAIStreamOptions = {},
 ) => {
   const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   return rawStream
-    .pipeThrough(createSSEProtocolTransformer(transformGoogleGenerativeAIStream, streamStack))
+    .pipeThrough(
+      createTokenSpeedCalculator(transformGoogleGenerativeAIStream, { inputStartAt, streamStack }),
+    )
+    .pipeThrough(createSSEProtocolTransformer((c) => c, streamStack))
     .pipeThrough(createCallbacksTransformer(callbacks));
 };

package/src/libs/agent-runtime/utils/streams/protocol.ts

@@ -298,17 +298,37 @@ export const TOKEN_SPEED_CHUNK_ID = 'output_speed';
  */
 export const createTokenSpeedCalculator = (
   transformer: (chunk: any, stack: StreamContext) => StreamProtocolChunk | StreamProtocolChunk[],
-  { streamStack, inputStartAt }: { inputStartAt?: number; streamStack?: StreamContext } = {},
+  { inputStartAt, streamStack }: { inputStartAt?: number; streamStack?: StreamContext } = {},
 ) => {
   let outputStartAt: number | undefined;
+  let outputThinking: boolean | undefined;
 
   const process = (chunk: StreamProtocolChunk) => {
     let result = [chunk];
-    // if the chunk is the first text chunk, set as output start
-    if (!outputStartAt && chunk.type === 'text') outputStartAt = Date.now();
+    // if the chunk is the first text or reasoning chunk, set as output start
+    if (!outputStartAt && (chunk.type === 'text' || chunk.type === 'reasoning')) {
+      outputStartAt = Date.now();
+    }
+
+    /**
+     * 部分 provider 在正式输出 reasoning 前，可能会先输出 content 为空字符串的 chunk，
+     * 其中 reasoning 可能为 null，会导致判断是否输出思考内容错误，所以过滤掉 null 或者空字符串。
+     * 也可能是某些特殊 token，所以不修改 outputStartAt 的逻辑。
+     */
+    if (
+      outputThinking === undefined &&
+      (chunk.type === 'text' || chunk.type === 'reasoning') &&
+      typeof chunk.data === 'string' &&
+      chunk.data.length > 0
+    ) {
+      outputThinking = chunk.type === 'reasoning';
+    }
     // if the chunk is the stop chunk, set as output finish
     if (inputStartAt && outputStartAt && chunk.type === 'usage') {
-      const outputTokens = chunk.data?.totalOutputTokens || chunk.data?.outputTextTokens;
+      const totalOutputTokens = chunk.data?.totalOutputTokens || chunk.data?.outputTextTokens;
+      const reasoningTokens = chunk.data?.outputReasoningTokens || 0;
+      const outputTokens =
+        (outputThinking ?? false) ? totalOutputTokens : totalOutputTokens - reasoningTokens;
       result.push({
         data: {
           tps: (outputTokens / (Date.now() - outputStartAt)) * 1000,

package/src/libs/agent-runtime/utils/streams/vertex-ai.test.ts

@@ -103,10 +103,12 @@ describe('VertexAIStream', () => {
     const onCompletionMock = vi.fn();
 
     const protocolStream = VertexAIStream(mockGoogleStream, {
-      onStart: onStartMock,
-      onText: onTextMock,
-      onToolsCalling: onToolCallMock,
-      onCompletion: onCompletionMock,
+      callbacks: {
+        onStart: onStartMock,
+        onText: onTextMock,
+        onToolsCalling: onToolCallMock,
+        onCompletion: onCompletionMock,
+      },
     });
 
     const decoder = new TextDecoder();
@@ -136,6 +138,7 @@ describe('VertexAIStream', () => {
 
   it('tool_calls', async () => {
     vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
     const rawChunks = [
       {
         candidates: [
@@ -204,10 +207,12 @@ describe('VertexAIStream', () => {
     const onCompletionMock = vi.fn();
 
     const protocolStream = VertexAIStream(mockGoogleStream, {
-      onStart: onStartMock,
-      onText: onTextMock,
-      onToolsCalling: onToolCallMock,
-      onCompletion: onCompletionMock,
+      callbacks: {
+        onStart: onStartMock,
+        onText: onTextMock,
+        onToolsCalling: onToolCallMock,
+        onCompletion: onCompletionMock,
+      },
     });
 
     const decoder = new TextDecoder();
@@ -223,10 +228,106 @@ describe('VertexAIStream', () => {
       'id: chat_1\n',
       'event: tool_calls\n',
       `data: [{"function":{"arguments":"{\\"city\\":\\"杭州\\"}","name":"realtime-weather____fetchCurrentWeather"},"id":"realtime-weather____fetchCurrentWeather_0","index":0,"type":"function"}]\n\n`,
+      'id: chat_1\n',
+      'event: stop\n',
+      'data: "STOP"\n\n',
+      'id: chat_1\n',
+      'event: usage\n',
+      'data: {"outputTextTokens":9,"totalInputTokens":95,"totalOutputTokens":9,"totalTokens":104}\n\n',
     ]);
 
     expect(onStartMock).toHaveBeenCalledTimes(1);
     expect(onToolCallMock).toHaveBeenCalledTimes(1);
     expect(onCompletionMock).toHaveBeenCalledTimes(1);
   });
+
+  it('should handle stop with content', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const data = [
+      {
+        candidates: [
+          {
+            content: { parts: [{ text: '234' }], role: 'model' },
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        text: () => '234',
+        usageMetadata: {
+          promptTokenCount: 20,
+          totalTokenCount: 20,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 20 }],
+        },
+        modelVersion: 'gemini-2.0-flash-exp-image-generation',
+      },
+      {
+        text: () => '567890\n',
+        candidates: [
+          {
+            content: { parts: [{ text: '567890\n' }], role: 'model' },
+            finishReason: 'STOP',
+            safetyRatings: [
+              { category: 'HARM_CATEGORY_HATE_SPEECH', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_DANGEROUS_CONTENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_HARASSMENT', probability: 'NEGLIGIBLE' },
+              { category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT', probability: 'NEGLIGIBLE' },
+            ],
+          },
+        ],
+        usageMetadata: {
+          promptTokenCount: 19,
+          candidatesTokenCount: 11,
+          totalTokenCount: 30,
+          promptTokensDetails: [{ modality: 'TEXT', tokenCount: 19 }],
+          candidatesTokensDetails: [{ modality: 'TEXT', tokenCount: 11 }],
+        },
+        modelVersion: 'gemini-2.0-flash-exp-image-generation',
+      },
+    ];
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        data.forEach((item) => {
+          controller.enqueue(item);
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = VertexAIStream(mockGoogleStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual(
+      [
+        'id: chat_1',
+        'event: text',
+        'data: "234"\n',
+
+        'id: chat_1',
+        'event: text',
+        `data: "567890\\n"\n`,
+        // stop
+        'id: chat_1',
+        'event: stop',
+        `data: "STOP"\n`,
+        // usage
+        'id: chat_1',
+        'event: usage',
+        `data: {"inputTextTokens":19,"outputTextTokens":11,"totalInputTokens":19,"totalOutputTokens":11,"totalTokens":30}\n`,
+      ].map((i) => i + '\n'),
+    );
+  });
 });

package/src/libs/agent-runtime/utils/streams/vertex-ai.ts

@@ -1,27 +1,58 @@
 import { EnhancedGenerateContentResponse, GenerateContentResponse } from '@google/generative-ai';
 
+import { ModelTokensUsage } from '@/types/message';
 import { nanoid } from '@/utils/uuid';
 
-import { ChatStreamCallbacks } from '../../types';
+import { type GoogleAIStreamOptions } from './google-ai';
 import {
   StreamContext,
   StreamProtocolChunk,
   createCallbacksTransformer,
   createSSEProtocolTransformer,
+  createTokenSpeedCalculator,
   generateToolCallId,
 } from './protocol';
 
 const transformVertexAIStream = (
   chunk: GenerateContentResponse,
-  stack: StreamContext,
-): StreamProtocolChunk => {
+  context: StreamContext,
+): StreamProtocolChunk | StreamProtocolChunk[] => {
   // maybe need another structure to add support for multiple choices
-  const candidates = chunk.candidates;
+  const candidate = chunk.candidates?.[0];
+  const usage = chunk.usageMetadata;
+  const usageChunks: StreamProtocolChunk[] = [];
+  if (candidate?.finishReason && usage) {
+    const outputReasoningTokens = (usage as any).thoughtsTokenCount || undefined;
+    const totalOutputTokens = (usage.candidatesTokenCount ?? 0) + (outputReasoningTokens ?? 0);
+
+    usageChunks.push(
+      { data: candidate.finishReason, id: context?.id, type: 'stop' },
+      {
+        data: {
+          // TODO: Google SDK 0.24.0 don't have promptTokensDetails types
+          inputImageTokens: (usage as any).promptTokensDetails?.find(
+            (i: any) => i.modality === 'IMAGE',
+          )?.tokenCount,
+          inputTextTokens: (usage as any).promptTokensDetails?.find(
+            (i: any) => i.modality === 'TEXT',
+          )?.tokenCount,
+          outputReasoningTokens,
+          outputTextTokens: totalOutputTokens - (outputReasoningTokens ?? 0),
+          totalInputTokens: usage.promptTokenCount,
+          totalOutputTokens,
+          totalTokens: usage.totalTokenCount,
+        } as ModelTokensUsage,
+        id: context?.id,
+        type: 'usage',
+      },
+    );
+  }
 
+  const candidates = chunk.candidates;
   if (!candidates)
     return {
       data: '',
-      id: stack?.id,
+      id: context?.id,
       type: 'text',
     };
 
@@ -32,44 +63,58 @@ const transformVertexAIStream = (
     if (part.functionCall) {
       const functionCall = part.functionCall;
 
-      return {
-        data: [
-          {
-            function: {
-              arguments: JSON.stringify(functionCall.args),
-              name: functionCall.name,
+      return [
+        {
+          data: [
+            {
+              function: {
+                arguments: JSON.stringify(functionCall.args),
+                name: functionCall.name,
+              },
+              id: generateToolCallId(0, functionCall.name),
+              index: 0,
+              type: 'function',
             },
-            id: generateToolCallId(0, functionCall.name),
-            index: 0,
-            type: 'function',
-          },
-        ],
-        id: stack?.id,
-        type: 'tool_calls',
-      };
+          ],
+          id: context?.id,
+          type: 'tool_calls',
+        },
+        ...usageChunks,
+      ];
+    }
+
+    if (item.finishReason) {
+      if (chunk.usageMetadata) {
+        return [
+          !!part.text ? { data: part.text, id: context?.id, type: 'text' } : undefined,
+          ...usageChunks,
+        ].filter(Boolean) as StreamProtocolChunk[];
+      }
+      return { data: item.finishReason, id: context?.id, type: 'stop' };
     }
 
     return {
       data: part.text,
-      id: stack?.id,
+      id: context?.id,
       type: 'text',
     };
   }
 
   return {
     data: '',
-    id: stack?.id,
+    id: context?.id,
     type: 'stop',
   };
 };
 
 export const VertexAIStream = (
   rawStream: ReadableStream<EnhancedGenerateContentResponse>,
-  callbacks?: ChatStreamCallbacks,
+  { callbacks, inputStartAt }: GoogleAIStreamOptions = {},
 ) => {
   const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   return rawStream
-    .pipeThrough(createSSEProtocolTransformer(transformVertexAIStream, streamStack))
+    .pipeThrough(createTokenSpeedCalculator(transformVertexAIStream, { inputStartAt, streamStack }))
+    .pipeThrough(createSSEProtocolTransformer((c) => c, streamStack))
     .pipeThrough(createCallbacksTransformer(callbacks));
 };

package/src/server/routers/lambda/__tests__/file.test.ts

@@ -0,0 +1,213 @@
+import { TRPCError } from '@trpc/server';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { fileRouter } from '@/server/routers/lambda/file';
+import { AsyncTaskStatus } from '@/types/asyncTask';
+
+// Patch: Use actual router context middleware to inject the correct models/services
+function createCallerWithCtx(partialCtx: any = {}) {
+  // All mocks are spies
+  const fileModel = {
+    checkHash: vi.fn().mockResolvedValue({ isExist: true }),
+    create: vi.fn().mockResolvedValue({ id: 'test-id' }),
+    findById: vi.fn().mockResolvedValue(undefined),
+    query: vi.fn().mockResolvedValue([]),
+    delete: vi.fn().mockResolvedValue(undefined),
+    deleteMany: vi.fn().mockResolvedValue([]),
+    clear: vi.fn().mockResolvedValue({} as any),
+  };
+
+  const fileService = {
+    getFullFileUrl: vi.fn().mockResolvedValue('full-url'),
+    deleteFile: vi.fn().mockResolvedValue(undefined),
+    deleteFiles: vi.fn().mockResolvedValue(undefined),
+  };
+
+  const chunkModel = {
+    countByFileIds: vi.fn().mockResolvedValue([{ id: 'test-id', count: 5 }]),
+    countByFileId: vi.fn().mockResolvedValue(5),
+  };
+
+  const asyncTaskModel = {
+    findByIds: vi.fn().mockResolvedValue([
+      {
+        id: 'test-task-id',
+        status: AsyncTaskStatus.Success,
+      },
+    ]),
+    findById: vi.fn(),
+    delete: vi.fn(),
+  };
+
+  const ctx = {
+    serverDB: {} as any,
+    userId: 'test-user',
+    asyncTaskModel,
+    chunkModel,
+    fileModel,
+    fileService,
+    ...partialCtx,
+  };
+
+  return { ctx, caller: fileRouter.createCaller(ctx) };
+}
+
+vi.mock('@/config/db', () => ({
+  serverDBEnv: {
+    REMOVE_GLOBAL_FILE: false,
+  },
+}));
+
+vi.mock('@/database/models/asyncTask', () => ({
+  AsyncTaskModel: vi.fn(() => ({
+    findById: vi.fn(),
+    findByIds: vi.fn(),
+    delete: vi.fn(),
+  })),
+}));
+
+vi.mock('@/database/models/chunk', () => ({
+  ChunkModel: vi.fn(() => ({
+    countByFileId: vi.fn(),
+    countByFileIds: vi.fn(),
+  })),
+}));
+
+vi.mock('@/database/models/file', () => ({
+  FileModel: vi.fn(() => ({
+    checkHash: vi.fn(),
+    create: vi.fn(),
+    delete: vi.fn(),
+    deleteMany: vi.fn(),
+    findById: vi.fn(),
+    query: vi.fn(),
+    clear: vi.fn(),
+  })),
+}));
+
+vi.mock('@/server/services/file', () => ({
+  FileService: vi.fn(() => ({
+    getFullFileUrl: vi.fn(),
+    deleteFile: vi.fn(),
+    deleteFiles: vi.fn(),
+  })),
+}));
+
+describe('fileRouter', () => {
+  let ctx: any;
+  let caller: any;
+  let mockFile: any;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    mockFile = {
+      id: 'test-id',
+      name: 'test.txt',
+      url: 'test-url',
+      createdAt: new Date(),
+      updatedAt: new Date(),
+      accessedAt: new Date(),
+      userId: 'test-user',
+      size: 100,
+      fileType: 'text',
+      metadata: {},
+      fileHash: null,
+      clientId: null,
+      chunkTaskId: null,
+      embeddingTaskId: null,
+    };
+
+    // Use actual context with default mocks
+    ({ ctx, caller } = createCallerWithCtx());
+  });
+
+  describe('checkFileHash', () => {
+    it('should handle when fileModel.checkHash returns undefined', async () => {
+      ctx.fileModel.checkHash.mockResolvedValue(undefined);
+      await expect(caller.checkFileHash({ hash: 'test-hash' })).resolves.toBeUndefined();
+    });
+  });
+
+  describe('createFile', () => {
+    it('should throw if fileModel.checkHash returns undefined', async () => {
+      ctx.fileModel.checkHash.mockResolvedValue(undefined);
+      await expect(
+        caller.createFile({
+          hash: 'test-hash',
+          fileType: 'text',
+          name: 'test.txt',
+          size: 100,
+          url: 'test-url',
+          metadata: {},
+        }),
+      ).rejects.toThrow();
+    });
+  });
+
+  describe('findById', () => {
+    it('should throw error when file not found', async () => {
+      ctx.fileModel.findById.mockResolvedValue(null);
+
+      await expect(caller.findById({ id: 'invalid-id' })).rejects.toThrow(TRPCError);
+    });
+  });
+
+  describe('getFileItemById', () => {
+    it('should throw error when file not found', async () => {
+      ctx.fileModel.findById.mockResolvedValue(null);
+
+      await expect(caller.getFileItemById({ id: 'invalid-id' })).rejects.toThrow(TRPCError);
+    });
+  });
+
+  describe('getFiles', () => {
+    it('should handle fileModel.query returning undefined', async () => {
+      ctx.fileModel.query.mockResolvedValue(undefined);
+
+      await expect(caller.getFiles({})).rejects.toThrow();
+    });
+  });
+
+  describe('removeFile', () => {
+    it('should do nothing when file not found', async () => {
+      ctx.fileModel.delete.mockResolvedValue(null);
+
+      await caller.removeFile({ id: 'invalid-id' });
+
+      expect(ctx.fileService.deleteFile).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('removeFiles', () => {
+    it('should do nothing when no files found', async () => {
+      ctx.fileModel.deleteMany.mockResolvedValue([]);
+
+      await caller.removeFiles({ ids: ['invalid-1', 'invalid-2'] });
+
+      expect(ctx.fileService.deleteFiles).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('removeFileAsyncTask', () => {
+    it('should do nothing when file not found', async () => {
+      ctx.fileModel.findById.mockResolvedValue(null);
+
+      await caller.removeFileAsyncTask({ id: 'test-id', type: 'chunk' });
+
+      expect(ctx.asyncTaskModel.delete).not.toHaveBeenCalled();
+    });
+
+    it('should do nothing when task id is missing', async () => {
+      ctx.fileModel.findById.mockResolvedValue(mockFile);
+
+      await caller.removeFileAsyncTask({ id: 'test-id', type: 'embedding' });
+
+      expect(ctx.asyncTaskModel.delete).not.toHaveBeenCalled();
+
+      await caller.removeFileAsyncTask({ id: 'test-id', type: 'chunk' });
+
+      expect(ctx.asyncTaskModel.delete).not.toHaveBeenCalled();
+    });
+  });
+});