From 4c3532d2b395859fb0d9db00bc5209f160ce2e29 Mon Sep 17 00:00:00 2001 From: darkcocoa <67269224+darkdarkcocoa@users.noreply.github.com> Date: Tue, 22 Jul 2025 06:57:11 +0900 Subject: [PATCH] fix: Add warning message for token limit truncation (#2260) Co-authored-by: Sandy Tao --- .../cli/src/ui/hooks/useGeminiStream.test.tsx | 238 +++++++++++++++++- packages/cli/src/ui/hooks/useGeminiStream.ts | 50 +++- packages/core/src/core/turn.test.ts | 159 ++++++++++++ packages/core/src/core/turn.ts | 18 ++ 4 files changed, 463 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index 7e45cab2..d7fd35c8 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -16,7 +16,12 @@ import { TrackedExecutingToolCall, TrackedCancelledToolCall, } from './useReactToolScheduler.js'; -import { Config, EditorType, AuthType } from '@google/gemini-cli-core'; +import { + Config, + EditorType, + AuthType, + GeminiEventType as ServerGeminiEventType, +} from '@google/gemini-cli-core'; import { Part, PartListUnion } from '@google/genai'; import { UseHistoryManagerReturn } from './useHistoryManager.js'; import { @@ -1178,4 +1183,235 @@ describe('useGeminiStream', () => { }); }); }); + + describe('handleFinishedEvent', () => { + it('should add info message for MAX_TOKENS finish reason', async () => { + // Setup mock to return a stream with MAX_TOKENS finish reason + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'This is a truncated response...', + }; + yield { type: ServerGeminiEventType.Finished, value: 'MAX_TOKENS' }; + })(), + ); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockSetShowHelp, + mockConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + ), + ); + + // Submit a query + await act(async () => { + await result.current.submitQuery('Generate long text'); + }); + + // Check that the info message was added + await waitFor(() => { + expect(mockAddItem).toHaveBeenCalledWith( + { + type: 'info', + text: '⚠️ Response truncated due to token limits.', + }, + expect.any(Number), + ); + }); + }); + + it('should not add message for STOP finish reason', async () => { + // Setup mock to return a stream with STOP finish reason + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'Complete response', + }; + yield { type: ServerGeminiEventType.Finished, value: 'STOP' }; + })(), + ); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockSetShowHelp, + mockConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + ), + ); + + // Submit a query + await act(async () => { + await result.current.submitQuery('Test normal completion'); + }); + + // Wait a bit to ensure no message is added + await new Promise((resolve) => setTimeout(resolve, 100)); + + // Check that no info message was added for STOP + const infoMessages = mockAddItem.mock.calls.filter( + (call) => call[0].type === 'info', + ); + expect(infoMessages).toHaveLength(0); + }); + + it('should not add message for FINISH_REASON_UNSPECIFIED', async () => { + // Setup mock to return a stream with FINISH_REASON_UNSPECIFIED + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'Response with unspecified finish', + }; + yield { + type: ServerGeminiEventType.Finished, + value: 'FINISH_REASON_UNSPECIFIED', + }; + })(), + ); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockSetShowHelp, + mockConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + ), + ); + + // Submit a query + await act(async () => { + await result.current.submitQuery('Test unspecified finish'); + }); + + // Wait a bit to ensure no message is added + await new Promise((resolve) => setTimeout(resolve, 100)); + + // Check that no info message was added + const infoMessages = mockAddItem.mock.calls.filter( + (call) => call[0].type === 'info', + ); + expect(infoMessages).toHaveLength(0); + }); + + it('should add appropriate messages for other finish reasons', async () => { + const testCases = [ + { + reason: 'SAFETY', + message: '⚠️ Response stopped due to safety reasons.', + }, + { + reason: 'RECITATION', + message: '⚠️ Response stopped due to recitation policy.', + }, + { + reason: 'LANGUAGE', + message: '⚠️ Response stopped due to unsupported language.', + }, + { + reason: 'BLOCKLIST', + message: '⚠️ Response stopped due to forbidden terms.', + }, + { + reason: 'PROHIBITED_CONTENT', + message: '⚠️ Response stopped due to prohibited content.', + }, + { + reason: 'SPII', + message: + '⚠️ Response stopped due to sensitive personally identifiable information.', + }, + { reason: 'OTHER', message: '⚠️ Response stopped for other reasons.' }, + { + reason: 'MALFORMED_FUNCTION_CALL', + message: '⚠️ Response stopped due to malformed function call.', + }, + { + reason: 'IMAGE_SAFETY', + message: '⚠️ Response stopped due to image safety violations.', + }, + { + reason: 'UNEXPECTED_TOOL_CALL', + message: '⚠️ Response stopped due to unexpected tool call.', + }, + ]; + + for (const { reason, message } of testCases) { + // Reset mocks for each test case + mockAddItem.mockClear(); + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: `Response for ${reason}`, + }; + yield { type: ServerGeminiEventType.Finished, value: reason }; + })(), + ); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockSetShowHelp, + mockConfig, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + ), + ); + + await act(async () => { + await result.current.submitQuery(`Test ${reason}`); + }); + + await waitFor(() => { + expect(mockAddItem).toHaveBeenCalledWith( + { + type: 'info', + text: message, + }, + expect.any(Number), + ); + }); + } + }); + }); }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 9fa23c52..295b5650 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -14,6 +14,7 @@ import { ServerGeminiContentEvent as ContentEvent, ServerGeminiErrorEvent as ErrorEvent, ServerGeminiChatCompressedEvent, + ServerGeminiFinishedEvent, getErrorMessage, isNodeError, MessageSenderType, @@ -26,7 +27,7 @@ import { UserPromptEvent, DEFAULT_GEMINI_FLASH_MODEL, } from '@google/gemini-cli-core'; -import { type Part, type PartListUnion } from '@google/genai'; +import { type Part, type PartListUnion, FinishReason } from '@google/genai'; import { StreamingState, HistoryItem, @@ -422,6 +423,46 @@ export const useGeminiStream = ( [addItem, pendingHistoryItemRef, setPendingHistoryItem, config], ); + const handleFinishedEvent = useCallback( + (event: ServerGeminiFinishedEvent, userMessageTimestamp: number) => { + const finishReason = event.value; + + const finishReasonMessages: Record = { + [FinishReason.FINISH_REASON_UNSPECIFIED]: undefined, + [FinishReason.STOP]: undefined, + [FinishReason.MAX_TOKENS]: 'Response truncated due to token limits.', + [FinishReason.SAFETY]: 'Response stopped due to safety reasons.', + [FinishReason.RECITATION]: 'Response stopped due to recitation policy.', + [FinishReason.LANGUAGE]: + 'Response stopped due to unsupported language.', + [FinishReason.BLOCKLIST]: 'Response stopped due to forbidden terms.', + [FinishReason.PROHIBITED_CONTENT]: + 'Response stopped due to prohibited content.', + [FinishReason.SPII]: + 'Response stopped due to sensitive personally identifiable information.', + [FinishReason.OTHER]: 'Response stopped for other reasons.', + [FinishReason.MALFORMED_FUNCTION_CALL]: + 'Response stopped due to malformed function call.', + [FinishReason.IMAGE_SAFETY]: + 'Response stopped due to image safety violations.', + [FinishReason.UNEXPECTED_TOOL_CALL]: + 'Response stopped due to unexpected tool call.', + }; + + const message = finishReasonMessages[finishReason]; + if (message) { + addItem( + { + type: 'info', + text: `⚠️ ${message}`, + }, + userMessageTimestamp, + ); + } + }, + [addItem], + ); + const handleChatCompressionEvent = useCallback( (eventValue: ServerGeminiChatCompressedEvent['value']) => addItem( @@ -501,6 +542,12 @@ export const useGeminiStream = ( case ServerGeminiEventType.MaxSessionTurns: handleMaxSessionTurnsEvent(); break; + case ServerGeminiEventType.Finished: + handleFinishedEvent( + event as ServerGeminiFinishedEvent, + userMessageTimestamp, + ); + break; case ServerGeminiEventType.LoopDetected: // handle later because we want to move pending history to history // before we add loop detected message to history @@ -524,6 +571,7 @@ export const useGeminiStream = ( handleErrorEvent, scheduleToolCalls, handleChatCompressionEvent, + handleFinishedEvent, handleMaxSessionTurnsEvent, ], ); 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);