fix: Add warning message for token limit truncation (#2260)

Co-authored-by: Sandy Tao <[email protected]>

2 files changed, 177 insertions, 0 deletions

diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts
index b0c27f7e..2a557927 100644
--- a/packages/core/src/core/turn.test.ts
+++ b/packages/core/src/core/turn.test.ts
@@ -282,6 +282,165 @@ describe('Turn', () => {
       expect(turn.pendingToolCalls[2]).toEqual(event3.value);
       expect(turn.getDebugResponses().length).toBe(1);
     });
+
+    it('should yield finished event when response has finish reason', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'Partial response' }] },
+              finishReason: 'STOP',
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      const reqParts: Part[] = [{ text: 'Test finish reason' }];
+      for await (const event of turn.run(
+        reqParts,
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        { type: GeminiEventType.Content, value: 'Partial response' },
+        { type: GeminiEventType.Finished, value: 'STOP' },
+      ]);
+    });
+
+    it('should yield finished event for MAX_TOKENS finish reason', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: {
+                parts: [
+                  { text: 'This is a long response that was cut off...' },
+                ],
+              },
+              finishReason: 'MAX_TOKENS',
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      const reqParts: Part[] = [{ text: 'Generate long text' }];
+      for await (const event of turn.run(
+        reqParts,
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        {
+          type: GeminiEventType.Content,
+          value: 'This is a long response that was cut off...',
+        },
+        { type: GeminiEventType.Finished, value: 'MAX_TOKENS' },
+      ]);
+    });
+
+    it('should yield finished event for SAFETY finish reason', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'Content blocked' }] },
+              finishReason: 'SAFETY',
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      const reqParts: Part[] = [{ text: 'Test safety' }];
+      for await (const event of turn.run(
+        reqParts,
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        { type: GeminiEventType.Content, value: 'Content blocked' },
+        { type: GeminiEventType.Finished, value: 'SAFETY' },
+      ]);
+    });
+
+    it('should not yield finished event when there is no finish reason', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'Response without finish reason' }] },
+              // No finishReason property
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      const reqParts: Part[] = [{ text: 'Test no finish reason' }];
+      for await (const event of turn.run(
+        reqParts,
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        {
+          type: GeminiEventType.Content,
+          value: 'Response without finish reason',
+        },
+      ]);
+      // No Finished event should be emitted
+    });
+
+    it('should handle multiple responses with different finish reasons', async () => {
+      const mockResponseStream = (async function* () {
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'First part' }] },
+              // No finish reason on first response
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+        yield {
+          candidates: [
+            {
+              content: { parts: [{ text: 'Second part' }] },
+              finishReason: 'OTHER',
+            },
+          ],
+        } as unknown as GenerateContentResponse;
+      })();
+      mockSendMessageStream.mockResolvedValue(mockResponseStream);
+
+      const events = [];
+      const reqParts: Part[] = [{ text: 'Test multiple responses' }];
+      for await (const event of turn.run(
+        reqParts,
+        new AbortController().signal,
+      )) {
+        events.push(event);
+      }
+
+      expect(events).toEqual([
+        { type: GeminiEventType.Content, value: 'First part' },
+        { type: GeminiEventType.Content, value: 'Second part' },
+        { type: GeminiEventType.Finished, value: 'OTHER' },
+      ]);
+    });
   });
 
   describe('getDebugResponses', () => {
diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts
index 836a162d..bea29b66 100644
--- a/packages/core/src/core/turn.ts
+++ b/packages/core/src/core/turn.ts
@@ -9,6 +9,7 @@ import {
   GenerateContentResponse,
   FunctionCall,
   FunctionDeclaration,
+  FinishReason,
 } from '@google/genai';
 import {
   ToolCallConfirmationDetails,
@@ -49,6 +50,7 @@ export enum GeminiEventType {
   ChatCompressed = 'chat_compressed',
   Thought = 'thought',
   MaxSessionTurns = 'max_session_turns',
+  Finished = 'finished',
   LoopDetected = 'loop_detected',
 }
 
@@ -134,6 +136,11 @@ export type ServerGeminiMaxSessionTurnsEvent = {
   type: GeminiEventType.MaxSessionTurns;
 };
 
+export type ServerGeminiFinishedEvent = {
+  type: GeminiEventType.Finished;
+  value: FinishReason;
+};
+
 export type ServerGeminiLoopDetectedEvent = {
   type: GeminiEventType.LoopDetected;
 };
@@ -149,6 +156,7 @@ export type ServerGeminiStreamEvent =
   | ServerGeminiChatCompressedEvent
   | ServerGeminiThoughtEvent
   | ServerGeminiMaxSessionTurnsEvent
+  | ServerGeminiFinishedEvent
   | ServerGeminiLoopDetectedEvent;
 
 // A turn manages the agentic loop turn within the server context.
@@ -222,6 +230,16 @@ export class Turn {
             yield event;
           }
         }
+
+        // Check if response was truncated or stopped for various reasons
+        const finishReason = resp.candidates?.[0]?.finishReason;
+
+        if (finishReason) {
+          yield {
+            type: GeminiEventType.Finished,
+            value: finishReason as FinishReason,
+          };
+        }
       }
     } catch (e) {
       const error = toFriendlyError(e);