diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 248736dc9ba..34b14b96918 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -2584,6 +2584,47 @@ describe('GeminiChat', async () => { }); }); + it('reattaches stored image markers on later below-threshold requests', async () => { + vi.mocked(mockConfig.getChatCompression).mockReturnValue({ + maxRecentImagesToRetain: 1, + imagePayloadThreshold: 1, + }); + chat.setHistory([ + { + role: 'user', + parts: [{ inlineData: { mimeType: 'image/png', data: 'old-shot' } }], + }, + { role: 'model', parts: [{ text: 'I see the image' }] }, + ]); + vi.mocked(mockContentGenerator.generateContentStream).mockImplementation( + async () => streamResponse(stopResponse([{ text: 'response' }])), + ); + + for (const [message, promptId] of [ + ['first question', 'prompt-id-image-refs-first'], + ['second question', 'prompt-id-image-refs-second'], + ] as const) { + const stream = await chat.sendMessageStream( + 'test-model', + { message }, + promptId, + ); + for await (const _ of stream) { + // consume stream + } + } + + const durable = JSON.stringify(chat.getHistory()); + expect(durable).toMatch(/Image #[a-f0-9]{12}/); + expect(durable).not.toContain('"data":"old-shot"'); + const secondRequest = vi.mocked( + mockContentGenerator.generateContentStream, + ).mock.calls[1]?.[0]; + expect(JSON.stringify(secondRequest?.contents)).toContain( + '"data":"old-shot"', + ); + }); + it('coalesces startup reminders with the first user prompt for provider requests', async () => { chat.setHistory([ { @@ -5687,19 +5728,29 @@ describe('GeminiChat', async () => { }); describe('getHistoryShallow', () => { - it('copies containers without structured-cloning large part payloads', () => { + it('copies Part containers without cloning large leaf payloads', () => { const payload = { output: 'x'.repeat(128 * 1024) }; + const topLevelInlineData = { + mimeType: 'image/png', + data: 'top-level-image', + }; + const nestedInlineData = { + mimeType: 'image/png', + data: 'nested-image', + }; + const topLevelPart: Part = { inlineData: topLevelInlineData }; + const nestedPart: Part = { inlineData: nestedInlineData }; + const functionResponsePart: Part = { + functionResponse: { + id: 'call-1', + name: 'read_file', + response: payload, + parts: [nestedPart], + }, + }; const content: Content = { role: 'user', - parts: [ - { - functionResponse: { - id: 'call-1', - name: 'read_file', - response: payload, - }, - }, - ], + parts: [topLevelPart, functionResponsePart], }; chat.addHistory(content); const structuredCloneSpy = vi @@ -5714,7 +5765,25 @@ describe('GeminiChat', async () => { expect(history).toEqual([content]); expect(history[0]).not.toBe(content); expect(history[0]!.parts).not.toBe(content.parts); - const response = history[0]!.parts![0] as { + expect(history[0]!.parts![0]).not.toBe(topLevelPart); + expect(history[0]!.parts![0]!.inlineData).toBe(topLevelInlineData); + const copiedFunctionResponsePart = history[0]!.parts![1]!; + expect(copiedFunctionResponsePart).not.toBe(functionResponsePart); + expect(copiedFunctionResponsePart.functionResponse).not.toBe( + functionResponsePart.functionResponse, + ); + const copiedNested = copiedFunctionResponsePart.functionResponse + ?.parts as Part[]; + expect(copiedNested).not.toBe( + functionResponsePart.functionResponse?.parts, + ); + expect(copiedNested[0]).not.toBe(nestedPart); + expect(copiedNested[0]!.inlineData).toBe(nestedInlineData); + delete history[0]!.parts![0]!.inlineData; + delete copiedNested[0]!.inlineData; + expect(topLevelPart.inlineData).toBe(topLevelInlineData); + expect(nestedPart.inlineData).toBe(nestedInlineData); + const response = copiedFunctionResponsePart as { functionResponse: { response: typeof payload }; }; expect(response.functionResponse.response).toBe(payload); diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 4f791d121b7..93758e4a238 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -77,6 +77,7 @@ import { } from '../services/chatCompressionService.js'; import { acquireSleepInhibitor } from '../services/sleepInhibitor.js'; import { + getFunctionResponseParts, resolveCompactionTuning, resolveSlimmingConfig, slimCompactionInput, @@ -1233,7 +1234,19 @@ function appendCuratedContent( function copyContentContainer(content: Content): Content { return { ...content, - ...(content.parts ? { parts: [...content.parts] } : {}), + ...(content.parts ? { parts: content.parts.map(copyPartContainer) } : {}), + }; +} + +function copyPartContainer(part: Part): Part { + const nested = getFunctionResponseParts(part); + if (!nested) return { ...part }; + return { + ...part, + functionResponse: { + ...part.functionResponse, + parts: nested.map((inner) => ({ ...inner })), + }, }; } @@ -1876,6 +1889,7 @@ export class GeminiChat { const { maxRecentImages, imagePayloadThreshold } = resolveCompactionTuning( this.config.getChatCompression(), ); + let replaced: ReturnType = []; if (countAllInlineImages(curatedHistory) >= imagePayloadThreshold) { const skipEntry = currentUserContent ? curatedHistory.find( @@ -1885,24 +1899,28 @@ export class GeminiChat { currentUserContent.parts?.some((p) => c.parts?.includes(p))), ) : undefined; - const replaced = replaceImagePayloadsInPlace( + replaced = replaceImagePayloadsInPlace( curatedHistory, this.imagePayloadStore, skipEntry, ); - const requestHistory = curatedHistory.map(copyContentContainer); - const reattachParts = buildReattachParts(replaced, maxRecentImages); - if (reattachParts.length > 0) { - const last = requestHistory.at(-1); - if (last?.role === 'user') { - last.parts = [...(last.parts ?? []), ...reattachParts]; - } else { - requestHistory.push({ role: 'user', parts: reattachParts }); - } + } + const requestHistory = curatedHistory.map(copyContentContainer); + const reattachParts = buildReattachParts( + replaced, + maxRecentImages, + requestHistory, + this.imagePayloadStore, + ); + if (reattachParts.length > 0) { + const last = requestHistory.at(-1); + if (last?.role === 'user') { + last.parts = [...(last.parts ?? []), ...reattachParts]; + } else { + requestHistory.push({ role: 'user', parts: reattachParts }); } - return requestHistory; } - return curatedHistory.map(copyContentContainer); + return requestHistory; } private getRequestHistoryForRoute( @@ -4047,9 +4065,9 @@ export class GeminiChat { } /** - * Returns a shallow copy of the history and each entry's parts array without - * cloning large part payloads. Use only for read-only consumers or consumers - * that replace touched entries before mutating them. + * Copies history containers, Part objects, and nested functionResponse parts + * without cloning large leaf payloads. Consumers must not mutate leaf + * payload objects. */ getHistoryShallow(curated: boolean = false): Content[] { const history = curated diff --git a/packages/core/src/services/image-payload-references.test.ts b/packages/core/src/services/image-payload-references.test.ts index 9e4b27bdfc0..95699c5884c 100644 --- a/packages/core/src/services/image-payload-references.test.ts +++ b/packages/core/src/services/image-payload-references.test.ts @@ -88,14 +88,16 @@ describe('prepareImagePayloadsForRequest', () => { store, }, ); - const id = JSON.stringify(firstPass).match(/Image #([a-f0-9]{12})/)?.[1]; - expect(id).toBeDefined(); + const marker = JSON.stringify(firstPass).match( + /\[Image #[a-f0-9]{12}: [^\]]+\]/, + )?.[0]; + expect(marker).toBeDefined(); const prepared = prepareImagePayloadsForRequest( [ oldImage, toolImageTurn('new-shot'), - { role: 'user', parts: [{ text: `inspect Image #${id}` }] }, + { role: 'user', parts: [{ text: `inspect ${marker}` }] }, ], { maxRecentImages: 0, @@ -117,11 +119,13 @@ describe('prepareImagePayloadsForRequest', () => { store, }, ); - const id = JSON.stringify(firstPass).match(/Image #([a-f0-9]{12})/)?.[1]; - expect(id).toBeDefined(); + const marker = JSON.stringify(firstPass).match( + /\[Image #[a-f0-9]{12}: [^\]]+\]/, + )?.[0]; + expect(marker).toBeDefined(); const prepared = prepareImagePayloadsForRequest( - [{ role: 'user', parts: [{ text: `inspect Image #${id}` }] }], + [{ role: 'user', parts: [{ text: `inspect ${marker}` }] }], { maxRecentImages: 0, store, @@ -133,6 +137,25 @@ describe('prepareImagePayloadsForRequest', () => { ]); }); + it('does not resurrect a stored image from a bare Image #id echo', () => { + const store = new InMemoryImagePayloadStore(); + const firstPass = prepareImagePayloadsForRequest( + [toolImageTurn('old-shot'), { role: 'model', parts: [{ text: 'ok' }] }], + { maxRecentImages: 0, store }, + ); + const id = JSON.stringify(firstPass).match(/Image #([a-f0-9]{12})/)?.[1]; + expect(id).toBeDefined(); + + // A model reply echoing just the id (not the full eviction marker) must + // not re-inject the stored payload. + const prepared = prepareImagePayloadsForRequest( + [{ role: 'user', parts: [{ text: `I saw Image #${id} earlier` }] }], + { maxRecentImages: 0, store }, + ); + + expect(imageParts(prepared)).toEqual([]); + }); + it('reattaches the most recent unique historical images', () => { const store = new InMemoryImagePayloadStore(); const prepared = prepareImagePayloadsForRequest( @@ -270,6 +293,43 @@ describe('replaceImagePayloadsInPlace', () => { expect(JSON.stringify(contents)).toContain('"data":"current-shot"'); expect(JSON.stringify(contents)).not.toContain('"data":"old-shot"'); }); + + it('rewrites shared top-level and nested Part objects', () => { + const store = new InMemoryImagePayloadStore(); + const topLevel: Part = { + inlineData: { mimeType: 'image/png', data: 'top-level' }, + }; + const nested: Part = { + inlineData: { mimeType: 'image/png', data: 'nested' }, + }; + const durable: Content[] = [ + { role: 'user', parts: [topLevel] }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call-1', + name: 'screenshot', + response: {}, + parts: [nested], + }, + }, + ], + }, + ]; + const curated: Content[] = durable.map((content) => ({ + ...content, + parts: [...(content.parts ?? [])], + })); + + replaceImagePayloadsInPlace(curated, store); + + expect(JSON.stringify(durable)).not.toContain('"data":'); + expect(JSON.stringify(durable).match(/Image #[a-f0-9]{12}/g)).toHaveLength( + 2, + ); + }); }); describe('buildReattachParts', () => { @@ -296,4 +356,84 @@ describe('buildReattachParts', () => { const replaced = replaceImagePayloadsInPlace([toolImageTurn('a')], store); expect(buildReattachParts(replaced, 0)).toEqual([]); }); + + it('resolves stored markers even when the current replacement pass is empty', () => { + const store = new InMemoryImagePayloadStore(); + const contents = [ + toolImageTurn('a'), + toolImageTurn('b'), + toolImageTurn('c'), + { role: 'user', parts: [{ text: 'continue' }] }, + ]; + replaceImagePayloadsInPlace(contents, store); + + const parts = buildReattachParts([], 2, contents, store); + + expect( + parts + .filter((part) => part.inlineData) + .map((part) => part.inlineData?.data), + ).toEqual(['b', 'c']); + }); + + it('reattaches a marker in the current user turn outside the recency cap', () => { + const store = new InMemoryImagePayloadStore(); + const contents = [toolImageTurn('current')]; + replaceImagePayloadsInPlace(contents, store); + + const parts = buildReattachParts([], 0, contents, store); + + expect(parts.at(-1)?.inlineData?.data).toBe('current'); + }); + + it('bounds current-turn marker reattachment to the recency cap', () => { + const store = new InMemoryImagePayloadStore(); + const contents: Content[] = [ + { + role: 'user', + parts: ['a', 'b', 'c'].map((data) => ({ + inlineData: { mimeType: 'image/png', data }, + })), + }, + ]; + replaceImagePayloadsInPlace(contents, store); + + const parts = buildReattachParts([], 1, contents, store); + + expect( + parts + .filter((part) => part.inlineData) + .map((part) => part.inlineData?.data), + ).toEqual(['c']); + }); + + it('does not reattach an image that is already inline', () => { + const store = new InMemoryImagePayloadStore(); + const markerContents = [toolImageTurn('same')]; + replaceImagePayloadsInPlace(markerContents, store); + const marker = markerContents[0]!.parts![0]!; + const referencedContents: Content[] = [ + { + role: 'user', + parts: [ + marker, + { inlineData: { mimeType: 'image/png', data: 'same' } }, + ], + }, + ]; + + expect(buildReattachParts([], 1, referencedContents, store)).toEqual([]); + }); + + it('does not reattach an image already inline in a tool response', () => { + const store = new InMemoryImagePayloadStore(); + const markerContents = [toolImageTurn('same')]; + replaceImagePayloadsInPlace(markerContents, store); + const referencedContents: Content[] = [ + markerContents[0]!, + toolImageTurn('same'), + ]; + + expect(buildReattachParts([], 1, referencedContents, store)).toEqual([]); + }); }); diff --git a/packages/core/src/services/image-payload-references.ts b/packages/core/src/services/image-payload-references.ts index 4f01e11b2cf..0e36f321dc3 100644 --- a/packages/core/src/services/image-payload-references.ts +++ b/packages/core/src/services/image-payload-references.ts @@ -11,8 +11,12 @@ import { approxBase64Bytes } from '../core/inlineMediaLimit.js'; import { getFunctionResponseParts } from './compactionInputSlimming.js'; const IMAGE_ID_LENGTH = 12; +// Anchor the match to the full output of `imageReferenceText` so only the +// markers eviction actually wrote resolve against the store. A bare +// `Image #` echo (a model reply quoting the id, or a post-compaction +// summary that retained marker text) must not resurrect the stored payload. const IMAGE_REFERENCE_PATTERN = new RegExp( - `Image #([a-f0-9]{${IMAGE_ID_LENGTH}})`, + `\\[Image #([a-f0-9]{${IMAGE_ID_LENGTH}}): [^\\]]+\\]`, 'gi', ); @@ -49,16 +53,7 @@ export class InMemoryImagePayloadStore implements ImagePayloadStore { export function countAllInlineImages(contents: Content[]): number { let count = 0; - for (const content of contents) { - for (const part of content.parts ?? []) { - if (part.inlineData?.mimeType?.startsWith('image/')) count++; - const nested = getFunctionResponseParts(part); - if (!nested) continue; - for (const inner of nested) { - if (inner.inlineData?.mimeType?.startsWith('image/')) count++; - } - } - } + for (const _part of inlineImageParts(contents)) count++; return count; } @@ -76,58 +71,63 @@ export function replaceImagePayloadsInPlace( skipContent?: Content, ): StoredImagePayload[] { const replaced: StoredImagePayload[] = []; - for (const content of contents) { - if (content === skipContent) continue; - if (!content.parts) continue; - for (let i = 0; i < content.parts.length; i++) { - const part = content.parts[i]!; - if ( - part.inlineData?.mimeType?.startsWith('image/') && - part.inlineData.data - ) { - const stored = store.put(part); - replaced.push(stored); - content.parts[i] = { text: imageReferenceText(stored) }; - continue; - } - const nested = getFunctionResponseParts(part); - if (!nested) continue; - for (let j = 0; j < nested.length; j++) { - const inner = nested[j]!; - if ( - inner.inlineData?.mimeType?.startsWith('image/') && - inner.inlineData.data - ) { - const stored = store.put(inner); - replaced.push(stored); - nested[j] = { text: imageReferenceText(stored) }; - } - } - } + for (const part of inlineImageParts(contents, skipContent)) { + const stored = store.put(part); + replaced.push(stored); + part.text = imageReferenceText(stored); + delete part.inlineData; } return replaced; } /** - * Build the reattach parts for the most recent unique images from a - * replacement pass. Used after `replaceImagePayloadsInPlace` to append - * recent image bytes to the outgoing request. + * Build reattach parts from images replaced in the current pass and stored + * payloads referenced by markers in `referencedContents`, even when the + * current pass replaced nothing. */ export function buildReattachParts( replaced: StoredImagePayload[], maxRecentImages: number, + referencedContents: Content[] = [], + store?: ImagePayloadStore, ): Part[] { - if (maxRecentImages <= 0 || replaced.length === 0) return []; - const recent: StoredImagePayload[] = []; - const seen = new Set(); - for (let i = replaced.length - 1; i >= 0; i--) { - const img = replaced[i]!; - if (seen.has(img.id)) continue; - seen.add(img.id); - recent.push(img); - if (recent.length === maxRecentImages) break; + const referencedIds = collectReferencedImageIds(referencedContents); + if (replaced.length === 0 && (!store || referencedIds.size === 0)) return []; + const inlineIds = collectInlineImageIds(referencedContents); + const last = referencedContents.at(-1); + const lastReferencedIds = collectReferencedImageIds( + last?.role === 'user' ? [last] : [], + ); + const candidates: CollectedImage[] = replaced + .filter( + (image) => !inlineIds.has(image.id) && !lastReferencedIds.has(image.id), + ) + .map((stored) => ({ stored })); + + if (store) { + for (const id of referencedIds) { + if (inlineIds.has(id) || lastReferencedIds.has(id)) continue; + const stored = store.get(id); + if (stored) candidates.push({ stored }); + } + } + const recent = recentUniqueImages(candidates, maxRecentImages).map( + ({ stored }) => stored, + ); + const reattachLimit = Math.max(maxRecentImages, 1); + if (store) { + for (const id of lastReferencedIds) { + if (inlineIds.has(id) || recent.some((image) => image.id === id)) { + continue; + } + const stored = store.get(id); + if (stored) { + if (recent.length >= reattachLimit) recent.shift(); + recent.push(stored); + } + } } - recent.reverse(); + if (recent.length === 0) return []; return [ { text: @@ -147,7 +147,9 @@ export function prepareImagePayloadsForRequest( store: ImagePayloadStore; }, ): Content[] { - const referencedIds = collectReferencedImageIds(contents.at(-1)); + const referencedIds = collectReferencedImageIds( + contents.at(-1) ? [contents.at(-1)!] : [], + ); const collected: CollectedImage[] = []; const transformed = contents.map((content, index) => { if (index === options.preserveImagePartsForContentIndex) { @@ -243,15 +245,52 @@ function transformPart( return part; } -function collectReferencedImageIds(content: Content | undefined): Set { +function collectInlineImageIds(contents: Content[]): Set { const ids = new Set(); - for (const part of content?.parts ?? []) { - const text = part.text; - if (!text) continue; - for (const match of text.matchAll(IMAGE_REFERENCE_PATTERN)) { - const id = match[1]; - if (id) ids.add(id.toLowerCase()); + for (const part of inlineImageParts(contents)) { + ids.add(imagePartToStoredPayload(part).id); + } + return ids; +} + +function* inlineImageParts( + contents: Content[], + skipContent?: Content, +): Generator { + for (const content of contents) { + if (content === skipContent) continue; + for (const part of content.parts ?? []) { + if ( + part.inlineData?.mimeType?.startsWith('image/') && + part.inlineData.data + ) { + yield part; + } + for (const inner of getFunctionResponseParts(part) ?? []) { + if ( + inner.inlineData?.mimeType?.startsWith('image/') && + inner.inlineData.data + ) { + yield inner; + } + } + } + } +} + +function collectReferencedImageIds(contents: Content[]): Set { + const ids = new Set(); + const collect = (parts: Part[] | undefined): void => { + for (const part of parts ?? []) { + for (const match of part.text?.matchAll(IMAGE_REFERENCE_PATTERN) ?? []) { + const id = match[1]; + if (id) ids.add(id.toLowerCase()); + } + collect(getFunctionResponseParts(part)); } + }; + for (const content of contents) { + collect(content.parts); } return ids; } diff --git a/packages/core/src/utils/forkedAgent.cache.test.ts b/packages/core/src/utils/forkedAgent.cache.test.ts index 883b189dff7..49c731761dd 100644 --- a/packages/core/src/utils/forkedAgent.cache.test.ts +++ b/packages/core/src/utils/forkedAgent.cache.test.ts @@ -110,11 +110,24 @@ describe('CacheSafeParams', () => { expect(rereadTools[0].functionDeclarations).toHaveLength(1); }); - it('copies history containers without cloning part payloads', () => { + it('copies history containers and Part objects', () => { const historyPart = { text: 'large history entry' }; + const nestedPart = { + inlineData: { mimeType: 'image/png', data: 'screenshot' }, + }; const historyEntry: Content = { role: 'user', - parts: [historyPart], + parts: [ + historyPart, + { + functionResponse: { + id: 'call-1', + name: 'screenshot', + response: {}, + parts: [nestedPart], + }, + }, + ], }; const historyEntryWithoutParts: Content = { role: 'model' }; const history: Content[] = [historyEntry, historyEntryWithoutParts]; @@ -127,9 +140,14 @@ describe('CacheSafeParams', () => { expect(params!.history).toHaveLength(2); expect(params!.history).not.toBe(history); expect(params!.history[0]).not.toBe(historyEntry); - expect(params!.history[0]!.parts).toHaveLength(1); + expect(params!.history[0]!.parts).toHaveLength(2); expect(params!.history[0]!.parts).not.toBe(historyEntry.parts); - expect(params!.history[0]!.parts![0]).toBe(historyPart); + expect(params!.history[0]!.parts![0]).not.toBe(historyPart); + expect(params!.history[0]!.parts![0]).toEqual(historyPart); + const copiedNested = params!.history[0]!.parts![1]!.functionResponse + ?.parts as Array; + expect(copiedNested[0]).not.toBe(nestedPart); + expect(copiedNested[0]).toEqual(nestedPart); expect(params!.history[1]).not.toBe(historyEntryWithoutParts); expect('parts' in params!.history[1]!).toBe(false); @@ -139,7 +157,7 @@ describe('CacheSafeParams', () => { }); params!.history[0]!.parts!.push({ text: 'returned part mutation' }); expect(getCacheSafeParams()!.history).toHaveLength(2); - expect(getCacheSafeParams()!.history[0]!.parts).toHaveLength(1); + expect(getCacheSafeParams()!.history[0]!.parts).toHaveLength(2); }); }); diff --git a/packages/core/src/utils/forkedAgent.ts b/packages/core/src/utils/forkedAgent.ts index dd186eb2ace..3b5c2719c02 100644 --- a/packages/core/src/utils/forkedAgent.ts +++ b/packages/core/src/utils/forkedAgent.ts @@ -33,6 +33,7 @@ import type { Content, GenerateContentConfig, GenerateContentResponseUsageMetadata, + Part, } from '@google/genai'; import { runWithRuntimeContentGenerator, @@ -61,6 +62,7 @@ import { type ResolvedModelId, } from './modelId.js'; import { ToolNames } from '../tools/tool-names.js'; +import { getFunctionResponseParts } from '../services/compactionInputSlimming.js'; import { runWithChatRecordingSuppressed } from './chat-recording-suppression-context.js'; const debugLogger = createDebugLogger('FORKED_AGENT'); @@ -78,8 +80,9 @@ export interface CacheSafeParams { /** Full generation config including systemInstruction and tools */ generationConfig: GenerateContentConfig; /** - * Curated conversation history with copied Content and parts containers. - * Part objects are shared by reference; consumers must not mutate them. + * Curated conversation history with copied Content, parts arrays, and Part + * objects. Nested functionResponse parts are copied too; leaf payloads remain + * shared and must not be mutated. */ history: Content[]; /** Model identifier */ @@ -100,10 +103,22 @@ export interface CacheSafeParams { let currentCacheSafeParams: CacheSafeParams | null = null; let currentVersion = 0; +function clonePart(part: Part): Part { + const nested = getFunctionResponseParts(part); + if (!nested) return { ...part }; + return { + ...part, + functionResponse: { + ...part.functionResponse, + parts: nested.map((inner) => ({ ...inner })), + }, + }; +} + function copyHistoryContainers(history: Content[]): Content[] { return history.map((content) => ({ ...content, - ...(content.parts ? { parts: [...content.parts] } : {}), + ...(content.parts ? { parts: content.parts.map(clonePart) } : {}), })); }