mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-19 09:33:53 +00:00
Do not call nextSpeakerCheck if there was an error processing the stream. (#7048)
Co-authored-by: christine betts <chrstn@uw.edu> Co-authored-by: Antonio Scandurra <me@as-cii.com> Co-authored-by: Arya Gummadi <aryagummadi@google.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Shreya Keshive <skeshive@gmail.com> Co-authored-by: Abhi <43648792+abhipatel12@users.noreply.github.com> Co-authored-by: shishu314 <shishu_1998@yahoo.com> Co-authored-by: Shi Shu <shii@google.com> Co-authored-by: Steven <steventohme59@gmail.com> Co-authored-by: Pascal Birchler <pascalb@google.com> Co-authored-by: N. Taylor Mullen <ntaylormullen@google.com>
This commit is contained in:
@@ -51,6 +51,8 @@ vi.mock('./turn', () => {
|
||||
GeminiEventType: {
|
||||
MaxSessionTurns: 'MaxSessionTurns',
|
||||
ChatCompressed: 'ChatCompressed',
|
||||
Error: 'error',
|
||||
Content: 'content',
|
||||
},
|
||||
};
|
||||
});
|
||||
@@ -1887,6 +1889,89 @@ ${JSON.stringify(
|
||||
expect(JSON.stringify(finalCall)).toContain('fileC.ts');
|
||||
});
|
||||
});
|
||||
|
||||
it('should not call checkNextSpeaker when turn.run() yields an error', async () => {
|
||||
// Arrange
|
||||
const { checkNextSpeaker } = await import(
|
||||
'../utils/nextSpeakerChecker.js'
|
||||
);
|
||||
const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker);
|
||||
|
||||
const mockStream = (async function* () {
|
||||
yield {
|
||||
type: GeminiEventType.Error,
|
||||
value: { error: { message: 'test error' } },
|
||||
};
|
||||
})();
|
||||
mockTurnRunFn.mockReturnValue(mockStream);
|
||||
|
||||
const mockChat: Partial<GeminiChat> = {
|
||||
addHistory: vi.fn(),
|
||||
getHistory: vi.fn().mockReturnValue([]),
|
||||
};
|
||||
client['chat'] = mockChat as GeminiChat;
|
||||
|
||||
const mockGenerator: Partial<ContentGenerator> = {
|
||||
countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
|
||||
generateContent: mockGenerateContentFn,
|
||||
};
|
||||
client['contentGenerator'] = mockGenerator as ContentGenerator;
|
||||
|
||||
// Act
|
||||
const stream = client.sendMessageStream(
|
||||
[{ text: 'Hi' }],
|
||||
new AbortController().signal,
|
||||
'prompt-id-error',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
// consume stream
|
||||
}
|
||||
|
||||
// Assert
|
||||
expect(mockCheckNextSpeaker).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should not call checkNextSpeaker when turn.run() yields a value then an error', async () => {
|
||||
// Arrange
|
||||
const { checkNextSpeaker } = await import(
|
||||
'../utils/nextSpeakerChecker.js'
|
||||
);
|
||||
const mockCheckNextSpeaker = vi.mocked(checkNextSpeaker);
|
||||
|
||||
const mockStream = (async function* () {
|
||||
yield { type: GeminiEventType.Content, value: 'some content' };
|
||||
yield {
|
||||
type: GeminiEventType.Error,
|
||||
value: { error: { message: 'test error' } },
|
||||
};
|
||||
})();
|
||||
mockTurnRunFn.mockReturnValue(mockStream);
|
||||
|
||||
const mockChat: Partial<GeminiChat> = {
|
||||
addHistory: vi.fn(),
|
||||
getHistory: vi.fn().mockReturnValue([]),
|
||||
};
|
||||
client['chat'] = mockChat as GeminiChat;
|
||||
|
||||
const mockGenerator: Partial<ContentGenerator> = {
|
||||
countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }),
|
||||
generateContent: mockGenerateContentFn,
|
||||
};
|
||||
client['contentGenerator'] = mockGenerator as ContentGenerator;
|
||||
|
||||
// Act
|
||||
const stream = client.sendMessageStream(
|
||||
[{ text: 'Hi' }],
|
||||
new AbortController().signal,
|
||||
'prompt-id-error',
|
||||
);
|
||||
for await (const _ of stream) {
|
||||
// consume stream
|
||||
}
|
||||
|
||||
// Assert
|
||||
expect(mockCheckNextSpeaker).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('generateContent', () => {
|
||||
|
||||
@@ -513,6 +513,9 @@ export class GeminiClient {
|
||||
return turn;
|
||||
}
|
||||
yield event;
|
||||
if (event.type === GeminiEventType.Error) {
|
||||
return turn;
|
||||
}
|
||||
}
|
||||
if (!turn.pendingToolCalls.length && signal && !signal.aborted) {
|
||||
// Check if model was switched during the call (likely due to quota error)
|
||||
|
||||
Reference in New Issue
Block a user