Skip to content

Commit 4c3532d

Browse files
darkdarkcocoaSandyTao520
darkdarkcocoa
and
SandyTao520
authored
fix: Add warning message for token limit truncation (google-gemini#2260)
Co-authored-by: Sandy Tao <sandytao520@icloud.com>
1 parent dc2ac14 commit 4c3532d

File tree

4 files changed

+463
-2
lines changed

4 files changed

+463
-2
lines changed

‎packages/cli/src/ui/hooks/useGeminiStream.test.tsx‎

Lines changed: 237 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,12 @@ import {
1616
TrackedExecutingToolCall,
1717
TrackedCancelledToolCall,
1818
} from './useReactToolScheduler.js';
19-
import { Config, EditorType, AuthType } from '@google/gemini-cli-core';
19+
import {
20+
Config,
21+
EditorType,
22+
AuthType,
23+
GeminiEventType as ServerGeminiEventType,
24+
} from '@google/gemini-cli-core';
2025
import { Part, PartListUnion } from '@google/genai';
2126
import { UseHistoryManagerReturn } from './useHistoryManager.js';
2227
import {
@@ -1178,4 +1183,235 @@ describe('useGeminiStream', () => {
11781183
});
11791184
});
11801185
});
1186+
1187+
describe('handleFinishedEvent', () => {
1188+
it('should add info message for MAX_TOKENS finish reason', async () => {
1189+
// Setup mock to return a stream with MAX_TOKENS finish reason
1190+
mockSendMessageStream.mockReturnValue(
1191+
(async function* () {
1192+
yield {
1193+
type: ServerGeminiEventType.Content,
1194+
value: 'This is a truncated response...',
1195+
};
1196+
yield { type: ServerGeminiEventType.Finished, value: 'MAX_TOKENS' };
1197+
})(),
1198+
);
1199+
1200+
const { result } = renderHook(() =>
1201+
useGeminiStream(
1202+
new MockedGeminiClientClass(mockConfig),
1203+
[],
1204+
mockAddItem,
1205+
mockSetShowHelp,
1206+
mockConfig,
1207+
mockOnDebugMessage,
1208+
mockHandleSlashCommand,
1209+
false,
1210+
() => 'vscode' as EditorType,
1211+
() => {},
1212+
() => Promise.resolve(),
1213+
false,
1214+
() => {},
1215+
),
1216+
);
1217+
1218+
// Submit a query
1219+
await act(async () => {
1220+
await result.current.submitQuery('Generate long text');
1221+
});
1222+
1223+
// Check that the info message was added
1224+
await waitFor(() => {
1225+
expect(mockAddItem).toHaveBeenCalledWith(
1226+
{
1227+
type: 'info',
1228+
text: '⚠️ Response truncated due to token limits.',
1229+
},
1230+
expect.any(Number),
1231+
);
1232+
});
1233+
});
1234+
1235+
it('should not add message for STOP finish reason', async () => {
1236+
// Setup mock to return a stream with STOP finish reason
1237+
mockSendMessageStream.mockReturnValue(
1238+
(async function* () {
1239+
yield {
1240+
type: ServerGeminiEventType.Content,
1241+
value: 'Complete response',
1242+
};
1243+
yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
1244+
})(),
1245+
);
1246+
1247+
const { result } = renderHook(() =>
1248+
useGeminiStream(
1249+
new MockedGeminiClientClass(mockConfig),
1250+
[],
1251+
mockAddItem,
1252+
mockSetShowHelp,
1253+
mockConfig,
1254+
mockOnDebugMessage,
1255+
mockHandleSlashCommand,
1256+
false,
1257+
() => 'vscode' as EditorType,
1258+
() => {},
1259+
() => Promise.resolve(),
1260+
false,
1261+
() => {},
1262+
),
1263+
);
1264+
1265+
// Submit a query
1266+
await act(async () => {
1267+
await result.current.submitQuery('Test normal completion');
1268+
});
1269+
1270+
// Wait a bit to ensure no message is added
1271+
await new Promise((resolve) => setTimeout(resolve, 100));
1272+
1273+
// Check that no info message was added for STOP
1274+
const infoMessages = mockAddItem.mock.calls.filter(
1275+
(call) => call[0].type === 'info',
1276+
);
1277+
expect(infoMessages).toHaveLength(0);
1278+
});
1279+
1280+
it('should not add message for FINISH_REASON_UNSPECIFIED', async () => {
1281+
// Setup mock to return a stream with FINISH_REASON_UNSPECIFIED
1282+
mockSendMessageStream.mockReturnValue(
1283+
(async function* () {
1284+
yield {
1285+
type: ServerGeminiEventType.Content,
1286+
value: 'Response with unspecified finish',
1287+
};
1288+
yield {
1289+
type: ServerGeminiEventType.Finished,
1290+
value: 'FINISH_REASON_UNSPECIFIED',
1291+
};
1292+
})(),
1293+
);
1294+
1295+
const { result } = renderHook(() =>
1296+
useGeminiStream(
1297+
new MockedGeminiClientClass(mockConfig),
1298+
[],
1299+
mockAddItem,
1300+
mockSetShowHelp,
1301+
mockConfig,
1302+
mockOnDebugMessage,
1303+
mockHandleSlashCommand,
1304+
false,
1305+
() => 'vscode' as EditorType,
1306+
() => {},
1307+
() => Promise.resolve(),
1308+
false,
1309+
() => {},
1310+
),
1311+
);
1312+
1313+
// Submit a query
1314+
await act(async () => {
1315+
await result.current.submitQuery('Test unspecified finish');
1316+
});
1317+
1318+
// Wait a bit to ensure no message is added
1319+
await new Promise((resolve) => setTimeout(resolve, 100));
1320+
1321+
// Check that no info message was added
1322+
const infoMessages = mockAddItem.mock.calls.filter(
1323+
(call) => call[0].type === 'info',
1324+
);
1325+
expect(infoMessages).toHaveLength(0);
1326+
});
1327+
1328+
it('should add appropriate messages for other finish reasons', async () => {
1329+
const testCases = [
1330+
{
1331+
reason: 'SAFETY',
1332+
message: '⚠️ Response stopped due to safety reasons.',
1333+
},
1334+
{
1335+
reason: 'RECITATION',
1336+
message: '⚠️ Response stopped due to recitation policy.',
1337+
},
1338+
{
1339+
reason: 'LANGUAGE',
1340+
message: '⚠️ Response stopped due to unsupported language.',
1341+
},
1342+
{
1343+
reason: 'BLOCKLIST',
1344+
message: '⚠️ Response stopped due to forbidden terms.',
1345+
},
1346+
{
1347+
reason: 'PROHIBITED_CONTENT',
1348+
message: '⚠️ Response stopped due to prohibited content.',
1349+
},
1350+
{
1351+
reason: 'SPII',
1352+
message:
1353+
'⚠️ Response stopped due to sensitive personally identifiable information.',
1354+
},
1355+
{ reason: 'OTHER', message: '⚠️ Response stopped for other reasons.' },
1356+
{
1357+
reason: 'MALFORMED_FUNCTION_CALL',
1358+
message: '⚠️ Response stopped due to malformed function call.',
1359+
},
1360+
{
1361+
reason: 'IMAGE_SAFETY',
1362+
message: '⚠️ Response stopped due to image safety violations.',
1363+
},
1364+
{
1365+
reason: 'UNEXPECTED_TOOL_CALL',
1366+
message: '⚠️ Response stopped due to unexpected tool call.',
1367+
},
1368+
];
1369+
1370+
for (const { reason, message } of testCases) {
1371+
// Reset mocks for each test case
1372+
mockAddItem.mockClear();
1373+
mockSendMessageStream.mockReturnValue(
1374+
(async function* () {
1375+
yield {
1376+
type: ServerGeminiEventType.Content,
1377+
value: `Response for ${reason}`,
1378+
};
1379+
yield { type: ServerGeminiEventType.Finished, value: reason };
1380+
})(),
1381+
);
1382+
1383+
const { result } = renderHook(() =>
1384+
useGeminiStream(
1385+
new MockedGeminiClientClass(mockConfig),
1386+
[],
1387+
mockAddItem,
1388+
mockSetShowHelp,
1389+
mockConfig,
1390+
mockOnDebugMessage,
1391+
mockHandleSlashCommand,
1392+
false,
1393+
() => 'vscode' as EditorType,
1394+
() => {},
1395+
() => Promise.resolve(),
1396+
false,
1397+
() => {},
1398+
),
1399+
);
1400+
1401+
await act(async () => {
1402+
await result.current.submitQuery(`Test ${reason}`);
1403+
});
1404+
1405+
await waitFor(() => {
1406+
expect(mockAddItem).toHaveBeenCalledWith(
1407+
{
1408+
type: 'info',
1409+
text: message,
1410+
},
1411+
expect.any(Number),
1412+
);
1413+
});
1414+
}
1415+
});
1416+
});
11811417
});

‎packages/cli/src/ui/hooks/useGeminiStream.ts‎

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import {
1414
ServerGeminiContentEvent as ContentEvent,
1515
ServerGeminiErrorEvent as ErrorEvent,
1616
ServerGeminiChatCompressedEvent,
17+
ServerGeminiFinishedEvent,
1718
getErrorMessage,
1819
isNodeError,
1920
MessageSenderType,
@@ -26,7 +27,7 @@ import {
2627
UserPromptEvent,
2728
DEFAULT_GEMINI_FLASH_MODEL,
2829
} from '@google/gemini-cli-core';
29-
import { type Part, type PartListUnion } from '@google/genai';
30+
import { type Part, type PartListUnion, FinishReason } from '@google/genai';
3031
import {
3132
StreamingState,
3233
HistoryItem,
@@ -422,6 +423,46 @@ export const useGeminiStream = (
422423
[addItem, pendingHistoryItemRef, setPendingHistoryItem, config],
423424
);
424425

426+
const handleFinishedEvent = useCallback(
427+
(event: ServerGeminiFinishedEvent, userMessageTimestamp: number) => {
428+
const finishReason = event.value;
429+
430+
const finishReasonMessages: Record<FinishReason, string | undefined> = {
431+
[FinishReason.FINISH_REASON_UNSPECIFIED]: undefined,
432+
[FinishReason.STOP]: undefined,
433+
[FinishReason.MAX_TOKENS]: 'Response truncated due to token limits.',
434+
[FinishReason.SAFETY]: 'Response stopped due to safety reasons.',
435+
[FinishReason.RECITATION]: 'Response stopped due to recitation policy.',
436+
[FinishReason.LANGUAGE]:
437+
'Response stopped due to unsupported language.',
438+
[FinishReason.BLOCKLIST]: 'Response stopped due to forbidden terms.',
439+
[FinishReason.PROHIBITED_CONTENT]:
440+
'Response stopped due to prohibited content.',
441+
[FinishReason.SPII]:
442+
'Response stopped due to sensitive personally identifiable information.',
443+
[FinishReason.OTHER]: 'Response stopped for other reasons.',
444+
[FinishReason.MALFORMED_FUNCTION_CALL]:
445+
'Response stopped due to malformed function call.',
446+
[FinishReason.IMAGE_SAFETY]:
447+
'Response stopped due to image safety violations.',
448+
[FinishReason.UNEXPECTED_TOOL_CALL]:
449+
'Response stopped due to unexpected tool call.',
450+
};
451+
452+
const message = finishReasonMessages[finishReason];
453+
if (message) {
454+
addItem(
455+
{
456+
type: 'info',
457+
text: `⚠️ ${message}`,
458+
},
459+
userMessageTimestamp,
460+
);
461+
}
462+
},
463+
[addItem],
464+
);
465+
425466
const handleChatCompressionEvent = useCallback(
426467
(eventValue: ServerGeminiChatCompressedEvent['value']) =>
427468
addItem(
@@ -501,6 +542,12 @@ export const useGeminiStream = (
501542
case ServerGeminiEventType.MaxSessionTurns:
502543
handleMaxSessionTurnsEvent();
503544
break;
545+
case ServerGeminiEventType.Finished:
546+
handleFinishedEvent(
547+
event as ServerGeminiFinishedEvent,
548+
userMessageTimestamp,
549+
);
550+
break;
504551
case ServerGeminiEventType.LoopDetected:
505552
// handle later because we want to move pending history to history
506553
// before we add loop detected message to history
@@ -524,6 +571,7 @@ export const useGeminiStream = (
524571
handleErrorEvent,
525572
scheduleToolCalls,
526573
handleChatCompressionEvent,
574+
handleFinishedEvent,
527575
handleMaxSessionTurnsEvent,
528576
],
529577
);

0 commit comments

Comments
 (0)