diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 2247576d9e..a3b86fd981 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -7956,6 +7956,378 @@ describe('GeminiChat', async () => { } }); + /** + * The JSONL transcript that `--resume` / `--continue` reads is written + * by `recordAssistantTurn`, not by `this.history`. The continuation + * attempt's own parts carry only the resumed remainder, so without + * merging the delivered prefix back in, the durable transcript starts + * the recovered turn mid-sentence. + * + * `processStreamResponse` folds the prefix into the response parts once, + * before it writes either layer, so these tests assert the record and + * history agree — not just that the record is merged. Two earlier + * shapes failed exactly there: deduping the record against the trimmed + * `contentText` while history used the raw part, and merging in the + * outer send loop after the record had already been appended. + */ + function chatWithRecorder(recordAssistantTurn: ReturnType) { + return new GeminiChat( + mockConfig, + config, + [], + { + recordAssistantTurn, + recordChatCompression: vi.fn(), + } as unknown as ConstructorParameters[3], + uiTelemetryService, + ); + } + + function recordedText( + recordAssistantTurn: ReturnType, + callIndex = 0, + ): string | undefined { + const message = recordAssistantTurn.mock.calls[callIndex]![0] + .message as Array<{ text?: string }>; + return message.find((part) => part.text !== undefined)?.text; + } + + it('records the delivered prefix with the resumed remainder in one turn', async () => { + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = chatWithRecorder(recordAssistantTurn); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(cutAfter([textChunk('first half ')])) + .mockResolvedValueOnce( + (async function* () { + yield textChunk('second half', 'STOP'); + })(), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-transport-continuation-record', + ); + await collectStreamWithFakeTimers(stream, 5_000); + + // One turn in, one turn on disk. + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordedText(recordAssistantTurn)).toBe( + 'first half second half', + ); + // The durable record and in-memory history must agree. + expect(chatWithRecording.getHistory().at(-1)).toEqual({ + role: 'model', + parts: [{ text: 'first half second half' }], + }); + } finally { + vi.useRealTimers(); + } + }); + + it('merges a whitespace-leading remainder identically in both layers', async () => { + // R1-1: the record used to dedupe against `contentText`, which is + // trimmed, while history merged the raw part. A cut landing on a token + // boundary (before a space) then fused the two words in the transcript + // only — "The result is" + " 42." recorded as "The result is42.". + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = chatWithRecorder(recordAssistantTurn); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(cutAfter([textChunk('The result is')])) + .mockResolvedValueOnce( + (async function* () { + yield textChunk(' 42.', 'STOP'); + })(), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-transport-continuation-record-boundary', + ); + await collectStreamWithFakeTimers(stream, 5_000); + + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordedText(recordAssistantTurn)).toBe('The result is 42.'); + expect(chatWithRecording.getHistory().at(-1)).toEqual({ + role: 'model', + parts: [{ text: 'The result is 42.' }], + }); + } finally { + vi.useRealTimers(); + } + }); + + it('keeps a whitespace-boundary overlap dedup consistent across layers', async () => { + // The dedup-divergence half of R1-1: " total" is a 6-byte overlap only + // while untrimmed, so trimming the operand lost the dedup entirely and + // recorded "The grand totaltotal sum is 9.". + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = chatWithRecorder(recordAssistantTurn); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(cutAfter([textChunk('The grand total')])) + .mockResolvedValueOnce( + (async function* () { + yield textChunk(' total sum is 9.', 'STOP'); + })(), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-transport-continuation-record-boundary-overlap', + ); + await collectStreamWithFakeTimers(stream, 5_000); + + const history = chatWithRecording.getHistory().at(-1); + const historyText = history?.parts?.find( + (part) => part.text !== undefined, + )?.text; + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + // Whatever the dedup decides, both layers must decide it the same. + expect(recordedText(recordAssistantTurn)).toBe(historyText); + expect(recordedText(recordAssistantTurn)).not.toContain('totaltotal'); + } finally { + vi.useRealTimers(); + } + }); + + it('agrees across layers when the consumer aborts at the deferred finish chunk', async () => { + // R2-2: on a tool-result continuation the finishReason is withheld and + // re-emitted as a synthetic chunk AFTER the history push — a + // suspension point. `Turn.run` returns at exactly that kind of chunk + // when the user hits Esc. While the merge lived in the outer send + // loop, abandoning here left a merged record against a remainder-only + // history, and the JSONL is append-only so nothing reconciles it. + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = chatWithRecorder(recordAssistantTurn); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(cutAfter([textChunk('Analysis: the file ')])) + .mockResolvedValueOnce( + (async function* () { + yield textChunk('contains the bug.', 'STOP'); + })(), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { + // A functionResponse turn is what makes this a tool-result + // continuation, which is what defers the finishReason. + message: [ + { + functionResponse: { + id: 'call_deferred_window', + name: 'read_file', + response: { output: 'file contents' }, + }, + }, + ], + }, + 'prompt-transport-continuation-record-deferred-abort', + ); + + const collecting = (async () => { + for await (const event of stream) { + // The pass-through chunks have their finishReason stripped, so + // this fires only on the synthetic deferred chunk. + if ( + event.type === StreamEventType.CHUNK && + event.value.candidates?.[0]?.finishReason + ) { + break; + } + } + })(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(5_000); + await collecting; + + const historyText = chatWithRecording + .getHistory() + .at(-1) + ?.parts?.find((part) => part.text !== undefined)?.text; + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordedText(recordAssistantTurn)).toBe( + 'Analysis: the file contains the bug.', + ); + // The durable record and in-memory history must not disagree, even + // though the send was abandoned before it could finish. + expect(historyText).toBe(recordedText(recordAssistantTurn)); + } finally { + vi.useRealTimers(); + } + }); + + it('dedupes replayed overlap in the recorded turn too', async () => { + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = chatWithRecorder(recordAssistantTurn); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + cutAfter([textChunk('The quick brown fox jumps over')]), + ) + .mockResolvedValueOnce( + (async function* () { + yield textChunk('jumps over the lazy dog.', 'STOP'); + })(), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-transport-continuation-record-overlap', + ); + await collectStreamWithFakeTimers(stream, 5_000); + + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordedText(recordAssistantTurn)).toBe( + 'The quick brown fox jumps over the lazy dog.', + ); + } finally { + vi.useRealTimers(); + } + }); + + it('records nothing of a continuation a fresh-restart retry discarded', async () => { + // The mirror of the merge: when the continuation is superseded, the + // delivered text is dropped from history, so it must stay out of the + // transcript too. Recording the prefix when the continuation is + // *scheduled* would fix `--resume` for the success case and duplicate + // the answer here. + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = chatWithRecorder(recordAssistantTurn); + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(cutAfter([textChunk('doomed fragment ')])) + .mockResolvedValueOnce( + (async function* () { + throw new InvalidStreamError( + 'Model stream ended with empty response text.', + 'NO_RESPONSE_TEXT', + ); + + yield {} as GenerateContentResponse; + })(), + ) + .mockResolvedValueOnce( + (async function* () { + yield textChunk('a clean answer', 'STOP'); + })(), + ); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-transport-continuation-record-superseded', + ); + await collectStreamWithFakeTimers(stream, 10_000); + + // Three attempts proves the continuation really was scheduled and + // then superseded, rather than never starting. + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(3); + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + expect(recordedText(recordAssistantTurn)).toBe('a clean answer'); + } finally { + vi.useRealTimers(); + } + }); + + it('keeps the record remainder-only when the continuation itself is cut after a tool call', async () => { + // The one case where the prefix and a deferred partial record are + // live at the same time. A continuation attempt that yields a + // functionCall and then dies is excluded from continuing again + // (`canContinueAfterTransportCut` requires !streamYieldedFunctionCall), + // so the prefix is still set while `pendingPartialAssistantRecord` + // stashes the partial turn. + // + // The attempt did not survive, so the prefix must stay out of BOTH + // layers: history keeps the remainder-only partial (the merge at the + // success exit never runs) and the flushed record has to match it. + // Merging unconditionally instead of on success only would put the + // delivered text in the transcript and not in history — the same + // desync this fix removes, pointing the other way. + vi.useFakeTimers(); + try { + const recordAssistantTurn = vi.fn(); + const chatWithRecording = chatWithRecorder(recordAssistantTurn); + const toolCallChunk = { + candidates: [ + { + content: { + role: 'model', + parts: [ + { + functionCall: { + id: 'call_after_continuation', + name: 'read_file', + args: { path: '/tmp/x.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce(cutAfter([textChunk('delivered half ')])) + .mockResolvedValueOnce(cutAfter([toolCallChunk])); + + const stream = await chatWithRecording.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-transport-continuation-record-fc-cut', + ); + // This send rejects, so it cannot use `collectStreamWithFakeTimers`: + // that helper returns the collecting promise only after advancing + // timers, and the cut lands during the advance — leaving the + // rejection momentarily unhandled. Attach the assertion first, like + // `expectStreamExhaustion` above. + const collecting = (async () => { + for await (const _ of stream) { + /* consume */ + } + })(); + const settled = (async () => { + await expect(collecting).rejects.toThrow('terminated'); + })(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(10_000); + await settled; + + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(2); + expect(recordAssistantTurn).toHaveBeenCalledTimes(1); + // No text part at all: the attempt yielded only a functionCall. + expect(recordedText(recordAssistantTurn)).toBeUndefined(); + + // And the durable record still matches what survives in memory. + const lastTurn = chatWithRecording.getHistory().at(-1); + expect(lastTurn?.role).toBe('model'); + expect( + lastTurn?.parts?.some((part) => + part.text?.includes('delivered half'), + ), + ).toBe(false); + } finally { + vi.useRealTimers(); + } + }); + it('drops replayed overlap when the model repeats its own tail', async () => { vi.useFakeTimers(); try { diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index a08d2b5f96..80512921e7 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -835,6 +835,27 @@ function getRecoveryContinuationSuffix( return continuationText; } +/** + * Join already-delivered text to the continuation that resumes it, dropping + * any tail the model replayed. + * + * The single definition of "merged turn text" for the transport-continuation + * path. Both the durable JSONL record and in-memory history are built from one + * call to this (see `processStreamResponse`), so the two storage layers cannot + * drift apart if the dedup rule ever changes — the same reason the + * `willPersistToHistory` gate is a shared binding rather than two copies of + * one expression. + */ +function mergeDeliveredPrefix( + deliveredText: string, + continuationText: string, +): string { + return ( + deliveredText + + getRecoveryContinuationSuffix(deliveredText, continuationText) + ); +} + function isPlainTextPart(part: Part | undefined): part is Part & { text: string; } { @@ -2761,6 +2782,12 @@ export class GeminiChat { prompt_id, requestOverrides, turnGoalContext, + // Captured by value, so the attempt records exactly the prefix + // `buildAttemptContents()` just asked the model to resume from, + // even if a later branch resets the continuation. + transportContinuationPrefix.length > 0 + ? transportContinuationPrefix + : undefined, ); lastFinishReason = undefined; @@ -2789,10 +2816,13 @@ export class GeminiChat { } lastError = null; - if (transportContinuationPrefix.length > 0) { - self.prependTextToLastModelTurn(transportContinuationPrefix); - transportContinuationPrefix = ''; - } + // The merge itself now happens inside `processStreamResponse`, + // which folds the prefix into the parts before it writes either + // the JSONL record or the history turn (issue #8094). Merging + // again here would risk double-applying it: the dedup helper only + // strips a replayed prefix that clears its significance floor, so + // a short prefix would survive the second pass and be doubled. + transportContinuationPrefix = ''; break; } catch (error) { lastError = error; @@ -3856,6 +3886,7 @@ export class GeminiChat { retryErrorCodes?: readonly number[]; }, goalContext?: GoalTurnPermit, + transportContinuationPrefix?: string, ): Promise> { const generator = overrides?.contentGenerator ?? this.config.getContentGenerator(); @@ -3932,7 +3963,12 @@ export class GeminiChat { }, }); - return this.processStreamResponse(model, streamResponse, goalContext); + return this.processStreamResponse( + model, + streamResponse, + goalContext, + transportContinuationPrefix, + ); } private async *makeFallbackStream( @@ -4394,10 +4430,19 @@ export class GeminiChat { } } + /** + * @param transportContinuationPrefix - Text a previous attempt already + * delivered before a socket cut, which this attempt was asked to resume + * from (issue #7832). On success it is folded into the response parts + * before either durable write, so the JSONL transcript and in-memory + * history carry the same merged turn (issue #8094). Undefined on every + * non-continuation send. + */ private async *processStreamResponse( model: string, streamResponse: AsyncGenerator, goalContext?: GoalTurnPermit, + transportContinuationPrefix?: string, ): AsyncGenerator { // Collect ALL parts from the model response (including thoughts for recording) const allModelParts: Part[] = []; @@ -4848,6 +4893,62 @@ export class GeminiChat { streamError === null || (hasToolCall && (thoughtContentPart || consolidatedHistoryParts.length > 0)); + // Transport-continuation merge (issue #8094). `allModelParts` is + // per-attempt, so a continuation's parts carry the resumed remainder only. + // Fold the already-delivered prefix back in HERE — into the parts + // themselves, before either durable write — so the JSONL record below and + // the `this.history.push` further down are derived from the same data and + // cannot disagree. Otherwise `--resume` rehydrates a turn that starts + // mid-sentence while the live session shows a coherent answer. + // + // Merging in one place is load-bearing, not tidiness: + // - Computing the record's text and history's text from separate + // expressions lets them drift. They already would: `contentText` is + // trimmed (see its definition above) while the pushed parts are raw, + // so deduping the record against the trimmed text fuses words when the + // remainder opens with whitespace ("The result is" + " 42." → + // "The result is42."). + // - Writing them at different times opens a window. The record is + // appended below, the history push happens after it, and a + // `deferredFinishReason` chunk is yielded after that — a suspension + // point. A consumer abandoning iteration there (an abort inside + // `Turn.run`) would strand a merged record against a remainder-only + // history, and the JSONL is append-only so nothing reconciles it. + // + // Placed after the stream-validation throws above so an empty continuation + // still fails validation on its own merits rather than being masked by the + // prefix. + // + // Success only. On `streamError !== null` the parts must keep matching the + // remainder-only partial that survives in history (the + // `pendingPartialAssistantRecord` path below) — the prefix belongs to an + // attempt that did not survive, and a fresh-restart retry discards it via + // `resetTransportContinuation`. + if (streamError === null && transportContinuationPrefix) { + const textIndex = consolidatedHistoryParts.findIndex(isPlainTextPart); + if (textIndex < 0) { + // Continuation returned no text of its own (e.g. only a functionCall). + // `thoughtContentPart` is prepended separately at the push below, so + // index 0 here is already "after any leading thought part". + consolidatedHistoryParts.unshift({ text: transportContinuationPrefix }); + } else { + const remainderPart = consolidatedHistoryParts[textIndex] as Part & { + text: string; + }; + consolidatedHistoryParts[textIndex] = { + ...remainderPart, + text: mergeDeliveredPrefix( + transportContinuationPrefix, + remainderPart.text, + ), + }; + } + contentText = consolidatedHistoryParts + .filter((part) => part.text) + .map((part) => part.text) + .join('') + .trim(); + } if ( willPersistToHistory && (thoughtContentPart || contentText || hasToolCall || usageMetadata) @@ -4970,53 +5071,6 @@ export class GeminiChat { } } - /** - * Prepend already-delivered text to the trailing model turn. - * - * Used by the transport-continuation path (issue #7832): after a socket cut - * mid-response, the caller has seen text that `processStreamResponse` - * deliberately did not persist, and the continuation attempt's own turn - * carries only the resumed remainder. Without this, durable history would - * hold an answer that starts mid-sentence — visibly wrong on `/compress`, - * `--resume`, and every later turn's context. - * - * The delivered text is merged into the turn's first plain-text part (kept - * after any leading thought part, matching the - * `[thoughtPart?, ...text]` shape `processStreamResponse` produces), or - * inserted as a new part when the continuation returned no text of its own. - * Overlap is deduped by the same helper the output-recovery merge uses, so a - * model that replays part of its previous tail does not double it. - */ - private prependTextToLastModelTurn(deliveredText: string): void { - if (deliveredText.length === 0) return; - const lastEntry = this.history.at(-1); - if (lastEntry?.role !== 'model') { - debugLogger.warn( - '[TRANSPORT_CONTINUATION] Trailing entry is not a model turn; ' + - 'dropping the delivered-text merge.', - { role: lastEntry?.role ?? 'undefined' }, - ); - return; - } - const parts = [...(lastEntry.parts ?? [])]; - const textIndex = parts.findIndex(isPlainTextPart); - if (textIndex < 0) { - const insertAt = parts.findIndex((part) => !part.thought); - parts.splice(insertAt < 0 ? parts.length : insertAt, 0, { - text: deliveredText, - }); - } else { - const continuationPart = parts[textIndex] as Part & { text: string }; - parts[textIndex] = { - ...continuationPart, - text: - deliveredText + - getRecoveryContinuationSuffix(deliveredText, continuationPart.text), - }; - } - lastEntry.parts = parts; - } - /** * Merge `pairCount` trailing (user_recovery, model_continuation) pairs back * into the model turn that precedes them. Used after the output-token