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

Co-authored-by: Sandy Tao <sandytao520@icloud.com>
This commit is contained in:
darkcocoa
2025-07-22 06:57:11 +09:00
committed by GitHub
parent dc2ac144b7
commit 4c3532d2b3
4 changed files with 463 additions and 2 deletions

View File

@@ -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', () => {

View File

@@ -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);