mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-10 14:10:37 -07:00
feat(cli): continue request after disabling loop detection (#11416)
This commit is contained in:
@@ -2677,7 +2677,8 @@ describe('useGeminiStream', () => {
|
||||
};
|
||||
mockConfig.getGeminiClient = vi.fn().mockReturnValue(mockClient);
|
||||
|
||||
mockSendMessageStream.mockReturnValue(
|
||||
// Mock for the initial request
|
||||
mockSendMessageStream.mockReturnValueOnce(
|
||||
(async function* () {
|
||||
yield {
|
||||
type: ServerGeminiEventType.LoopDetected,
|
||||
@@ -2685,6 +2686,20 @@ describe('useGeminiStream', () => {
|
||||
})(),
|
||||
);
|
||||
|
||||
// Mock for the retry request
|
||||
mockSendMessageStream.mockReturnValueOnce(
|
||||
(async function* () {
|
||||
yield {
|
||||
type: ServerGeminiEventType.Content,
|
||||
value: 'Retry successful',
|
||||
};
|
||||
yield {
|
||||
type: ServerGeminiEventType.Finished,
|
||||
value: { reason: 'STOP' },
|
||||
};
|
||||
})(),
|
||||
);
|
||||
|
||||
const { result } = renderTestHook();
|
||||
|
||||
await act(async () => {
|
||||
@@ -2715,10 +2730,21 @@ describe('useGeminiStream', () => {
|
||||
expect(mockAddItem).toHaveBeenCalledWith(
|
||||
{
|
||||
type: 'info',
|
||||
text: 'Loop detection has been disabled for this session. Please try your request again.',
|
||||
text: 'Loop detection has been disabled for this session. Retrying request...',
|
||||
},
|
||||
expect.any(Number),
|
||||
);
|
||||
|
||||
// Verify that the request was retried
|
||||
await waitFor(() => {
|
||||
expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
|
||||
expect(mockSendMessageStream).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
'test query',
|
||||
expect.any(AbortSignal),
|
||||
expect.any(String),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it('should keep loop detection enabled and show message when user selects "keep"', async () => {
|
||||
@@ -2771,6 +2797,9 @@ describe('useGeminiStream', () => {
|
||||
},
|
||||
expect.any(Number),
|
||||
);
|
||||
|
||||
// Verify that the request was NOT retried
|
||||
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should handle multiple loop detection events properly', async () => {
|
||||
@@ -2821,6 +2850,20 @@ describe('useGeminiStream', () => {
|
||||
})(),
|
||||
);
|
||||
|
||||
// Mock for the retry request
|
||||
mockSendMessageStream.mockReturnValueOnce(
|
||||
(async function* () {
|
||||
yield {
|
||||
type: ServerGeminiEventType.Content,
|
||||
value: 'Retry successful',
|
||||
};
|
||||
yield {
|
||||
type: ServerGeminiEventType.Finished,
|
||||
value: { reason: 'STOP' },
|
||||
};
|
||||
})(),
|
||||
);
|
||||
|
||||
// Second loop detection
|
||||
await act(async () => {
|
||||
await result.current.submitQuery('second query');
|
||||
@@ -2843,10 +2886,21 @@ describe('useGeminiStream', () => {
|
||||
expect(mockAddItem).toHaveBeenCalledWith(
|
||||
{
|
||||
type: 'info',
|
||||
text: 'Loop detection has been disabled for this session. Please try your request again.',
|
||||
text: 'Loop detection has been disabled for this session. Retrying request...',
|
||||
},
|
||||
expect.any(Number),
|
||||
);
|
||||
|
||||
// Verify that the request was retried
|
||||
await waitFor(() => {
|
||||
expect(mockSendMessageStream).toHaveBeenCalledTimes(3); // 1st query, 2nd query, retry of 2nd query
|
||||
expect(mockSendMessageStream).toHaveBeenNthCalledWith(
|
||||
3,
|
||||
'second query',
|
||||
expect.any(AbortSignal),
|
||||
expect.any(String),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it('should process LoopDetected event after moving pending history to history', async () => {
|
||||
|
||||
@@ -185,6 +185,8 @@ export const useGeminiStream = (
|
||||
return undefined;
|
||||
}, [toolCalls]);
|
||||
|
||||
const lastQueryRef = useRef<PartListUnion | null>(null);
|
||||
const lastPromptIdRef = useRef<string | null>(null);
|
||||
const loopDetectedRef = useRef(false);
|
||||
const [
|
||||
loopDetectionConfirmationRequest,
|
||||
@@ -668,39 +670,6 @@ export const useGeminiStream = (
|
||||
[addItem, onCancelSubmit, config],
|
||||
);
|
||||
|
||||
const handleLoopDetectionConfirmation = useCallback(
|
||||
(result: { userSelection: 'disable' | 'keep' }) => {
|
||||
setLoopDetectionConfirmationRequest(null);
|
||||
|
||||
if (result.userSelection === 'disable') {
|
||||
config.getGeminiClient().getLoopDetectionService().disableForSession();
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `Loop detection has been disabled for this session. Please try your request again.`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
} else {
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
}
|
||||
},
|
||||
[config, addItem],
|
||||
);
|
||||
|
||||
const handleLoopDetectedEvent = useCallback(() => {
|
||||
// Show the confirmation dialog to choose whether to disable loop detection
|
||||
setLoopDetectionConfirmationRequest({
|
||||
onComplete: handleLoopDetectionConfirmation,
|
||||
});
|
||||
}, [handleLoopDetectionConfirmation]);
|
||||
|
||||
const processGeminiStreamEvents = useCallback(
|
||||
async (
|
||||
stream: AsyncIterable<GeminiEvent>,
|
||||
@@ -850,6 +819,10 @@ export const useGeminiStream = (
|
||||
setIsResponding(true);
|
||||
setInitError(null);
|
||||
|
||||
// Store query and prompt_id for potential retry on loop detection
|
||||
lastQueryRef.current = queryToSend;
|
||||
lastPromptIdRef.current = prompt_id;
|
||||
|
||||
try {
|
||||
const stream = geminiClient.sendMessageStream(
|
||||
queryToSend,
|
||||
@@ -872,7 +845,42 @@ export const useGeminiStream = (
|
||||
}
|
||||
if (loopDetectedRef.current) {
|
||||
loopDetectedRef.current = false;
|
||||
handleLoopDetectedEvent();
|
||||
// Show the confirmation dialog to choose whether to disable loop detection
|
||||
setLoopDetectionConfirmationRequest({
|
||||
onComplete: (result: { userSelection: 'disable' | 'keep' }) => {
|
||||
setLoopDetectionConfirmationRequest(null);
|
||||
|
||||
if (result.userSelection === 'disable') {
|
||||
config
|
||||
.getGeminiClient()
|
||||
.getLoopDetectionService()
|
||||
.disableForSession();
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `Loop detection has been disabled for this session. Retrying request...`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
|
||||
if (lastQueryRef.current && lastPromptIdRef.current) {
|
||||
submitQuery(
|
||||
lastQueryRef.current,
|
||||
{ isContinuation: true },
|
||||
lastPromptIdRef.current,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
addItem(
|
||||
{
|
||||
type: 'info',
|
||||
text: `A potential loop was detected. This can happen due to repetitive tool calls or other model behavior. The request has been halted.`,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
} catch (error: unknown) {
|
||||
if (error instanceof UnauthorizedError) {
|
||||
@@ -911,7 +919,6 @@ export const useGeminiStream = (
|
||||
config,
|
||||
startNewPrompt,
|
||||
getPromptCount,
|
||||
handleLoopDetectedEvent,
|
||||
],
|
||||
);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user