diff --git a/docs/design/full-turn-multimodal-routing.md b/docs/design/full-turn-multimodal-routing.md index f1cbea31507..eea8b585b7e 100644 --- a/docs/design/full-turn-multimodal-routing.md +++ b/docs/design/full-turn-multimodal-routing.md @@ -43,3 +43,9 @@ LLM-based automatic chat compression remains on the primary-model path. A full-t Phase 1 covers the interactive TUI, ACP, and non-interactive CLI. Textual `@` paths are resolved to their canonical target before MIME detection, workspace checks, ignore filtering, and file reads. Both the user-supplied alias and canonical target must pass ignore filtering, so a symlink cannot disguise an ignored file or a non-image target. Hardlinks are not resolved by `realpath` and are not covered by this check. + +## Durable visual context + +Historical images are represented by deterministic `Image #` references. Below the payload threshold all historical images remain attached; at or above it only the configured recent images are attached. A prompt that explicitly names one or more image IDs attaches only those images, then follows the same bridge or full-turn routing decision as a newly supplied image. + +Resume rebuilds the session-scoped image store from the original JSONL messages, including images older than the latest compression checkpoint. Compaction restoration headers retain the same IDs, and `/clear` drops the store. Raw bytes stay in the attachment layer and are never inserted into text sent to a text-only model. diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 851c239d9b7..f9c6ce62d43 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -517,6 +517,9 @@ describe('Session', () => { }; let mockLlmClient: { getChat: ReturnType; + setHistory: ReturnType; + truncateHistory: ReturnType; + stripOrphanedUserEntriesFromHistory: ReturnType; isInitialized: ReturnType; refreshSystemInstruction: ReturnType; setTools: ReturnType; @@ -737,6 +740,7 @@ describe('Session', () => { .fn() .mockReturnValue(new Map()), getLastModelMessageText: vi.fn().mockReturnValue(''), + resolveImageReferences: vi.fn((parts) => parts), setHistory: vi.fn(), truncateHistory: vi.fn(), stripThoughtsFromHistory: vi.fn(), @@ -745,6 +749,13 @@ describe('Session', () => { } as unknown as LlmChat; mockLlmClient = { getChat: vi.fn().mockReturnValue(mockChat), + setHistory: vi.fn(), + truncateHistory: vi.fn((keepCount) => + mockChat.truncateHistory(keepCount), + ), + stripOrphanedUserEntriesFromHistory: vi.fn(() => + mockChat.stripOrphanedUserEntriesFromHistory(), + ), isInitialized: vi.fn().mockReturnValue(true), refreshSystemInstruction: vi.fn().mockResolvedValue(undefined), setTools: vi.fn().mockResolvedValue(undefined), @@ -6077,6 +6088,7 @@ describe('Session', () => { const result = session.rewindToTurn(1); expect(result).toEqual({ targetTurnIndex: 1, apiTruncateIndex: 2 }); + expect(mockGeminiClient.truncateHistory).toHaveBeenCalledWith(2); expect(mockChat.truncateHistory).toHaveBeenCalledWith(2); expect(mockChat.stripThoughtsFromHistory).toHaveBeenCalled(); expect(mockChatRecordingService.rewindRecording).toHaveBeenCalledWith( @@ -6421,7 +6433,7 @@ describe('Session', () => { session.restoreHistory(snapshot); expect(snapshot).toEqual(history); - expect(mockChat.setHistory).toHaveBeenCalledWith(history); + expect(mockGeminiClient.setHistory).toHaveBeenCalledWith(history); expect(mockChat.getHistory).not.toHaveBeenCalled(); }); @@ -6468,7 +6480,7 @@ describe('Session', () => { expect(() => session.restoreHistory([])).toThrow( 'Cannot restore history while a prompt is running', ); - expect(mockChat.setHistory).not.toHaveBeenCalled(); + expect(mockGeminiClient.setHistory).not.toHaveBeenCalled(); }); it('rejects history restore while a cron prompt is mutating history', () => { @@ -6477,7 +6489,7 @@ describe('Session', () => { expect(() => session.restoreHistory([])).toThrow( 'Cannot restore history while a prompt is running', ); - expect(mockChat.setHistory).not.toHaveBeenCalled(); + expect(mockGeminiClient.setHistory).not.toHaveBeenCalled(); }); it('rejects history restore while a cron abort is active', () => { @@ -6488,7 +6500,7 @@ describe('Session', () => { expect(() => session.restoreHistory([])).toThrow( 'Cannot restore history while a prompt is running', ); - expect(mockChat.setHistory).not.toHaveBeenCalled(); + expect(mockGeminiClient.setHistory).not.toHaveBeenCalled(); }); it('rejects history restore while a notification prompt is processing', () => { @@ -6499,7 +6511,7 @@ describe('Session', () => { expect(() => session.restoreHistory([])).toThrow( 'Cannot restore history while a prompt is running', ); - expect(mockChat.setHistory).not.toHaveBeenCalled(); + expect(mockGeminiClient.setHistory).not.toHaveBeenCalled(); }); it('rejects history restore while a notification abort controller is active', () => { @@ -6510,7 +6522,7 @@ describe('Session', () => { expect(() => session.restoreHistory([])).toThrow( 'Cannot restore history while a prompt is running', ); - expect(mockChat.setHistory).not.toHaveBeenCalled(); + expect(mockGeminiClient.setHistory).not.toHaveBeenCalled(); }); }); @@ -11752,6 +11764,45 @@ describe('Session', () => { expect(sent.some((part) => 'inlineData' in part)).toBe(false); }); + it('resolves a stored image id before applying the ACP vision bridge', async () => { + const resolvedImage = { + inlineData: { mimeType: 'image/png', data: 'stored-image' }, + }; + mockChat.resolveImageReferences = vi + .fn() + .mockReturnValue([ + { text: 'inspect Image #abc123abc123' }, + resolvedImage, + ]); + mockConfig.getEffectiveInputModalities = vi.fn().mockReturnValue({}); + mockConfig.getDefaultVisionBridgeModel = vi + .fn() + .mockReturnValue({ id: 'qwen3.7-plus' }); + runVisionBridgeSpy.mockResolvedValue({ + applied: true, + status: 'ok', + parts: [{ text: '[focused transcription]' }], + convertedCount: 1, + omittedCount: 0, + modelId: 'qwen3.7-plus', + }); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValue(createEmptyStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'inspect Image #abc123abc123' }], + }); + + expect(mockChat.resolveImageReferences).toHaveBeenCalled(); + expect(runVisionBridgeSpy).toHaveBeenCalledWith( + expect.objectContaining({ + parts: expect.arrayContaining([resolvedImage]), + }), + ); + }); + it('routes an agent-capable image prompt for that ACP prompt only', async () => { const runtimeView = { contentGenerator: {}, @@ -15032,6 +15083,7 @@ describe('Session', () => { getHistory: vi.fn().mockReturnValue([]), getHistoryShallow: vi.fn().mockReturnValue([]), getLastModelMessageText: vi.fn().mockReturnValue(''), + resolveImageReferences: vi.fn((parts) => parts), } as unknown as LlmChat; mockChat.sendMessageStream = vi @@ -16160,6 +16212,7 @@ describe('Session', () => { getHistory: vi.fn().mockReturnValue([]), getHistoryShallow: vi.fn().mockReturnValue([]), getLastModelMessageText: vi.fn().mockReturnValue(''), + resolveImageReferences: vi.fn((parts) => parts), } as unknown as LlmChat; mockConfig.getSessionTokenLimit = vi.fn().mockReturnValue(100); let routeIdentity = 'route-a'; diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 1e882863250..c91da8f175e 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -4274,7 +4274,8 @@ export class Session implements SessionContext { ); } - const chat = this.config.getLlmClient()!.getChat(); + const llmClient = this.config.getLlmClient()!; + const chat = llmClient.getChat(); const apiHistory = chat.getHistoryShallow(); const apiTruncateIndex = this.#computeApiTruncationIndexForUserTurn( apiHistory, @@ -4288,7 +4289,7 @@ export class Session implements SessionContext { ); } - chat.truncateHistory(apiTruncateIndex); + llmClient.truncateHistory(apiTruncateIndex); chat.stripThoughtsFromHistory(); this.clearActiveTodoPlanRevision(); const preserveQueuedPromptPriority = this.todoStopGuardQueuedPromptPriority; @@ -4352,7 +4353,7 @@ export class Session implements SessionContext { ); } - this.config.getLlmClient()!.getChat().setHistory(structuredClone(history)); + this.config.getLlmClient()!.setHistory(structuredClone(history)); this.clearActiveTodoPlanRevision(); this.#clearTodoStopGuardTrustAndDrainAutomaticQueues(); } @@ -13932,6 +13933,9 @@ export class Session implements SessionContext { abortSignal: AbortSignal, onFullTurnModel?: (model: string) => boolean, ): Promise { + originalParts = this.#getCurrentChat().resolveImageReferences( + originalParts, + ) as Part[]; const parts = await this.#applyVoiceBridgeIfNeeded( originalParts, abortSignal, diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index e71f9f839e6..fcd8131e928 100644 --- a/packages/cli/src/nonInteractiveCli.test.ts +++ b/packages/cli/src/nonInteractiveCli.test.ts @@ -250,6 +250,7 @@ describe('runNonInteractive', () => { let processStderrSpy: MockInstance; let mockLlmClient: { sendMessageStream: Mock; + resolveImageReferences: Mock; getChatRecordingService: Mock; getChat: Mock; stripOrphanedUserEntriesFromHistory: Mock; @@ -320,6 +321,7 @@ describe('runNonInteractive', () => { mockLlmClient = { sendMessageStream: vi.fn(), + resolveImageReferences: vi.fn((parts) => parts), consumePendingMemoryTaskPromises: vi.fn().mockReturnValue([]), recordCompletedToolCall: vi.fn(), addHistory: vi.fn(), @@ -4274,6 +4276,342 @@ describe('runNonInteractive', () => { headlessImageParts, expect.any(AbortSignal), 'prompt-vision-route', + { type: SendMessageType.UserQuery, modelOverride: selector }, + ); + expect(mockLlmClient.sendMessageStream).toHaveBeenNthCalledWith( + 2, + [{ text: 'tool response' }], + expect.any(AbortSignal), + 'prompt-vision-route', + { type: SendMessageType.ToolResult, modelOverride: selector }, + ); + expect(mockCoreExecuteToolCall).toHaveBeenCalledWith( + mockConfig, + expect.objectContaining({ callId: 'vision-tool-1' }), + expect.any(AbortSignal), + expect.objectContaining({ runtimeView }), + ); + expect(processStderrSpy).toHaveBeenCalledWith( + expect.stringContaining('Routing this image turn to vision-agent'), + ); + }); + + it('does not leak a headless image route into a notification drain', async () => { + setupMetricsMock(); + await mockHeadlessImageInput(); + configureHeadlessVisionModel({ + id: 'vision-agent', + agentCapable: true, + }); + mockBackgroundTaskRegistry.setNotificationCallback.mockImplementation( + (callback) => { + callback?.('Task finished', 'task result', { + agentId: 'agent-1', + toolUseId: 'agent-tool-1', + status: 'completed', + }); + }, + ); + const drainToolCall: ServerLlmStreamEvent = { + type: LlmEventType.ToolCallRequest, + value: { + callId: 'drain-tool-1', + name: 'testTool', + args: {}, + isClientInitiated: false, + prompt_id: 'prompt-drain-isolation', + }, + }; + mockCoreExecuteToolCall.mockResolvedValue({ + responseParts: [{ text: 'drain tool response' }], + }); + mockLlmClient.sendMessageStream + .mockReturnValueOnce(createStreamFromEvents(finishedEvents)) + .mockReturnValueOnce(createStreamFromEvents([drainToolCall])) + .mockReturnValueOnce(createStreamFromEvents(finishedEvents)); + + await runNonInteractive( + mockConfig, + mockSettings, + 'inspect @image.png', + 'prompt-drain-isolation', + ); + + expect(mockLlmClient.sendMessageStream).toHaveBeenNthCalledWith( + 2, + [{ text: 'task result' }], + expect.any(AbortSignal), + 'prompt-drain-isolation/automatic/2', + expect.objectContaining({ + type: SendMessageType.Notification, + modelOverride: undefined, + }), + ); + expect(mockLlmClient.sendMessageStream).toHaveBeenNthCalledWith( + 3, + [{ text: 'drain tool response' }], + expect.any(AbortSignal), + 'prompt-drain-isolation/automatic/2', + { type: SendMessageType.ToolResult, modelOverride: undefined }, + ); + expect(mockCoreExecuteToolCall).toHaveBeenCalledWith( + mockConfig, + expect.objectContaining({ callId: 'drain-tool-1' }), + expect.any(AbortSignal), + expect.objectContaining({ runtimeView: undefined }), + ); + }); + + it('converts headless images through a non-agent vision bridge', async () => { + setupMetricsMock(); + await mockHeadlessImageInput(); + configureHeadlessVisionModel({ id: 'vision-bridge' }); + runVisionBridgeSpy.mockResolvedValue({ + applied: true, + status: 'ok', + parts: [{ text: 'machine transcription' }], + transcript: 'machine transcription', + convertedCount: 1, + omittedCount: 0, + modelId: 'vision-bridge', + egressOccurred: true, + }); + mockLlmClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(finishedEvents), + ); + + await runNonInteractive( + mockConfig, + mockSettings, + 'inspect @image.png', + 'prompt-vision-bridge', + ); + + expect(runVisionBridgeSpy).toHaveBeenCalledWith({ + config: mockConfig, + parts: headlessImageParts, + signal: expect.any(AbortSignal), + }); + expect(mockLlmClient.sendMessageStream).toHaveBeenCalledWith( + [{ text: 'machine transcription' }], + expect.any(AbortSignal), + 'prompt-vision-bridge', + { type: SendMessageType.UserQuery }, + ); + expect(processStderrSpy).toHaveBeenCalledWith( + expect.stringContaining('Converted 1 image(s)'), + ); + }); + + it('resolves a stored image id before applying the headless vision bridge', async () => { + setupMetricsMock(); + configureHeadlessVisionModel({ id: 'vision-bridge' }); + mockLlmClient.resolveImageReferences.mockReturnValue(headlessImageParts); + runVisionBridgeSpy.mockResolvedValue({ + applied: true, + status: 'ok', + parts: [{ text: 'focused transcription' }], + convertedCount: 1, + omittedCount: 0, + modelId: 'vision-bridge', + }); + mockLlmClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(finishedEvents), + ); + + await runNonInteractive( + mockConfig, + mockSettings, + 'inspect Image #abc123abc123', + 'prompt-stored-image', + ); + + expect(mockLlmClient.resolveImageReferences).toHaveBeenCalled(); + expect(runVisionBridgeSpy).toHaveBeenCalledWith({ + config: mockConfig, + parts: headlessImageParts, + signal: expect.any(AbortSignal), + }); + }); + + it('emits a stream-json system message for headless vision bridge notices', async () => { + setupMetricsMock(); + await mockHeadlessImageInput(); + (mockConfig.getOutputFormat as Mock).mockReturnValue( + OutputFormat.STREAM_JSON, + ); + configureHeadlessVisionModel({ id: 'vision-bridge' }); + runVisionBridgeSpy.mockResolvedValue({ + applied: true, + status: 'ok', + parts: [{ text: 'machine transcription' }], + transcript: 'machine transcription', + convertedCount: 1, + omittedCount: 0, + modelId: 'vision-bridge', + egressOccurred: true, + }); + mockLlmClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(finishedEvents), + ); + const writes: string[] = []; + processStdoutSpy.mockImplementation((chunk: string | Uint8Array) => { + writes.push( + typeof chunk === 'string' ? chunk : Buffer.from(chunk).toString('utf8'), + ); + return true; + }); + + await runNonInteractive( + mockConfig, + mockSettings, + 'inspect @image.png', + 'prompt-vision-bridge-json', + ); + + const systemMessage = writes + .join('') + .split('\n') + .filter(Boolean) + .map((line) => JSON.parse(line)) + .find( + (message) => + message.type === 'system' && message.subtype === 'vision_bridge', + ); + expect(systemMessage).toMatchObject({ + subtype: 'vision_bridge', + data: { notice: expect.stringContaining('Converted 1 image(s)') }, + }); + }); + + it('emits a notice when a non-agent vision bridge fails', async () => { + setupMetricsMock(); + await mockHeadlessImageInput(); + configureHeadlessVisionModel({ id: 'vision-bridge' }); + runVisionBridgeSpy.mockRejectedValue(new Error('bridge unavailable')); + mockLlmClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(finishedEvents), + ); + + await runNonInteractive( + mockConfig, + mockSettings, + 'inspect @image.png', + 'prompt-vision-bridge-failed', + ); + + expect(mockLlmClient.sendMessageStream).toHaveBeenCalledWith( + [{ text: 'inspect this image' }], + expect.any(AbortSignal), + 'prompt-vision-bridge-failed', + { type: SendMessageType.UserQuery }, + ); + expect(processStderrSpy).toHaveBeenCalledWith( + 'Vision bridge failed; proceeding without the image(s).\n', + ); + }); + + it('strips headless images when a non-agent vision bridge is skipped', async () => { + setupMetricsMock(); + await mockHeadlessImageInput(); + configureHeadlessVisionModel({ id: 'vision-bridge' }); + runVisionBridgeSpy.mockResolvedValue({ + applied: false, + status: 'skipped', + convertedCount: 0, + omittedCount: 0, + modelId: 'vision-bridge', + egressOccurred: true, + }); + mockLlmClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(finishedEvents), + ); + + await runNonInteractive( + mockConfig, + mockSettings, + 'inspect @image.png', + 'prompt-vision-bridge-skipped', + ); + + expect(mockLlmClient.sendMessageStream).toHaveBeenCalledWith( + [{ text: 'inspect this image' }], + expect.any(AbortSignal), + 'prompt-vision-bridge-skipped', + { type: SendMessageType.UserQuery }, + ); + expect(processStderrSpy).toHaveBeenCalledWith( + expect.stringContaining('Vision bridge cancelled.'), + ); + expect(processStderrSpy).toHaveBeenCalledWith( + expect.stringContaining('were sent to vision-bridge'), + ); + }); + + it('does not select a headless image route after clamping removes the image', async () => { + vi.stubEnv('QWEN_CODE_MAX_INLINE_MEDIA_BYTES', '1'); + setupMetricsMock(); + await mockHeadlessImageInput(); + const { resolveForModel } = configureHeadlessVisionModel({ + id: 'vision-agent', + agentCapable: true, + }); + mockLlmClient.sendMessageStream.mockReturnValue( + createStreamFromEvents(finishedEvents), + ); + + try { + await runNonInteractive( + mockConfig, + mockSettings, + 'inspect @image.png', + 'prompt-oversized-image', + ); + } finally { + vi.unstubAllEnvs(); + } + + expect(resolveForModel).not.toHaveBeenCalled(); + expect(mockLlmClient.sendMessageStream).toHaveBeenCalledWith( + [ + { text: 'inspect this image' }, + expect.objectContaining({ + text: expect.stringContaining('Media omitted'), + }), + ], + expect.any(AbortSignal), + 'prompt-oversized-image', + { type: SendMessageType.UserQuery }, + ); + }); + + it('fails closed when the headless image route cannot be resolved', async () => { + setupMetricsMock(); + await mockHeadlessImageInput(); + const { resolveForModel } = configureHeadlessVisionModel({ + id: 'vision-agent', + agentCapable: true, + }); + resolveForModel.mockRejectedValue(new Error('route unavailable')); + + await expect( + runNonInteractive( + mockConfig, + mockSettings, + 'inspect @image.png', + 'prompt-route-failure', + ), + ).rejects.toThrow('route unavailable'); + + expect(resolveForModel).toHaveBeenCalledWith('vision-agent', { + failClosed: true, + }); + expect(mockLlmClient.sendMessageStream).not.toHaveBeenCalled(); + }); + + it('should process input and write JSON output with stats', async () => { + const events: ServerLlmStreamEvent[] = [ + { type: LlmEventType.Content, value: 'Hello World' }, { type: SendMessageType.UserQuery, modelOverride: selector, diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index f5adcaffb13..c67c17bfede 100644 --- a/packages/cli/src/nonInteractiveCli.ts +++ b/packages/cli/src/nonInteractiveCli.ts @@ -1345,7 +1345,9 @@ export async function runNonInteractive( initialPartList = withReminder(initialPartList, recoveredAgentsNotice); } - let initialParts = normalizePartList(initialPartList); + let initialParts = normalizePartList( + llmClient.resolveImageReferences(initialPartList), + ); let fullTurnModelOverride: string | undefined; let fullTurnRuntimeView: RuntimeContentGeneratorView | undefined; const emitVisionNotice = (subtype: string, notice: string) => { diff --git a/packages/cli/src/ui/hooks/use-llm-stream.test.tsx b/packages/cli/src/ui/hooks/use-llm-stream.test.tsx index ccac4db6e88..2a80372503c 100644 --- a/packages/cli/src/ui/hooks/use-llm-stream.test.tsx +++ b/packages/cli/src/ui/hooks/use-llm-stream.test.tsx @@ -67,6 +67,7 @@ const MockedLlmClientClass = vi.hoisted(() => // _config this.startChat = mockStartChat; this.sendMessageStream = mockSendMessageStream; + this.resolveImageReferences = vi.fn((parts) => parts); this.addHistory = vi.fn(); this.consumePendingMemoryTaskPromises = vi.fn().mockReturnValue([]); this.recordCompletedToolCall = vi.fn(); @@ -878,6 +879,36 @@ describe('useLlmStream', () => { ); }); + it('resolves a stored image id before applying the vision bridge', async () => { + enableBridge(); + const client = mockConfig.getGeminiClient() as any; + client.resolveImageReferences = vi + .fn() + .mockReturnValue([{ text: 'inspect Image #abc123abc123' }, imagePart]); + mockRunVisionBridge.mockResolvedValue({ + applied: true, + status: 'ok', + parts: [{ text: '[focused transcription]' }], + convertedCount: 1, + omittedCount: 0, + modelId: 'vm', + }); + const { result } = renderTestHook([], client); + + await act(async () => { + await result.current.submitQuery('inspect Image #abc123abc123'); + }); + + expect(client.resolveImageReferences).toHaveBeenCalledWith( + 'inspect Image #abc123abc123', + ); + expect(mockRunVisionBridge).toHaveBeenCalledWith( + expect.objectContaining({ + parts: expect.arrayContaining([imagePart]), + }), + ); + }); + it('keeps an agent-capable image route through tools and retry, then clears it', async () => { enableBridge(); mockHandleSlashCommand.mockResolvedValue({ diff --git a/packages/cli/src/ui/hooks/use-llm-stream.ts b/packages/cli/src/ui/hooks/use-llm-stream.ts index a36893084a8..dffd2dcd34a 100644 --- a/packages/cli/src/ui/hooks/use-llm-stream.ts +++ b/packages/cli/src/ui/hooks/use-llm-stream.ts @@ -1436,9 +1436,11 @@ export const useLlmStream = ( timestamp: number, signal: AbortSignal, ): Promise<{ parts: PartListUnion | null; shouldProceed: boolean }> => { - if (parts === null || !hasImageParts(parts)) { + if (parts === null) { return { parts, shouldProceed: true }; } + parts = geminiClient.resolveImageReferences(parts); + if (!hasImageParts(parts)) return { parts, shouldProceed: true }; if (modelOverrideRef.current?.endsWith('\0')) { return { parts, shouldProceed: true }; } @@ -1512,7 +1514,7 @@ export const useLlmStream = ( ? { parts: textOnly, shouldProceed: true } : { parts: null, shouldProceed: false }; }, - [addItem, config], + [addItem, config, geminiClient], ); const prepareQueryForLlm = useCallback( diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index d4f42044e72..83a463e6c28 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -115,6 +115,7 @@ import { clearCacheSafeParams, getCacheSafeParams, } from '../agents/forkedAgent.js'; +import { imagePartToStoredPayload } from '../services/image-payload-references.js'; // Mock fs module to prevent actual file system operations during tests const mockFileSystem = new Map(); @@ -950,6 +951,70 @@ describe('Gemini Client (client.ts)', () => { expect(resumedClient['recentCompletedToolNames']).toEqual(['read_file']); }); + it('restores image payloads from the full resumed transcript', async () => { + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'pre-compact-shot' }, + }; + const discardedImagePart = { + inlineData: { mimeType: 'image/png', data: 'discarded-shot' }, + }; + const imageId = imagePartToStoredPayload(imagePart).id; + const discardedImageId = imagePartToStoredPayload(discardedImagePart).id; + vi.mocked(mockConfig.getResumedSessionData).mockReturnValue({ + conversation: { + sessionId: 'resumed-session-id', + projectHash: 'project-hash', + startTime: new Date(0).toISOString(), + lastUpdated: new Date(0).toISOString(), + messages: [ + { + message: { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call-screenshot', + name: 'computer_use__get_app_state', + response: { output: 'captured' }, + parts: [imagePart, discardedImagePart], + }, + }, + ], + }, + }, + { + type: 'system', + subtype: 'chat_compression', + systemPayload: { + compressedHistory: [ + { + role: 'user', + parts: [{ text: `Earlier Image #${imageId}` }], + }, + ], + }, + }, + ], + }, + filePath: '/test/session.jsonl', + lastCompletedUuid: null, + } as unknown as ReturnType); + + const resumedClient = new GeminiClient(mockConfig); + await resumedClient.initialize(); + + expect( + JSON.stringify( + resumedClient.resolveImageReferences(`inspect Image #${imageId}`), + ), + ).toContain('"data":"pre-compact-shot"'); + expect( + resumedClient.resolveImageReferences( + `inspect Image #${discardedImageId}`, + ), + ).toBe(`inspect Image #${discardedImageId}`); + }); + it('uses Startup SessionStart source for non-resumed initialize without explicit source', async () => { const hookSystem = { fireSessionStartEvent: vi.fn().mockResolvedValue( @@ -3068,6 +3133,27 @@ describe('Gemini Client (client.ts)', () => { expect(JSON.stringify(newHistory)).not.toContain('some old message'); }); + it('forgets remembered image payloads when resetting chat', async () => { + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'cleared-shot' }, + }; + const imageId = imagePartToStoredPayload(imagePart).id; + client + .getChat() + .rememberImagePayloads([{ role: 'user', parts: [imagePart] }]); + expect( + JSON.stringify( + client.resolveImageReferences(`inspect Image #${imageId}`), + ), + ).toContain('"data":"cleared-shot"'); + + await client.resetChat(); + + expect(client.resolveImageReferences(`inspect Image #${imageId}`)).toBe( + `inspect Image #${imageId}`, + ); + }); + it('clears the FileReadCache so post-reset Reads re-emit content', async () => { const cacheClear = mockFileReadCacheClear(); @@ -3182,6 +3268,7 @@ describe('Gemini Client (client.ts)', () => { const cacheClear = mockFileReadCacheClear(); client['chat'] = { setHistory: vi.fn(), + reconcileImagePayloads: vi.fn(), } as unknown as LlmChat; client.setHistory([{ role: 'user', parts: [{ text: 'replaced' }] }]); @@ -3189,6 +3276,89 @@ describe('Gemini Client (client.ts)', () => { expect(cacheClear).toHaveBeenCalled(); }); + it('setHistory forgets image payloads removed by restore', () => { + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'rewound-shot' }, + }; + const imageId = imagePartToStoredPayload(imagePart).id; + client + .getChat() + .rememberImagePayloads([{ role: 'user', parts: [imagePart] }]); + + client.setHistory([{ role: 'user', parts: [{ text: 'restored' }] }]); + + expect(client.resolveImageReferences(`inspect Image #${imageId}`)).toBe( + `inspect Image #${imageId}`, + ); + }); + + it('setHistory retains payloads still referenced after compaction', () => { + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'compressed-shot' }, + }; + const imageId = imagePartToStoredPayload(imagePart).id; + client + .getChat() + .rememberImagePayloads([{ role: 'user', parts: [imagePart] }]); + + client.setHistory([ + { + role: 'user', + parts: [{ text: `Earlier Image #${imageId}` }], + }, + ]); + + expect( + JSON.stringify( + client.resolveImageReferences(`inspect Image #${imageId}`), + ), + ).toContain('"data":"compressed-shot"'); + }); + + it('truncateHistory drops removed image payloads but keeps survivors', () => { + const keptImage = { + inlineData: { mimeType: 'image/png', data: 'kept-shot' }, + }; + const removedImage = { + inlineData: { mimeType: 'image/png', data: 'removed-shot' }, + }; + const keptId = imagePartToStoredPayload(keptImage).id; + const removedId = imagePartToStoredPayload(removedImage).id; + client.setHistory([ + { role: 'user', parts: [keptImage] }, + { role: 'model', parts: [{ text: 'seen' }] }, + { role: 'user', parts: [removedImage] }, + ]); + + client.truncateHistory(2); + + expect( + JSON.stringify( + client.resolveImageReferences(`inspect Image #${keptId}`), + ), + ).toContain('"data":"kept-shot"'); + expect(client.resolveImageReferences(`inspect Image #${removedId}`)).toBe( + `inspect Image #${removedId}`, + ); + }); + + it('stripOrphanedUserEntries drops image payloads from removed turns', () => { + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'orphan-shot' }, + }; + const imageId = imagePartToStoredPayload(imagePart).id; + client.setHistory([ + { role: 'model', parts: [{ text: 'done' }] }, + { role: 'user', parts: [imagePart] }, + ]); + + client.stripOrphanedUserEntriesFromHistory(); + + expect(client.resolveImageReferences(`inspect Image #${imageId}`)).toBe( + `inspect Image #${imageId}`, + ); + }); + /** * Test helper: mock a LlmChat whose history length goes from * `before` to `after` across truncateHistory(). The first @@ -3717,10 +3887,12 @@ describe('Gemini Client (client.ts)', () => { const { history } = await makeReadFileResponses(6); const setHistory = vi.fn(); + const reconcileImagePayloads = vi.fn(); client['chat'] = { addHistory: vi.fn(), getHistory: vi.fn().mockReturnValue(history), setHistory, + reconcileImagePayloads, } as unknown as LlmChat; client['lastApiCompletionTimestamp'] = Date.now() - 90 * 60_000; @@ -3735,6 +3907,9 @@ describe('Gemini Client (client.ts)', () => { } expect(setHistory).toHaveBeenCalled(); + expect(reconcileImagePayloads).toHaveBeenCalledWith( + setHistory.mock.calls[0][0], + ); // The blanket wipe is gone — read-before-write state is preserved. expect(clear).not.toHaveBeenCalled(); // Exactly the one blanked file (oldest of 6, keepRecent=5) had its @@ -4750,6 +4925,7 @@ describe('Gemini Client (client.ts)', () => { }), isLastPromptTokenCountEstimated: vi.fn().mockReturnValue(false), getHistory: vi.fn().mockReturnValue([]), + copyImagePayloadsTo: vi.fn(), } as unknown as LlmChat; client['forceFullIdeContext'] = false; @@ -4796,6 +4972,38 @@ describe('Gemini Client (client.ts)', () => { expect(client['forceFullIdeContext']).toBe(true); }); + it('preserves remembered image payloads across manual compression', async () => { + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'pre-compress-shot' }, + }; + const imageId = imagePartToStoredPayload(imagePart).id; + const originalChat = client.getChat(); + originalChat.rememberImagePayloads([ + { role: 'user', parts: [imagePart] }, + ]); + vi.spyOn(originalChat, 'tryCompress').mockImplementation(async () => { + originalChat.setHistory([ + { + role: 'user', + parts: [{ text: `Earlier Image #${imageId}` }], + }, + ]); + return { + originalTokenCount: 1000, + newTokenCount: 200, + compressionStatus: CompressionStatus.COMPRESSED, + }; + }); + + await client.tryCompressChat('p4'); + + expect( + JSON.stringify( + client.resolveImageReferences(`inspect Image #${imageId}`), + ), + ).toContain('"data":"pre-compress-shot"'); + }); + it('preserves Compact SessionStart additionalContext on the new chat', async () => { const compressedHistory: Content[] = [ { role: 'user', parts: [{ text: 'summary' }] }, diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 2630a262b14..81e0863fd75 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -574,6 +574,12 @@ export class LlmClient { ); this.restoreLoadedSkillsFromHistory(resumedHistory); const chat = this.getChat(); + chat.rememberImagePayloads( + resumedSessionData.conversation.messages.flatMap((record) => + record.message ? [record.message] : [], + ), + ); + chat.reconcileImagePayloads(chat.getHistory()); if (resumeTokenCounts) { chat.seedResumeTokenCounts( resumeTokenCounts.promptTokenCount, @@ -651,6 +657,10 @@ export class LlmClient { return this.chat; } + resolveImageReferences(message: PartListUnion): PartListUnion { + return this.getChat().resolveImageReferences(message); + } + isInitialized(): boolean { return this.chat !== undefined; } @@ -1059,7 +1069,9 @@ export class LlmClient { } setHistory(history: Content[]) { - this.getChat().setHistory(history); + const chat = this.getChat(); + chat.setHistory(history); + chat.reconcileImagePayloads(history); // Replacing history wholesale drops any prior read_file tool // results the FileReadCache still believes the model has seen. // Without clearing, a follow-up Read of an unchanged file would @@ -2739,7 +2751,9 @@ export class LlmClient { const changed = m.tokensSaved > 0; if (changed) { // setHistory conservatively clears loaded-skill tracking. - this.getChat().setHistory(mcResult.history); + const chat = this.getChat(); + chat.setHistory(mcResult.history); + chat.reconcileImagePayloads?.(mcResult.history); await this.disarmFileReadCacheAfterEviction(m, 'microcompaction'); } if (m.triggerReason === 'size') { @@ -5009,6 +5023,7 @@ export class LlmClient { const compressedHistory = previousChat.getHistoryShallow?.() ?? previousChat.getHistory(); await this.startChat(compressedHistory, SessionStartSource.Compact); + chat.copyImagePayloadsTo(this.getChat()); if ( !this.lastSessionStartContext && previousSessionStartContext && diff --git a/packages/core/src/core/llm-chat.test.ts b/packages/core/src/core/llm-chat.test.ts index 390c8680115..c278c7314d5 100644 --- a/packages/core/src/core/llm-chat.test.ts +++ b/packages/core/src/core/llm-chat.test.ts @@ -50,6 +50,7 @@ import { getToolCallPreparations, setToolCallPreparations, } from './tool-call-preparation.js'; +import { imagePartToStoredPayload } from '../services/image-payload-references.js'; import { ApprovalMode } from '../config/approval-mode.js'; // Mock fs module to prevent actual file system operations during tests @@ -200,6 +201,10 @@ describe('LlmChat', async () => { toolResultsNumToKeep: 1, }), getAutoCompactThreshold: vi.fn().mockReturnValue(undefined), + getClearContextOnIdle: vi.fn().mockReturnValue({ + toolResultsThresholdMinutes: 60, + toolResultsNumToKeep: 5, + }), getHookSystem: vi.fn().mockReturnValue(undefined), getDebugLogger: vi .fn() @@ -3684,6 +3689,118 @@ describe('LlmChat', async () => { }); }); + it('adds stable refs before the image payload threshold', async () => { + vi.mocked(mockConfig.getChatCompression).mockReturnValue({ + maxRecentImagesToRetain: 1, + imagePayloadThreshold: 20, + }); + chat.setHistory([ + { + role: 'user', + parts: [{ inlineData: { mimeType: 'image/png', data: 'only-shot' } }], + }, + { role: 'model', parts: [{ text: 'seen' }] }, + ]); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + streamResponse(stopResponse([{ text: 'response' }])), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'continue' }, + 'prompt-id-image-ref-below-threshold', + ); + for await (const _ of stream) { + // consume stream + } + + const contents = vi.mocked(mockContentGenerator.generateContentStream) + .mock.calls[0]?.[0].contents as Content[]; + const serialized = JSON.stringify(contents); + expect(serialized).toMatch(/Image #[a-f0-9]{12}/); + expect(serialized.match(/"data":"only-shot"/g)).toHaveLength(1); + }); + + it('resolves remembered image ids and forgets them on clear', () => { + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'resumed-shot' }, + }; + const stored = imagePartToStoredPayload(imagePart); + const refHistory: Content[] = [ + { + role: 'user', + parts: [ + { + text: `[Image #${stored.id}: image/png, ${stored.bytes} bytes]`, + }, + ], + }, + { role: 'model', parts: [{ text: 'seen' }] }, + ]; + chat.rememberImagePayloads([{ role: 'user', parts: [imagePart] }]); + chat.setHistory(refHistory); + + expect( + JSON.stringify( + chat.resolveImageReferences(`inspect Image #${stored.id}`), + ), + ).toContain('"data":"resumed-shot"'); + + chat.clearHistory(); + chat.setHistory(refHistory); + expect(chat.resolveImageReferences(`inspect Image #${stored.id}`)).toBe( + `inspect Image #${stored.id}`, + ); + }); + + it('sends only explicitly selected images below the payload threshold', async () => { + vi.mocked(mockConfig.getChatCompression).mockReturnValue({ + maxRecentImagesToRetain: 3, + imagePayloadThreshold: 20, + }); + const images = ['shot-a', 'shot-b', 'shot-c'].map((data) => ({ + inlineData: { mimeType: 'image/png', data }, + })); + const ids = images.map((image) => imagePartToStoredPayload(image).id); + const history: Content[] = [ + { role: 'user', parts: [images[0]] }, + { role: 'model', parts: [{ text: 'seen 0' }] }, + { role: 'user', parts: [images[1]] }, + { role: 'model', parts: [{ text: 'seen 1' }] }, + { + role: 'user', + parts: [ + { text: `Recent image restored: Image #${ids[2]}` }, + images[2], + ], + }, + ]; + chat.setHistory(history); + chat.rememberImagePayloads(history); + const resolved = chat.resolveImageReferences( + `compare Image #${ids[0]} with Image #${ids[2]}`, + ); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + streamResponse(stopResponse([{ text: 'response' }])), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: resolved }, + 'prompt-id-multi-image-focus', + ); + for await (const _ of stream) { + // consume stream + } + + const contents = vi.mocked(mockContentGenerator.generateContentStream) + .mock.calls[0]?.[0].contents as Content[]; + const serialized = JSON.stringify(contents); + expect(serialized.match(/"data":"shot-a"/g)).toHaveLength(1); + expect(serialized).not.toContain('"data":"shot-b"'); + expect(serialized.match(/"data":"shot-c"/g)).toHaveLength(1); + }); + it('reattaches stored image markers on later below-threshold requests', async () => { vi.mocked(mockConfig.getChatCompression).mockReturnValue({ maxRecentImagesToRetain: 1, @@ -16681,6 +16798,41 @@ describe('LlmChat', async () => { expect(chat.getLastPromptTokenCount()).toBe(200); }); + it('drops image payloads removed by compression', async () => { + mockCompressionService('compressed'); + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'discarded-shot' }, + }; + const imageId = imagePartToStoredPayload(imagePart).id; + chat.rememberImagePayloads([{ role: 'user', parts: [imagePart] }]); + + await chat.tryCompress('p-image', 'm1'); + + expect(chat.resolveImageReferences(`inspect Image #${imageId}`)).toBe( + `inspect Image #${imageId}`, + ); + }); + + it('drops image payloads removed by fast compression', () => { + const imagePart = { + inlineData: { mimeType: 'image/png', data: 'discarded-fast-shot' }, + }; + const imageId = imagePartToStoredPayload(imagePart).id; + chat.rememberImagePayloads([{ role: 'user', parts: [imagePart] }]); + chat.setHistory([ + { + role: 'model', + parts: [{ text: 'thinking', thought: true }, { text: 'answer' }], + }, + ]); + + chat.compressFast(); + + expect(chat.resolveImageReferences(`inspect Image #${imageId}`)).toBe( + `inspect Image #${imageId}`, + ); + }); + it('mirrors lastPromptTokenCount to the global telemetry only when wired', async () => { mockCompressionService('compressed'); // chat under test was constructed with telemetryService=uiTelemetryService. diff --git a/packages/core/src/core/llm-chat.ts b/packages/core/src/core/llm-chat.ts index 29753038313..ca26a1ebf69 100644 --- a/packages/core/src/core/llm-chat.ts +++ b/packages/core/src/core/llm-chat.ts @@ -14,6 +14,7 @@ import type { FunctionCall, SendMessageParameters, Part, + PartListUnion, Tool, GenerateContentResponseUsageMetadata, } from '@google/genai'; @@ -86,7 +87,10 @@ import { import { InMemoryImagePayloadStore, buildReattachParts, + collectReferencedImageIds, countAllInlineImages, + prepareImagePayloadsForRequest, + rememberImagePayloads, replaceImagePayloadsInPlace, } from '../services/image-payload-references.js'; import { @@ -2222,26 +2226,39 @@ export class LlmChat { const { maxRecentImages, imagePayloadThreshold } = resolveCompactionTuning( this.config.getChatCompression(), ); - let replaced: ReturnType = []; - if (countAllInlineImages(curatedHistory) >= imagePayloadThreshold) { - const skipEntry = currentUserContent - ? curatedHistory.find( - (c) => - c === currentUserContent || - (c.role === 'user' && - currentUserContent.parts?.some((p) => c.parts?.includes(p))), - ) - : undefined; - replaced = replaceImagePayloadsInPlace( - curatedHistory, - this.imagePayloadStore, - skipEntry, - ); - } + // History always holds `Image #` markers rather than raw bytes: that + // is what survives compaction, truncation and resume. The threshold only + // decides how many payloads are reattached to the outgoing request. + const imageCount = countAllInlineImages(curatedHistory); + const skipEntry = currentUserContent + ? curatedHistory.find( + (c) => + c === currentUserContent || + (c.role === 'user' && + currentUserContent.parts?.some((p) => c.parts?.includes(p))), + ) + : undefined; + const replaced = replaceImagePayloadsInPlace( + curatedHistory, + this.imagePayloadStore, + skipEntry, + ); const requestHistory = curatedHistory.map(copyContentContainer); + // A prompt naming explicit image ids attaches only those, so the recent + // window contributes nothing; below the threshold every historical image + // stays attached; at or above it only the configured recent ones do. + const hasExplicitReferences = + collectReferencedImageIds( + requestHistory.at(-1) ? [requestHistory.at(-1)!] : [], + ).size > 0; + const reattachCount = hasExplicitReferences + ? 0 + : imageCount >= imagePayloadThreshold + ? maxRecentImages + : imageCount; const reattachParts = buildReattachParts( replaced, - maxRecentImages, + reattachCount, requestHistory, this.imagePayloadStore, ); @@ -2256,6 +2273,40 @@ export class LlmChat { return requestHistory; } + /** + * Resolve `Image #` references the user typed into the outgoing message + * back into their stored payloads, so an explicitly named historical image + * is re-sent even though history carries only its marker. + */ + resolveImageReferences(message: PartListUnion): PartListUnion { + const current = createUserContent(message); + if (collectReferencedImageIds([current]).size === 0) { + return message; + } + const history = extractCuratedHistory(this.history); + const resolved = prepareImagePayloadsForRequest([...history, current], { + maxRecentImages: 0, + preserveImagePartsForContentIndex: history.length, + store: this.imagePayloadStore, + }).at(-1)?.parts; + return resolved?.some((part) => part.inlineData) ? resolved : message; + } + + /** Absorb raw payloads in `contents` into the store without rewriting them. */ + rememberImagePayloads(contents: Content[]): void { + rememberImagePayloads(contents, this.imagePayloadStore); + } + + /** Drop stored payloads no longer referenced by `contents`. */ + reconcileImagePayloads(contents: Content[]): void { + this.imagePayloadStore.reconcile(contents); + } + + /** Seed a forked chat's store so it can resolve the parent's references. */ + copyImagePayloadsTo(target: LlmChat): void { + this.imagePayloadStore.copyTo(target.imagePayloadStore); + } + private getRequestHistoryForRoute( currentUserContent: Content | undefined, supportedModalities: InputModalities, @@ -2426,6 +2477,7 @@ export class LlmChat { }); } this.setHistory(newHistory); + this.reconcileImagePayloads(newHistory); debugLogger.debug('[FILE_READ_CACHE] clear after auto tryCompress'); this.config.getFileReadCache().clear(); // Compression rewrote the shared history every retained entry sizes, @@ -2565,6 +2617,7 @@ export class LlmChat { }), ); this.setHistory(newHistory); + this.reconcileImagePayloads(newHistory); this.lastPromptTokenCount = adjustedTokenCount; this.lastPromptTokenCountIsEstimated = true; this.tokenCountsRouteKey = this.currentRouteKey(); @@ -4886,6 +4939,7 @@ export class LlmChat { */ clearHistory(): void { this.history = []; + this.imagePayloadStore.clear(); // Any pending partial-push state points into the now-empty history; // resetting prevents `popPendingPartialAssistantTurn` from splicing whatever // shows up at that index in a future send (defense-in-depth — the @@ -5061,6 +5115,7 @@ export class LlmChat { truncateHistory(keepCount: number): void { const prevLen = this.history.length; this.history = this.history.slice(0, keepCount); + this.reconcileImagePayloads(this.history); // Truncation can drop the entry the partial-push marker points at, // or leave it valid but shift the meaning of nearby indices. Reset // both fields rather than try to fix them up — they're per-send and @@ -5082,6 +5137,7 @@ export class LlmChat { this.history = this.history .map(stripThoughtPartsFromContent) .filter((content): content is Content => content !== null); + this.reconcileImagePayloads(this.history); // Filter+map replaces `this.history` with a new array, so any pending // partial-push marker is now indexed against an array that no longer // exists. Clear it for the same reason setHistory does — and drop @@ -5142,6 +5198,9 @@ export class LlmChat { ); } this.clearPendingPartialState(); + if (strippedEntries.length > 0) { + this.reconcileImagePayloads(this.history); + } return strippedEntries; } diff --git a/packages/core/src/services/image-payload-references.test.ts b/packages/core/src/services/image-payload-references.test.ts index 95699c5884c..3f38ef9be3f 100644 --- a/packages/core/src/services/image-payload-references.test.ts +++ b/packages/core/src/services/image-payload-references.test.ts @@ -9,6 +9,7 @@ import { describe, expect, it } from 'vitest'; import { InMemoryImagePayloadStore, buildReattachParts, + collectReferencedImageIds, countAllInlineImages, prepareImagePayloadsForRequest, replaceImagePayloadsInPlace, @@ -49,6 +50,23 @@ function imageParts(contents: Content[]): Part[] { } describe('prepareImagePayloadsForRequest', () => { + it('collects image references nested in function responses', () => { + const ids = collectReferencedImageIds({ + role: 'user', + parts: [ + { + functionResponse: { + name: 'screenshot', + response: {}, + parts: [{ text: 'Inspect Image #abcdef123456' }], + }, + }, + ], + }); + + expect([...ids]).toEqual(['abcdef123456']); + }); + it('replaces historical image positions with stable refs and reattaches only the most recent images', () => { const store = new InMemoryImagePayloadStore(); const history: Content[] = [ @@ -110,6 +128,36 @@ describe('prepareImagePayloadsForRequest', () => { ]); }); + it('reattaches only the explicitly referenced images', () => { + const store = new InMemoryImagePayloadStore(); + const history = prepareImagePayloadsForRequest( + [ + toolImageTurn('shot-a'), + toolImageTurn('shot-b'), + toolImageTurn('shot-c'), + { role: 'model', parts: [{ text: 'done' }] }, + ], + { maxRecentImages: 0, store }, + ); + const ids = JSON.stringify(history).match(/Image #([a-f0-9]{12})/g) ?? []; + + const prepared = prepareImagePayloadsForRequest( + [ + ...history, + { + role: 'user', + parts: [{ text: `compare ${ids[0]} with ${ids[2]}` }], + }, + ], + { maxRecentImages: 0, store }, + ); + + expect(imageParts(prepared).map((part) => part.inlineData?.data)).toEqual([ + 'shot-a', + 'shot-c', + ]); + }); + it('reattaches a stored image when only its stable reference remains in history', () => { const store = new InMemoryImagePayloadStore(); const firstPass = prepareImagePayloadsForRequest( @@ -180,31 +228,24 @@ describe('prepareImagePayloadsForRequest', () => { ]); }); - it('preserves images in the current user request when maxRecentImages is zero', () => { + it('does not append historical image payloads to the live current turn', () => { const store = new InMemoryImagePayloadStore(); + const current: Content = { + role: 'user', + parts: [{ text: 'continue' }], + }; + const prepared = prepareImagePayloadsForRequest( - [ - toolImageTurn('old-shot'), - { - role: 'user', - parts: [ - { text: 'inspect this' }, - { inlineData: { mimeType: 'image/png', data: 'current-shot' } }, - ], - }, - ], + [toolImageTurn('old-shot'), current], { - maxRecentImages: 0, - preserveLastUserImagePartCount: 2, + maxRecentImages: 1, + preserveImagePartsForContentIndex: 1, store, }, ); - const serialized = JSON.stringify(prepared); - expect(serialized).not.toContain('"data":"old-shot"'); - expect(imageParts(prepared).map((part) => part.inlineData?.data)).toEqual([ - 'current-shot', - ]); + expect(prepared.at(-1)?.parts).toHaveLength(3); + expect(current.parts).toEqual([{ text: 'continue' }]); }); it('does not echo tool-controlled image metadata into text references', () => { diff --git a/packages/core/src/services/image-payload-references.ts b/packages/core/src/services/image-payload-references.ts index 0e36f321dc3..dacaae3ed7b 100644 --- a/packages/core/src/services/image-payload-references.ts +++ b/packages/core/src/services/image-payload-references.ts @@ -49,6 +49,30 @@ export class InMemoryImagePayloadStore implements ImagePayloadStore { get(id: string): StoredImagePayload | undefined { return this.images.get(id); } + + clear(): void { + this.images.clear(); + } + + copyTo(target: InMemoryImagePayloadStore): void { + for (const [id, image] of this.images) { + target.images.set(id, image); + } + } + + /** + * Drop payloads no longer referenced by `contents` and absorb any raw + * payloads still inline in them. Call after any operation that replaces + * history wholesale (compaction, truncation, thought stripping) so evicted + * references do not pin their bytes for the rest of the session. + */ + reconcile(contents: Content[]): void { + const referencedIds = collectReferencedImageIds(contents); + for (const id of this.images.keys()) { + if (!referencedIds.has(id)) this.images.delete(id); + } + rememberImagePayloads(contents, this); + } } export function countAllInlineImages(contents: Content[]): number { @@ -114,7 +138,9 @@ export function buildReattachParts( const recent = recentUniqueImages(candidates, maxRecentImages).map( ({ stored }) => stored, ); - const reattachLimit = Math.max(maxRecentImages, 1); + // An explicit multi-image prompt must keep every id it named; the recent + // window alone would shift all but the last one back out. + const reattachLimit = Math.max(maxRecentImages, lastReferencedIds.size, 1); if (store) { for (const id of lastReferencedIds) { if (inlineIds.has(id) || recent.some((image) => image.id === id)) { @@ -278,7 +304,7 @@ function* inlineImageParts( } } -function collectReferencedImageIds(contents: Content[]): Set { +export function collectReferencedImageIds(contents: Content[]): Set { const ids = new Set(); const collect = (parts: Part[] | undefined): void => { for (const part of parts ?? []) { @@ -314,7 +340,7 @@ function recentUniqueImages( return recent.reverse(); } -function imagePartToStoredPayload(part: Part): StoredImagePayload { +export function imagePartToStoredPayload(part: Part): StoredImagePayload { const data = part.inlineData?.data ?? ''; const mimeType = part.inlineData?.mimeType ?? 'application/octet-stream'; const hash = createHash('sha256') @@ -350,3 +376,32 @@ function storedImageToPart(stored: StoredImagePayload): Part { }, }; } + +/** + * Absorb every raw image payload in `contents` into `store` without rewriting + * the contents. Used on resume, where history is rebuilt from the original + * JSONL and the store must be repopulated before references are resolved. + */ +export function rememberImagePayloads( + contents: Content[], + store: ImagePayloadStore, +): void { + for (const content of contents) { + for (const part of content.parts ?? []) { + if ( + part.inlineData?.mimeType?.startsWith('image/') && + part.inlineData.data + ) { + store.put(part); + } + for (const nested of getFunctionResponseParts(part) ?? []) { + if ( + nested.inlineData?.mimeType?.startsWith('image/') && + nested.inlineData.data + ) { + store.put(nested); + } + } + } + } +} diff --git a/packages/core/src/services/memoryPressureMonitor.test.ts b/packages/core/src/services/memoryPressureMonitor.test.ts index af5e0d9a944..268d4891e97 100644 --- a/packages/core/src/services/memoryPressureMonitor.test.ts +++ b/packages/core/src/services/memoryPressureMonitor.test.ts @@ -143,6 +143,7 @@ function createMockConfig( getHistoryShallow?: () => unknown[]; getHistory?: () => unknown[]; setHistory?: (h: unknown[]) => void; + reconcileImagePayloads?: (h: unknown[]) => void; }; } | null; clearContextOnIdle?: { @@ -160,6 +161,7 @@ function createMockConfig( getHistoryShallow: () => [], getHistory: () => [], setHistory: vi.fn(), + reconcileImagePayloads: vi.fn(), }), } : overrides.llmClient; @@ -1328,9 +1330,11 @@ describe('MemoryPressureMonitor', () => { { ...DEFAULT_PRESSURE_CONFIG, cleanupCooldownMs: 0 }, ); - setMemUsage(11 * 1024 * 1024 * 1024); // 11/16 = 0.6875 >= 0.65: hard pressure - monitor.performCheck(); - await drainCleanupMeasurement(); + ( + monitor as unknown as { + executeStep(step: 'compact_history'): void; + } + ).executeStep('compact_history'); // compact_history error is caught and logged without propagating, // so subsequent cleanup steps (like trigger_gc) can still run. @@ -1346,15 +1350,18 @@ describe('MemoryPressureMonitor', () => { { ...DEFAULT_PRESSURE_CONFIG, cleanupCooldownMs: 0 }, ); - setMemUsage(11 * 1024 * 1024 * 1024); // 11/16 = 0.6875 >= 0.65: hard pressure - monitor.performCheck(); - await drainCleanupMeasurement(); + ( + monitor as unknown as { + executeStep(step: 'compact_history'): void; + } + ).executeStep('compact_history'); expect(setHistory).not.toHaveBeenCalled(); }); it('compacts history and clears fileReadCache when meta is non-null', async () => { const setHistory = vi.fn(); + const reconcileImagePayloads = vi.fn(); const clearCache = vi.fn(); // Build history with 7 read_file tool results (keep=5, so 2 get cleared) const toolHistory: Content[] = []; @@ -1397,6 +1404,7 @@ describe('MemoryPressureMonitor', () => { getChat: () => ({ getHistoryShallow: () => toolHistory, setHistory, + reconcileImagePayloads, }), }, fileReadCache: { @@ -1416,6 +1424,9 @@ describe('MemoryPressureMonitor', () => { await drainCleanupMeasurement(); expect(setHistory).toHaveBeenCalled(); + expect(reconcileImagePayloads).toHaveBeenCalledWith( + setHistory.mock.calls[0][0], + ); expect(clearCache).toHaveBeenCalled(); const compacted = setHistory.mock.calls[0][0] as Content[]; // microcompactHistory blanks old tool responses with a cleared message diff --git a/packages/core/src/services/memoryPressureMonitor.ts b/packages/core/src/services/memoryPressureMonitor.ts index 393cf2f3f6a..f74bff2a381 100644 --- a/packages/core/src/services/memoryPressureMonitor.ts +++ b/packages/core/src/services/memoryPressureMonitor.ts @@ -736,6 +736,7 @@ export class MemoryPressureMonitor extends EventEmitter { ); if (result.meta) { chat.setHistory(result.history); + chat.reconcileImagePayloads?.(result.history); // Explicitly clear fileReadCache here instead of relying on // the subsequent clear_file_cache step. This removes the // implicit coupling between step ordering. diff --git a/packages/core/src/services/postCompactAttachments.test.ts b/packages/core/src/services/postCompactAttachments.test.ts index 76d1acebd82..30a5a41d236 100644 --- a/packages/core/src/services/postCompactAttachments.test.ts +++ b/packages/core/src/services/postCompactAttachments.test.ts @@ -646,6 +646,7 @@ import { buildImageRestorationBlock, type ExtractedImage, } from './postCompactAttachments.js'; +import { imagePartToStoredPayload } from './image-payload-references.js'; describe('buildImageRestorationBlock', () => { it('returns null when no images are provided', () => { @@ -674,6 +675,12 @@ describe('buildImageRestorationBlock', () => { const header = (block!.parts![0] as { text: string }).text; expect(header).toContain('Recent visual snapshots'); + expect(header).toContain( + `- Image #${imagePartToStoredPayload(images[0].part).id}, turn 5: computer_use__get_app_state args={"app":"Safari"}`, + ); + expect(header).toContain( + `- Image #${imagePartToStoredPayload(images[1].part).id}, turn 11: computer_use__get_app_state args={"app":"Mail"}`, + ); expect(header).toContain('turn 5'); expect(header).toContain('mcp__node-repl__node_repl'); expect(header).toContain('"app":"Safari"'); @@ -693,8 +700,9 @@ describe('buildImageRestorationBlock', () => { ]; const block = buildImageRestorationBlock(images); const header = (block!.parts![0] as { text: string }).text; - expect(header).toContain('turn 3'); - expect(header).toContain('user-provided'); // labeled instead of tool name + expect(header).toContain( + `- Image #${imagePartToStoredPayload(images[0].part).id}, turn 3: user-provided image`, + ); }); }); diff --git a/packages/core/src/services/postCompactAttachments.ts b/packages/core/src/services/postCompactAttachments.ts index 5cc4f7f7896..592925122b2 100644 --- a/packages/core/src/services/postCompactAttachments.ts +++ b/packages/core/src/services/postCompactAttachments.ts @@ -21,6 +21,7 @@ import { realpathSync } from 'node:fs'; import { resolve as resolvePath, sep as pathSep } from 'node:path'; import { CHARS_PER_TOKEN } from './tokenEstimation.js'; import { getFunctionResponseParts } from './compactionInputSlimming.js'; +import { imagePartToStoredPayload } from './image-payload-references.js'; import { escapeXml } from '../utils/xml.js'; import { ToolNames } from '../tools/tool-names.js'; @@ -443,13 +444,16 @@ export function buildImageRestorationBlock( '', ]; for (const img of images) { + const imageId = imagePartToStoredPayload(img.part).id; if (img.sourceToolName) { const argsStr = JSON.stringify(img.sourceToolArgs ?? {}); lines.push( - `- turn ${img.turnIndex}: ${img.sourceToolName} args=${argsStr}`, + `- Image #${imageId}, turn ${img.turnIndex}: ${img.sourceToolName} args=${argsStr}`, ); } else { - lines.push(`- turn ${img.turnIndex}: user-provided image`); + lines.push( + `- Image #${imageId}, turn ${img.turnIndex}: user-provided image`, + ); } }