skills/google-gemini__gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-09-18 10:48:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

(FIX) history rollback and retry nudge optimizations (#28934)

Browse Source
This commit is contained in:
David Pierce
2026-08-24 18:41:54 +00:00
committed by GitHub
parent af687798af
commit 812f7a2bcf
9 changed files with 695 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/cli/src/nonInteractiveCliAgentSession.test.ts
+11 -2
View File
@@ -841,9 +841,18 @@ describe('runNonInteractive', () => {
},
];
// Third call handles the auto-nudge recovery (when no response is received after tools)
const thirdCallEvents: ServerGeminiStreamEvent[] = [
{
type: GeminiEventType.Finished,
value: { reason: undefined, usageMetadata: { totalTokenCount: 1 } },
},
];
mockGeminiClient.sendMessageStream
.mockReturnValueOnce(createStreamFromEvents(firstCallEvents))
.mockReturnValueOnce(createStreamFromEvents(secondCallEvents));
.mockReturnValueOnce(createStreamFromEvents(secondCallEvents))
.mockReturnValueOnce(createStreamFromEvents(thirdCallEvents));
vi.mocked(mockConfig.getOutputFormat).mockReturnValue(OutputFormat.JSON);
vi.spyOn(uiTelemetryService, 'getMetrics').mockReturnValue(
@@ -857,7 +866,7 @@ describe('runNonInteractive', () => {
prompt_id: 'prompt-id-tool-only',
});
expect(mockGeminiClient.sendMessageStream).toHaveBeenCalledTimes(2);
expect(mockGeminiClient.sendMessageStream).toHaveBeenCalledTimes(3);
expect(mockSchedulerSchedule).toHaveBeenCalledWith(
[expect.objectContaining({ name: 'testTool' })],
expect.any(AbortSignal),
packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+157 -30
View File
@@ -86,6 +86,11 @@ const MockedGeminiClientClass = vi.hoisted(() =>
this.startChat = mockStartChat;
this.sendMessageStream = mockSendMessageStream;
this.addHistory = vi.fn();
let mockHistory: any[] = [];
this.getHistory = vi.fn().mockImplementation(() => mockHistory);
this.setHistory = vi.fn().mockImplementation((newHistory: any[]) => {
mockHistory = [...newHistory];
});
this.generateContent = vi.fn().mockResolvedValue({
candidates: [
{ content: { parts: [{ text: 'Got it. Focusing on tests only.' }] } },
@@ -761,6 +766,15 @@ describe('useGeminiStream', () => {
];
});
mockSendMessageStream.mockReturnValue(
(async function* () {
yield {
type: ServerGeminiEventType.Content,
value: 'Visible response text',
};
})(),
);
await renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
@@ -927,7 +941,106 @@ describe('useGeminiStream', () => {
});
});
it('should handle all tool calls being cancelled', async () => {
it('should auto-nudge the model when tool execution succeeds but model stream is empty', async () => {
const toolCallResponseParts: Part[] = [{ text: 'tool final response' }];
const completedToolCalls: TrackedToolCall[] = [
{
request: {
callId: 'call1',
name: 'tool1',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-ack',
},
status: 'success',
responseSubmittedToGemini: false,
response: {
callId: 'call1',
responseParts: toolCallResponseParts,
errorType: undefined,
},
tool: {
displayName: 'MockTool',
},
invocation: {
getDescription: () => `Mock description`,
} as unknown as AnyToolInvocation,
} as TrackedCompletedToolCall,
];
let callCount = 0;
mockSendMessageStream.mockImplementation(() => {
callCount += 1;
if (callCount === 1) {
return (async function* () {})();
} else {
return (async function* () {
yield {
type: ServerGeminiEventType.Content,
value:
'I have analyzed the empty response. Here is the final answer.',
};
})();
}
});
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
| null = null;
mockUseToolScheduler.mockImplementation((onComplete) => {
capturedOnComplete = onComplete;
return [
[],
mockScheduleToolCalls,
mockMarkToolsAsSubmitted,
vi.fn(),
mockCancelAllToolCalls,
0,
];
});
await renderHookWithProviders(() =>
useGeminiStream(
new MockedGeminiClientClass(mockConfig),
[],
mockAddItem,
mockConfig,
mockLoadedSettings,
mockOnDebugMessage,
mockHandleSlashCommand,
false,
() => 'vscode' as EditorType,
() => {},
() => Promise.resolve(),
false,
() => {},
() => {},
() => {},
80,
24,
undefined,
() => 'focus on tests only',
),
);
await act(async () => {
if (capturedOnComplete) {
await new Promise((resolve) => setTimeout(resolve, 0));
await capturedOnComplete(completedToolCalls);
}
});
await waitFor(() => {
expect(mockSendMessageStream).toHaveBeenCalledTimes(2);
});
const sentParts = mockSendMessageStream.mock.calls[1][0] as Part[];
expect(sentParts[0].text).toContain(
'[System: You successfully executed a tool but returned an empty response. Please analyze the tool output and explain your progress or final answer.]',
);
});
it('should handle all tool calls being cancelled by rolling back the history', async () => {
const cancelledToolCalls: TrackedToolCall[] = [
{
request: {
@@ -977,6 +1090,7 @@ describe('useGeminiStream', () => {
} as any,
];
const client = new MockedGeminiClientClass(mockConfig);
client.setHistory([{ role: 'user', parts: [{ text: 'User prompt' }] }]);
// Capture the onComplete callback
let capturedOnComplete:
@@ -995,7 +1109,7 @@ describe('useGeminiStream', () => {
];
});
await renderHookWithProviders(() =>
const { result } = await renderHookWithProviders(() =>
useGeminiStream(
client,
[],
@@ -1017,6 +1131,21 @@ describe('useGeminiStream', () => {
),
);
// Call submitQuery to populate the user turn and set historyLengthAfterUserPromptRef
await act(async () => {
// eslint-disable-next-line @typescript-eslint/no-floating-promises
result.current.submitQuery('User prompt');
});
// Model issues a functionCall request, which appends to history
client.setHistory([
{ role: 'user', parts: [{ text: 'User prompt' }] },
{
role: 'model',
parts: [{ functionCall: { name: 'testTool', args: {} } }],
},
]);
// Trigger the onComplete callback with cancelled tools
await act(async () => {
if (capturedOnComplete) {
@@ -1028,21 +1157,12 @@ describe('useGeminiStream', () => {
await waitFor(() => {
expect(mockMarkToolsAsSubmitted).toHaveBeenCalledWith(['topic1', '1']);
expect(client.addHistory).toHaveBeenCalledWith({
role: 'user',
parts: [
{
functionResponse: {
name: UPDATE_TOPIC_TOOL_NAME,
id: 'topic1',
response: {},
},
},
{ text: CoreToolCallStatus.Cancelled },
],
});
// Ensure we do NOT call back to the API
expect(mockSendMessageStream).not.toHaveBeenCalled();
// Should NOT have appended cancellations via addHistory
expect(client.addHistory).not.toHaveBeenCalled();
// Should have rolled history back to pre-model length (1)
expect(client.getHistory().length).toBe(1);
// Ensure we do NOT call back to the API a second time (only the initial user turn was sent)
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
});
});
@@ -1375,7 +1495,7 @@ describe('useGeminiStream', () => {
expect(noteIndex).toBeLessThan(stopIndex);
});
it('should group multiple cancelled tool call responses into a single history entry', async () => {
it('should rollback multiple cancelled tool calls rather than appending them to history', async () => {
const cancelledToolCall1: TrackedCancelledToolCall = {
request: {
callId: 'cancel-1',
@@ -1436,6 +1556,7 @@ describe('useGeminiStream', () => {
};
const allCancelledTools = [cancelledToolCall1, cancelledToolCall2];
const client = new MockedGeminiClientClass(mockConfig);
client.setHistory([{ role: 'user', parts: [{ text: 'User prompt' }] }]);
let capturedOnComplete:
| ((completedTools: TrackedToolCall[]) => Promise<void>)
@@ -1453,7 +1574,7 @@ describe('useGeminiStream', () => {
];
});
await renderHookWithProviders(() =>
const { result } = await renderHookWithProviders(() =>
useGeminiStream(
client,
[],
@@ -1475,6 +1596,18 @@ describe('useGeminiStream', () => {
),
);
// Call submitQuery to populate the user turn and set historyLengthAfterUserPromptRef
await act(async () => {
// eslint-disable-next-line @typescript-eslint/no-floating-promises
result.current.submitQuery('User prompt');
});
// Model issues model turns, which appends to history
client.setHistory([
{ role: 'user', parts: [{ text: 'User prompt' }] },
{ role: 'model', parts: [{ functionCall: { name: 'toolA', args: {} } }] },
]);
// Trigger the onComplete callback with multiple cancelled tools
await act(async () => {
if (capturedOnComplete) {
@@ -1491,20 +1624,14 @@ describe('useGeminiStream', () => {
'cancel-2',
]);
// Crucially, addHistory should be called only ONCE
expect(client.addHistory).toHaveBeenCalledTimes(1);
// Crucially, addHistory should NOT be called
expect(client.addHistory).not.toHaveBeenCalled();
// And that single call should contain BOTH function responses
expect(client.addHistory).toHaveBeenCalledWith({
role: 'user',
parts: [
...cancelledToolCall1.response.responseParts,
...cancelledToolCall2.response.responseParts,
],
});
// Instead, history should be rolled back to pre-model length (1)
expect(client.getHistory().length).toBe(1);
// No message should be sent back to the API for a turn with only cancellations
expect(mockSendMessageStream).not.toHaveBeenCalled();
expect(mockSendMessageStream).toHaveBeenCalledTimes(1);
});
});
packages/cli/src/ui/hooks/useGeminiStream.ts
+51 -11
View File
@@ -258,6 +258,7 @@ export const useGeminiStream = (
const abortControllerRef = useRef<AbortController | null>(null);
const turnCancelledRef = useRef(false);
const activeQueryIdRef = useRef<string | null>(null);
const historyLengthAfterUserPromptRef = useRef<number | undefined>(undefined);
const previousApprovalModeRef = useRef<ApprovalMode>(
config.getApprovalMode(),
);
@@ -717,6 +718,16 @@ export const useGeminiStream = (
const lastQueryRef = useRef<PartListUnion | null>(null);
const lastPromptIdRef = useRef<string | null>(null);
const loopDetectedRef = useRef(false);
const autoNudgeAttemptCountRef = useRef<number>(0);
const MAX_AUTO_NUDGE_ATTEMPTS = 2;
const submitQueryRef = useRef<
| ((
query: PartListUnion,
options?: { isContinuation: boolean },
prompt_id?: string,
) => Promise<void>)
| null
>(null);
const [
loopDetectionConfirmationRequest,
setLoopDetectionConfirmationRequest,
@@ -1524,6 +1535,8 @@ export const useGeminiStream = (
stream: AsyncIterable<GeminiEvent>,
userMessageTimestamp: number,
signal: AbortSignal,
isContinuation?: boolean,
prompt_id?: string,
): Promise<StreamProcessingStatus> => {
let geminiMessageBuffer = '';
const toolCallRequests: ToolCallRequestInfo[] = [];
@@ -1622,6 +1635,25 @@ export const useGeminiStream = (
setPendingHistoryItem(null);
}
await scheduleToolCalls(toolCallRequests, signal);
} else {
const hasVisibleText = geminiMessageBuffer.trim().length > 0;
if (
isContinuation &&
!hasVisibleText &&
autoNudgeAttemptCountRef.current < MAX_AUTO_NUDGE_ATTEMPTS
) {
autoNudgeAttemptCountRef.current += 1;
const nudgeMessage =
'[System: You successfully executed a tool but returned an empty response. Please analyze the tool output and explain your progress or final answer.]';
// Automatically continue the query with the nudge
// eslint-disable-next-line @typescript-eslint/no-floating-promises
submitQueryRef.current?.(
[{ text: nudgeMessage }],
{ isContinuation: true },
prompt_id,
);
}
}
return StreamProcessingStatus.Completed;
},
@@ -1677,6 +1709,7 @@ export const useGeminiStream = (
// Reset quota error flag when starting a new query (not a continuation)
if (!options?.isContinuation) {
autoNudgeAttemptCountRef.current = 0;
setModelSwitchedFromQuotaError(false);
config.setQuotaErrorOccurred(false);
config.resetBillingTurnState(
@@ -1706,6 +1739,11 @@ export const useGeminiStream = (
return;
}
if (geminiClient) {
historyLengthAfterUserPromptRef.current =
geminiClient.getHistory().length;
}
if (!options?.isContinuation) {
if (typeof queryToSend === 'string') {
// logging the text prompts only for now
@@ -1743,6 +1781,8 @@ export const useGeminiStream = (
stream,
userMessageTimestamp,
abortSignal,
options?.isContinuation,
prompt_id,
);
if (processingStatus === StreamProcessingStatus.UserCancelled) {
@@ -1845,6 +1885,7 @@ export const useGeminiStream = (
setIsResponding,
],
);
submitQueryRef.current = submitQuery;
const handleApprovalModeChange = useCallback(
async (newApprovalMode: ApprovalMode) => {
@@ -2086,17 +2127,16 @@ export const useGeminiStream = (
}
setIsResponding(false);
if (geminiClient) {
// We need to manually add the function responses to the history
// so the model knows the tools were cancelled.
const combinedParts = geminiTools.flatMap(
(toolCall) => toolCall.response.responseParts,
);
// eslint-disable-next-line @typescript-eslint/no-floating-promises
geminiClient.addHistory({
role: 'user',
parts: combinedParts,
});
if (
geminiClient &&
historyLengthAfterUserPromptRef.current !== undefined
) {
const targetLength = historyLengthAfterUserPromptRef.current;
if (geminiClient.getHistory().length > targetLength) {
geminiClient.setHistory(
geminiClient.getHistory().slice(0, targetLength),
);
}
}
const callIdsToMarkAsSubmitted = geminiTools.map(
packages/core/src/agent/legacy-agent-session.test.ts
+84
View File
@@ -601,6 +601,90 @@ describe('LegacyAgentSession', () => {
expect(streamEnd?.reason).toBe('failed');
expect(sendMock).toHaveBeenCalledTimes(1);
});
it('nudges the model if it returns an empty response after a tool execution', async () => {
const sendMock = deps.client.sendMessageStream as ReturnType<
typeof vi.fn
>;
// First turn: model requests a tool
sendMock.mockReturnValueOnce(
makeStream([
{
type: GeminiEventType.ToolCallRequest,
value: makeToolRequest('call-1', 'read_file'),
},
{
type: GeminiEventType.Finished,
value: { reason: FinishReason.STOP, usageMetadata: undefined },
},
]),
);
// Second turn: model goes silent (returns empty text response, no tool calls)
sendMock.mockReturnValueOnce(
makeStream([
{
type: GeminiEventType.Finished,
value: { reason: FinishReason.STOP, usageMetadata: undefined },
},
]),
);
// Third turn (retry after nudge): model finally provides final answer
sendMock.mockReturnValueOnce(
makeStream([
{ type: GeminiEventType.Content, value: 'Analysis completed.' },
{
type: GeminiEventType.Finished,
value: { reason: FinishReason.STOP, usageMetadata: undefined },
},
]),
);
const scheduleMock = deps.scheduler.schedule as ReturnType<typeof vi.fn>;
scheduleMock.mockResolvedValueOnce([
makeCompletedToolCall('call-1', 'read_file', 'file contents'),
]);
const session = new LegacyAgentSession(deps);
await session.send(makeMessageSend('read a file'));
const events = await collectEvents(session);
const types = events.map((e) => e.type);
expect(types).toContain('tool_request');
expect(types).toContain('tool_response');
expect(types).toContain('agent_end');
// The nudged turn should have called sendMessageStream a third time!
expect(sendMock).toHaveBeenCalledTimes(3);
// Verify that the third call received the correct system nudge message
expect(sendMock).toHaveBeenLastCalledWith(
[
{
text: '[System: You successfully executed a tool but returned an empty response. Please analyze the tool output and explain your progress or final answer.]',
},
],
expect.any(AbortSignal),
expect.any(String),
undefined,
undefined,
);
const messages = events.filter(
(e): e is AgentEvent<'message'> =>
e.type === 'message' && e.role === 'agent',
);
expect(
messages.some(
(m) =>
m.content[0] &&
'text' in m.content[0] &&
m.content[0].text === 'Analysis completed.',
),
).toBe(true);
});
});
describe('stream - terminal events', () => {
packages/core/src/agent/legacy-agent-session.ts
+19
View File
@@ -179,6 +179,7 @@ export class LegacyAgentProtocol implements AgentProtocol {
let currentDisplayContent = initialDisplayContent;
let turnCount = 0;
const maxTurns = this._config.getMaxSessionTurns();
let isAfterToolResponse = false;
while (true) {
turnCount++;
@@ -193,6 +194,8 @@ export class LegacyAgentProtocol implements AgentProtocol {
const toolCallRequests: ToolCallRequestInfo[] = [];
let finishedReason: FinishReason | undefined = undefined;
let hasVisibleText = false;
const responseStream = this._client.sendMessageStream(
currentParts,
this._abortController.signal,
@@ -212,6 +215,12 @@ export class LegacyAgentProtocol implements AgentProtocol {
toolCallRequests.push(event.value);
}
if (event.type === GeminiEventType.Content) {
if (typeof event.value === 'string' && event.value.trim() !== '') {
hasVisibleText = true;
}
}
this._emit(translateEvent(event, this._translationState));
switch (event.type) {
@@ -239,6 +248,15 @@ export class LegacyAgentProtocol implements AgentProtocol {
}
if (toolCallRequests.length === 0) {
if (isAfterToolResponse && !hasVisibleText) {
const nudgeMessage =
'[System: You successfully executed a tool but returned an empty response. Please analyze the tool output and explain your progress or final answer.]';
currentParts = [{ text: nudgeMessage }];
isAfterToolResponse = false;
continue;
}
if (finishedReason !== undefined) {
this._finishStream(mapFinishReason(finishedReason));
} else {
@@ -321,6 +339,7 @@ export class LegacyAgentProtocol implements AgentProtocol {
}
currentParts = toolResponseParts;
isAfterToolResponse = completedToolCalls.length > 0;
}
}
packages/core/src/core/geminiChat.test.ts
+194 -5
View File
@@ -24,6 +24,9 @@ import {
type HistoryTurn,
coalesceConsecutiveRoles,
stripThoughts,
THINKING_ONLY_NUDGE_MESSAGE,
NO_RESPONSE_TEXT_NUDGE_MESSAGE,
applyRetryNudge,
} from './geminiChat.js';
import {
type CompletedToolCall,
@@ -849,7 +852,7 @@ describe('GeminiChat', () => {
})(),
).resolves.not.toThrow();
// Verify history now ends with a successful model turn (with empty parts)
// Verify history now ends with a successful model turn containing the empty parts array
const lastTurn = chat.agentHistory.get()[chat.agentHistory.length - 1];
expect(lastTurn.content.role).toBe('model');
expect(lastTurn.content.parts).toEqual([]);
@@ -2717,7 +2720,7 @@ describe('GeminiChat', () => {
);
});
it('should append nudge message to systemInstruction on retry when InvalidStreamError occurs', async () => {
it('should append nudge message on retry when InvalidStreamError occurs without altering systemInstruction', async () => {
vi.mocked(mockContentGenerator.generateContentStream)
.mockImplementationOnce(async () =>
(async function* () {
@@ -2779,22 +2782,116 @@ describe('GeminiChat', () => {
LlmRole.MAIN,
);
// Second call (retry) should have nudge message appended to systemInstruction
// Second call (retry) should preserve systemInstruction and append nudge to contents
expect(
mockContentGenerator.generateContentStream,
).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
config: expect.objectContaining({
systemInstruction:
'Initial instruction\n[System: You previously generated thoughts but failed to provide a final user-facing response. Please ensure you provide your final answer or call a tool now.]',
systemInstruction: 'Initial instruction',
}),
contents: [
expect.objectContaining({
role: 'user',
parts: [
{ text: 'test' },
{ text: '\n' + THINKING_ONLY_NUDGE_MESSAGE },
],
}),
],
}),
'prompt-id-retry-nudge',
LlmRole.MAIN,
);
});
it('should re-apply nudge message on retry if a BeforeModel hook returns modifiedContents', async () => {
vi.mocked(mockConfig.getEnableHooks).mockReturnValue(true);
const modifiedHookContents: Content[] = [
{
role: 'user',
parts: [{ text: 'hook-modified-prompt' }],
},
];
const mockHookSystem = {
fireBeforeModelEvent: vi.fn().mockResolvedValue({
blocked: false,
modifiedContents: modifiedHookContents,
}),
fireAfterModelEvent: vi.fn().mockResolvedValue({ response: {} }),
fireBeforeToolSelectionEvent: vi.fn().mockResolvedValue({}),
} as unknown as HookSystem;
mockConfig.getHookSystem = vi.fn().mockReturnValue(mockHookSystem);
vi.mocked(mockContentGenerator.generateContentStream)
.mockImplementationOnce(async () =>
(async function* () {
yield {
candidates: [
{
content: {
role: 'model',
parts: [{ thought: true, text: 'thinking...' }],
},
finishReason: 'STOP',
},
],
} as unknown as GenerateContentResponse;
})(),
)
.mockImplementationOnce(async () =>
(async function* () {
yield {
candidates: [
{
content: { parts: [{ text: 'valid response' }] },
finishReason: 'STOP',
},
],
} as unknown as GenerateContentResponse;
})(),
);
const stream = await chat.sendMessageStream(
{ model: 'gemini-2.5-pro' },
'original-test-prompt',
'prompt-id-retry-hook-modified',
new AbortController().signal,
LlmRole.MAIN,
);
for await (const _ of stream) {
// consume
}
expect(mockContentGenerator.generateContentStream).toHaveBeenCalledTimes(
2,
);
// The second call (retry) should have hook-modified contents WITH the nudge message appended!
expect(
mockContentGenerator.generateContentStream,
).toHaveBeenNthCalledWith(
2,
expect.objectContaining({
contents: [
expect.objectContaining({
role: 'user',
parts: [
{ text: 'hook-modified-prompt' },
{ text: '\n' + THINKING_ONLY_NUDGE_MESSAGE },
],
}),
],
}),
'prompt-id-retry-hook-modified',
LlmRole.MAIN,
);
});
it('should fail after all retries on persistent invalid content and report metrics', async () => {
vi.mocked(mockContentGenerator.generateContentStream).mockImplementation(
async () =>
@@ -4854,4 +4951,96 @@ describe('GeminiChat', () => {
expect(history[1].role).toBe('model');
});
});
describe('applyRetryNudge', () => {
it('should return original contents if nudge message is empty', () => {
const contents: Content[] = [
{ role: 'user', parts: [{ text: 'Hello' }] },
];
const result = applyRetryNudge(contents, '');
expect(result).toEqual(contents);
});
it('should append THINKING_ONLY_NUDGE_MESSAGE to the final user turn', () => {
const contents: Content[] = [
{ role: 'user', parts: [{ text: 'Hello' }] },
];
const result = applyRetryNudge(contents, THINKING_ONLY_NUDGE_MESSAGE);
expect(result).toHaveLength(1);
expect(result[0].parts).toHaveLength(2);
expect(result[0].parts![0].text).toBe('Hello');
expect(result[0].parts![1].text).toBe('\n' + THINKING_ONLY_NUDGE_MESSAGE);
});
it('should append NO_RESPONSE_TEXT_NUDGE_MESSAGE to the final user turn', () => {
const contents: Content[] = [
{ role: 'user', parts: [{ text: 'Hello' }] },
];
const result = applyRetryNudge(contents, NO_RESPONSE_TEXT_NUDGE_MESSAGE);
expect(result).toHaveLength(1);
expect(result[0].parts).toHaveLength(2);
expect(result[0].parts![0].text).toBe('Hello');
expect(result[0].parts![1].text).toBe(
'\n' + NO_RESPONSE_TEXT_NUDGE_MESSAGE,
);
});
it('should create a new user turn if history is empty', () => {
const result = applyRetryNudge([], THINKING_ONLY_NUDGE_MESSAGE);
expect(result).toHaveLength(1);
expect(result[0].role).toBe('user');
expect(result[0].parts).toEqual([{ text: THINKING_ONLY_NUDGE_MESSAGE }]);
});
it('should create a new user turn if the last turn is from model', () => {
const contents: Content[] = [
{ role: 'model', parts: [{ text: 'AI response' }] },
];
const result = applyRetryNudge(contents, NO_RESPONSE_TEXT_NUDGE_MESSAGE);
expect(result).toHaveLength(2);
expect(result[0]).toEqual(contents[0]);
expect(result[1].role).toBe('user');
expect(result[1].parts).toEqual([
{ text: NO_RESPONSE_TEXT_NUDGE_MESSAGE },
]);
});
it('should insert synthetic model turn and dedicated user turn if the last turn is user with functionResponse', () => {
const contents: Content[] = [
{
role: 'user',
parts: [
{
functionResponse: {
name: 'Edit',
response: { result: 'success' },
},
},
],
},
];
const result = applyRetryNudge(contents, NO_RESPONSE_TEXT_NUDGE_MESSAGE);
expect(result).toHaveLength(3);
expect(result[0]).toEqual(contents[0]);
expect(result[1].role).toBe('model');
expect(result[1].parts).toEqual([
{ text: '[Tool execution completed.]' },
]);
expect(result[2].role).toBe('user');
expect(result[2].parts).toEqual([
{ text: NO_RESPONSE_TEXT_NUDGE_MESSAGE },
]);
});
it('should not duplicate the nudge message if it is already present in contents', () => {
const contents: Content[] = [
{
role: 'user',
parts: [{ text: 'Hello\n' + THINKING_ONLY_NUDGE_MESSAGE }],
},
];
const result = applyRetryNudge(contents, THINKING_ONLY_NUDGE_MESSAGE);
expect(result).toEqual(contents);
});
});
});
packages/core/src/core/geminiChat.ts
+133 -42
View File
@@ -241,6 +241,69 @@ function extractCuratedHistory(
return curatedHistory;
}
/**
* Nudge message used during retry after an empty response with thoughts.
*/
export const THINKING_ONLY_NUDGE_MESSAGE =
'[System: You previously generated thoughts but failed to provide a final user-facing response. Please ensure you provide your final answer or call a tool now.]';
/**
* Nudge message used during retry after an empty response with no text or thoughts.
*/
export const NO_RESPONSE_TEXT_NUDGE_MESSAGE =
'[System: You previously returned an empty response with no text or thoughts. Please ensure you provide your final answer or call a tool now.]';
/**
* Appends an on-retry nudge message to the final user turn in contents (or adds a user turn)
* so that the model observes it at the end of the context window without altering systemInstruction.
*/
export function applyRetryNudge(
contents: Content[],
nudgeMessage: string,
): Content[] {
if (!nudgeMessage) {
return contents;
}
const lastTurn = contents[contents.length - 1];
const hasNudge = lastTurn?.parts?.some((p) => p.text?.includes(nudgeMessage));
if (hasNudge) {
return contents;
}
const cloned: Content[] = contents.map((c) => ({
...c,
parts: c.parts ? [...c.parts] : [],
}));
const clonedLastTurn = cloned[cloned.length - 1];
const hasFunctionResponse = clonedLastTurn?.parts?.some(
(p) => p.functionResponse,
);
if (clonedLastTurn?.role === 'user' && hasFunctionResponse) {
// Satisfy strict role alternation invariants of the Gemini API by inserting
// a neutral, synthetic model turn between the tool response and the nudge prompt.
cloned.push({
role: 'model',
parts: [{ text: '[Tool execution completed.]' }],
});
cloned.push({
role: 'user',
parts: [{ text: nudgeMessage }],
});
} else if (clonedLastTurn?.role === 'user') {
if (!clonedLastTurn.parts) {
clonedLastTurn.parts = [];
}
clonedLastTurn.parts.push({ text: `\n${nudgeMessage}` });
} else {
cloned.push({
role: 'user',
parts: [{ text: nudgeMessage }],
});
}
return cloned;
}
/**
* Custom error to signal that a stream completed with invalid content,
* which should trigger a retry.
@@ -898,27 +961,20 @@ export class GeminiChat {
abortSignal,
};
// Apply Context-Aware Retries (On-Retry Nudging) to guide the model out of silent loops
// Apply Context-Aware Retries (On-Retry Nudging) to guide the model out of silent loops.
// The nudge message is appended to the contents array (end of conversation) rather than modifying
// systemInstruction. This preserves the prefix cache and ensures the nudge is directly observed
// by the model at the end of the context window.
let nudgeMessage = '';
if (
modelConfigKey.isRetry &&
modelConfigKey.lastStreamError instanceof InvalidStreamError
) {
const lastError = modelConfigKey.lastStreamError;
let nudgeMessage = '';
if (lastError.type === 'THINKING_ONLY_RESPONSE') {
nudgeMessage =
'\n[System: You previously generated thoughts but failed to provide a final user-facing response. Please ensure you provide your final answer or call a tool now.]';
nudgeMessage = THINKING_ONLY_NUDGE_MESSAGE;
} else if (lastError.type === 'NO_RESPONSE_TEXT') {
nudgeMessage =
'\n[System: You previously returned an empty response with no text or thoughts. Please ensure you provide your final answer or call a tool now.]';
}
if (nudgeMessage) {
if (typeof config.systemInstruction === 'string') {
config.systemInstruction += nudgeMessage;
} else if (config.systemInstruction === undefined) {
config.systemInstruction = nudgeMessage;
}
nudgeMessage = NO_RESPONSE_TEXT_NUDGE_MESSAGE;
}
}
@@ -927,6 +983,10 @@ export class GeminiChat {
? [...contentsForPreviewModel]
: [...requestContents];
if (nudgeMessage) {
contentsToUse = applyRetryNudge(contentsToUse, nudgeMessage);
}
const hookSystem = this.context.config.getHookSystem();
if (hookSystem) {
const beforeModelResult = await hookSystem.fireBeforeModelEvent({
@@ -971,6 +1031,9 @@ export class GeminiChat {
supportsModernFeatures(modelToUse) || isGemini2Model(modelToUse)
? [...contentsForPreviewModel]
: [...requestContents];
if (nudgeMessage) {
contentsToUse = applyRetryNudge(contentsToUse, nudgeMessage);
}
}
if (beforeModelResult.modifiedConfig) {
Object.assign(config, beforeModelResult.modifiedConfig);
@@ -981,6 +1044,9 @@ export class GeminiChat {
) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
contentsToUse = beforeModelResult.modifiedContents as Content[];
if (nudgeMessage) {
contentsToUse = applyRetryNudge(contentsToUse, nudgeMessage);
}
}
const toolSelectionResult =
@@ -1578,10 +1644,10 @@ export class GeminiChat {
}
}
let id: string;
// Record model response text from the collected parts.
// Also flush when there are thoughts or a tool call (even with no text)
// so that BeforeTool hooks always see the latest transcript state.
let id: string;
if (responseText || hasThoughts || hasToolCall) {
id = this.chatRecordingService.recordMessage({
model,
@@ -1688,35 +1754,60 @@ export function isInvalidArgumentError(errorMessage: string): boolean {
}
export function stripToolCallIdPrefixes(contents: Content[]): Content[] {
return contents.map((content) => ({
...content,
parts: (content.parts || []).map((part) => {
const newPart = { ...part };
if (newPart.functionCall) {
const fc = newPart.functionCall;
const name = fc.name?.trim() || 'generic_tool';
if (fc.id && fc.id.startsWith(`${name}__`)) {
newPart.functionCall = {
name: fc.name,
args: fc.args,
id: fc.id.substring(name.length + 2),
};
return contents.map((content) => {
const parts = (content.parts || [])
.map((part) => {
const newPart = { ...part };
if (newPart.functionCall) {
const fc = newPart.functionCall;
const name = fc.name?.trim() || 'generic_tool';
if (fc.id && fc.id.startsWith(`${name}__`)) {
newPart.functionCall = {
name: fc.name,
args: fc.args,
id: fc.id.substring(name.length + 2),
};
}
}
}
if (newPart.functionResponse) {
const fr = newPart.functionResponse;
const name = fr.name?.trim() || 'generic_tool';
if (fr.id && fr.id.startsWith(`${name}__`)) {
newPart.functionResponse = {
name: fr.name,
response: fr.response,
id: fr.id.substring(name.length + 2),
};
if (newPart.functionResponse) {
const fr = newPart.functionResponse;
const name = fr.name?.trim() || 'generic_tool';
if (fr.id && fr.id.startsWith(`${name}__`)) {
newPart.functionResponse = {
name: fr.name,
response: fr.response,
id: fr.id.substring(name.length + 2),
};
}
}
}
return newPart;
}),
}));
// If there's an empty text key alongside other active properties, remove the empty text key
// so it doesn't trigger "contains empty parts" validation errors on the Gemini API.
const hasOtherKeys = Object.keys(newPart).some(
(key) => key !== 'text' && key !== 'thought' && key !== 'callIndex',
);
if (newPart.text !== undefined && newPart.text === '' && hasOtherKeys) {
delete newPart.text;
}
return newPart;
})
.filter((part) => {
// Filter out truly empty parts that have only text: '' and no payload
const hasOtherKeys = Object.keys(part).some(
(key) => key !== 'text' && key !== 'thought' && key !== 'callIndex',
);
if (part.text !== undefined && part.text === '' && !hasOtherKeys) {
return false;
}
return true;
});
return {
...content,
parts,
};
});
}
export function coalesceConsecutiveRoles(
packages/core/src/tools/edit.test.ts
+26 -4
View File
@@ -817,10 +817,8 @@ function doIt() {
const result = await invocation.execute({
abortSignal: new AbortController().signal,
});
expect(result.llmContent).toMatch(/0 occurrences found for old_string/);
expect(result.returnDisplay).toMatch(
/Failed to edit, could not find the string to replace./,
);
expect(result.llmContent).toMatch(/Could not find an exact match/);
expect(result.returnDisplay).toMatch(/Could not find an exact match/);
expect(mockFixLLMEditWithInstruction).toHaveBeenCalled();
});
@@ -1369,6 +1367,30 @@ function doIt() {
expect(mockFixLLMEditWithInstruction).not.toHaveBeenCalled();
});
it('fails fast without calling FixLLMEditWithInstruction when old_string is empty', async () => {
const filePath = path.join(rootDir, 'empty_old_string_test.txt');
fs.writeFileSync(filePath, 'Some content.', 'utf8');
// Enable LLM correction for this test
(mockConfig.getDisableLLMCorrection as Mock).mockReturnValue(false);
const params = {
file_path: filePath,
instruction: 'Replace empty text',
old_string: ' ',
new_string: 'replacement',
};
const invocation = tool.build(params);
const result = await invocation.execute({
abortSignal: new AbortController().signal,
});
expect(mockFixLLMEditWithInstruction).not.toHaveBeenCalled();
expect(result.error?.type).toBe(ToolErrorType.INVALID_TOOL_PARAMS);
expect(result.error?.message).toContain('ReadFile');
});
});
describe('JIT context discovery', () => {
packages/core/src/tools/edit.ts
+20 -8
View File
@@ -53,11 +53,7 @@ import {
logEditCorrectionEvent,
} from '../telemetry/loggers.js';
import {
EDIT_TOOL_NAME,
READ_FILE_TOOL_NAME,
EDIT_DISPLAY_NAME,
} from './tool-names.js';
import { EDIT_TOOL_NAME, EDIT_DISPLAY_NAME } from './tool-names.js';
import { debugLogger } from '../utils/debugLogger.js';
import levenshtein from 'fast-levenshtein';
import { EDIT_DEFINITION } from './definitions/coreTools.js';
@@ -366,8 +362,8 @@ export function getErrorReplaceResult(
undefined;
if (occurrences === 0) {
error = {
display: `Failed to edit, could not find the string to replace.`,
raw: `Failed to edit, 0 occurrences found for old_string in ${params.file_path}. Ensure you're not escaping content incorrectly and check whitespace, indentation, and context. Use ${READ_FILE_TOOL_NAME} tool to verify.`,
display: `Could not find an exact match for old_string in '${params.file_path}'.`,
raw: `Could not find an exact match for 'old_string' in '${params.file_path}'. If previous edits modified the file or you are modifying lines outside your recent read window, please use ReadFile to inspect the target lines before retrying with an exact 'old_string'.`,
type: ToolErrorType.EDIT_NO_OCCURRENCE_FOUND,
};
} else if (!params.allow_multiple && occurrences !== 1) {
@@ -552,6 +548,22 @@ class EditToolInvocation
abortSignal: AbortSignal,
originalLineEnding: '\r\n' | '\n',
): Promise<CalculatedEdit> {
// Fail fast without invoking FixLLMEditWithInstruction when old_string is empty
if (!params.old_string || params.old_string.trim() === '') {
return {
currentContent,
newContent: currentContent,
occurrences: 0,
isNewFile: false,
error: {
display: "Edit failed: 'old_string' cannot be empty.",
raw: "The 'old_string' parameter is required. The Edit tool performs localized search-and-replace. If you are modifying a section of the file you have not viewed recently, call ReadFile on the target line range to inspect the current code, then provide the exact matching lines in 'old_string'.",
type: ToolErrorType.INVALID_TOOL_PARAMS,
},
originalLineEnding,
};
}
// In order to keep from clobbering edits made outside our system,
// check if the file has been modified since we first read it.
let errorForLlmEditFixer = initialError.raw;
@@ -1023,7 +1035,7 @@ ${snippet}`);
}
if (this.params.modified_by_user) {
llmSuccessMessageParts.push(
`User modified the \`new_string\` content to be: ${this.params.new_string}.`,
`The confirmation step modified the \`new_string\` content to be: ${this.params.new_string}.`,
);
}