diff --git a/.qwen/e2e-tests/terminal-inline-images.md b/.qwen/e2e-tests/terminal-inline-images.md new file mode 100644 index 00000000000..bcf3506ea5b --- /dev/null +++ b/.qwen/e2e-tests/terminal-inline-images.md @@ -0,0 +1,97 @@ +# Terminal Inline Images E2E Plan + +## Baseline + +1. Start the interactive CLI from `main`. +2. Return an assistant response containing + `text -> inlineData(image/png) -> text`. +3. Run a tool whose `functionResponse.parts` contains an `image/png`. +4. Confirm that `main` omits the assistant image and reduces tool media to + text, while the separate `display_image` tool from #8217 can display a + workspace PNG path. + +The global `qwen` executable is unavailable in the current environment, so the +baseline is grounded in issue #8090 and the unchanged `main` event mapping. + +## Verification + +### Shared renderer regression + +1. Ask the model to call `display_image` for a workspace PNG. +2. In direct Kitty or Ghostty, confirm the existing native preview still + renders. +3. In a non-native terminal with `chafa` installed, confirm the existing ANSI + preview still renders. +4. Confirm the file-path tool retains its workspace, file-size, and PNG + validation behavior. + +### Inline assistant and tool images + +1. Return a 1x1 PNG between two assistant text parts. +2. Confirm the transcript order is `text -> image -> text`. +3. Return a PNG in a successful tool's top-level and nested + `functionResponse.parts`. +4. Confirm the successful tool row retains its images. +5. Open Ctrl+O and resume the session; confirm successful tool image order is + reconstructed from persisted parts. Assistant output resumes its persisted + text; assistant inline images are not persisted by the current Core recorder. +6. Return six images in one assistant output and one tool response; confirm the + first four render and the row ends with `[+2 more images]`. + +### Kitty/Ghostty + +1. Start the CLI in direct Kitty or Ghostty without tmux or SSH. +2. Repeat the assistant and tool cases. +3. Confirm inline PNGs use the same virtual placement and Unicode placeholders + as `display_image`. +4. Confirm remounting a history row does not retransmit an already-written + payload. + +### chafa + +1. Start the CLI in Warp, iTerm2, tmux, SSH, or another non-native environment + with `chafa` installed. +2. Repeat the assistant and tool cases. +3. Confirm PNG bytes render as ANSI symbol rows and stay aligned during normal + scrolling. +4. Confirm image data is supplied through stdin and no model-controlled value + is used as a command argument. + +### Placeholders and accessibility + +1. Repeat without `chafa` and outside direct Kitty/Ghostty. +2. Confirm a valid PNG displays `[image: x png]`. +3. Repeat with malformed base64, a payload above 8 MiB, invalid IHDR + dimensions, and a non-PNG MIME type. +4. Confirm no raw image sequence is written. Confirm the oversized payload is + dropped before UI history, while admitted malformed/non-PNG data uses a + deterministic placeholder. +5. Repeat with `INK_SCREEN_READER=true`; confirm only the placeholder is + emitted. + +### Stream lifecycle + +1. Exercise a fresh retry, continuation retry, model fallback, cancellation, + stream boundary, tool call boundary, and displayed goal-state event. +2. Confirm fresh attempts discard all staged image/text runs from the failed + attempt. +3. Confirm continuations preserve partial output. +4. Confirm every normal boundary commits earlier runs before later status or + tool rows. + +## Automated Evidence + +Record final results for: + +- focused CLI renderer, component, stream, tool, resume, and compaction tests; +- Core `Turn`, `display_image`, config, scheduler, and tool tests; +- `npm run lint:ci`; +- `npm run typecheck`; +- `npm run build`; +- `npm run check:serve-fast-path-bundle`. + +Real-terminal output remains a reviewer hardware step when the local environment +has no Kitty/Ghostty session. Reusing #8217's renderer means its existing manual +Kitty, Ghostty, cmux, and Warp evidence continues to cover the terminal protocol +layer; this plan focuses new manual verification on the inline-data entry point +and transcript lifecycle. diff --git a/docs/design/terminal-inline-images.md b/docs/design/terminal-inline-images.md new file mode 100644 index 00000000000..3e646902030 --- /dev/null +++ b/docs/design/terminal-inline-images.md @@ -0,0 +1,138 @@ +# Terminal Inline Images + +## Problem + +The interactive CLI drops model `inlineData` image parts at the +`Turn`-to-TUI boundary. Images nested in tool `functionResponse.parts` +survive in model history, but the tool display reduces them to text. As a +result, image-generating models and screenshot-producing tools cannot show +their output in the conversation. + +PR #8217 introduced the path-based `display_image` tool and established the +project's terminal image infrastructure: `TerminalImage`, +`terminal-image-renderer`, native Kitty/Ghostty placement, and `chafa` +symbol output. This change extends that infrastructure to in-memory model and +tool image parts instead of adding another renderer. + +## Scope + +This is the render-and-forget slice requested by issue #8090: + +- preserve ordered text and image parts on content events without changing the + existing concatenated `value` contract; +- render live assistant PNGs and restored successful tool PNGs through the + #8217 component and renderer; +- render PNGs nested in successful tool responses; +- keep text/image ordering across retry, model fallback, cancellation, stream + boundaries, and goal-state events; +- bound retained image payloads during UI history compaction; +- render at most four images per assistant output or tool row and collapse the + remainder into a `[+K more images]` marker; +- show a deterministic text placeholder when an image cannot be rendered. + +Kitty deletion, resize-driven replacement, terminal cell pixel queries, and +global scroll lifecycle ownership remain out of scope and are tracked in +#8520. + +## Data Flow + +### Model output + +`ServerGeminiContentEvent.value` remains the concatenated text consumed by +existing clients. When a response chunk contains image `inlineData`, the +event also carries an optional ordered `parts` field containing displayable +non-thought text and image parts. + +Only the interactive TUI reads `parts`. It stages text and image history +items in their original order. A fresh retry or model fallback discards the +failed attempt's staged output, while a normal response boundary commits it. +The TUI admits at most four images for one assistant output and represents any +remaining image parts with a small overflow marker. Visible status rows end the +current assistant display block, so later text starts with the normal assistant +prefix instead of being attached to the status row as a continuation. +Text-only events keep their existing runtime shape, so non-interactive output, +SDK, ACP, daemon, channel, Web UI, and VS Code consumers continue using +`value` unchanged. + +After a thrown stream, staged output remains transient so an explicit retry can +discard the failed attempt. If an out-of-band shell or slash-command item is +added before the next model submit, that item can enter committed history before +the staged output. The next model submit commits the staged output; history +clear, resume, branch, restore, rewind, and Ctrl+L paths discard it instead. + +Resume logic reconstructs ordered text/image runs when persisted parts contain +them. Tool responses retain nested image parts in the session record. The +current Core recorder flattens assistant output to text, so live assistant +images are not restored by `--continue`; assistant-image persistence is outside +this slice and tracked in #8521. + +### Tool output + +Tool media is stored in `functionResponse.parts`. A CLI extractor reads image +`inlineData` from top-level and nested response parts. Live scheduler mapping +and resume mapping attach the images to the existing +`IndividualToolCallDisplay`. Each successful tool row keeps the first four +images and an overflow count for the rest. Failed and cancelled tool records +currently do not carry inline image parts from Core, so their image handling is +defensive rather than a supported output path in this slice. + +Tools carrying images render individually even when their text-only form would +normally collapse into a read/search summary. `ToolMessage` routes the images +through the same `TerminalImage` component used by assistant messages. + +## Rendering + +The existing #8217 file-path entry point is unchanged. The shared renderer +adds an in-memory PNG entry point that: + +1. validates bounded base64 before decoding; +2. verifies the PNG signature and IHDR dimensions; +3. rejects payloads above 8 MiB, dimensions above 1,000,000 pixels, or images + above 64 million total pixels; +4. reuses the existing terminal sizing and bounded render cache; +5. uses native Kitty placement in direct Kitty/Ghostty sessions; +6. passes PNG bytes to `chafa` over stdin in other supported environments; +7. returns a text placeholder when rendering is unavailable. + +No temporary file is created. The inline payload is never used as a command +argument, and `chafa` receives the same allowlisted environment as the +path-based renderer. + +The fallback format is `[image: x png]`. Invalid PNG data +becomes `[image: png]`; unsupported image MIME types retain their sanitized +format label, such as `[image: jpeg]`. Screen-reader mode always uses the text +placeholder and emits no raw image sequence. + +The same encoded-length limit is applied before inline data enters CLI history +or tool-display state. Payloads that exceed the renderer's 8 MiB decoded-image +budget are dropped before rendering and do not produce a placeholder. + +The first slice renders validated PNG data only. Other image MIME types remain +visible as deterministic placeholders rather than entering a second protocol +or decoding path. + +## Memory + +Encoded images are much larger than ordinary history text. UI compaction drops +payloads from old assistant image items while retaining the 20 most recent +items. Cleared images leave a visible marker instead of becoming blank rows. +Tool image payloads participate in the existing tool-result compaction limit. +The four-image admission cap also bounds synchronous Kitty/chafa rendering for +each assistant output and tool row. + +## Test Plan + +- Keep every #8217 renderer and `display_image` test green. +- Verify inline PNG validation, Kitty rendering, `chafa` stdin rendering, + screen-reader output, and unavailable-renderer placeholders. +- Verify `Turn` preserves mixed `text -> image -> text` ordering while + retaining the old `value` and text-only event shape. +- Verify live TUI ordering across retry, fallback, cancellation, stream + boundaries, and goal-state events. +- Verify live and restored successful tool responses expose nested images. +- Verify the resume parser preserves assistant text/image ordering for records + that already contain persisted parts; the current recorder's assistant-image + persistence gap remains outside this slice. +- Verify live assistant and successful tool output enforce the image cap and + expose the overflow count; oversized payloads are dropped before UI history. +- Verify memory compaction clears old assistant and tool image payloads. diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx index 6a66052271a..0d9988e84e9 100644 --- a/packages/cli/src/ui/AppContainer.test.tsx +++ b/packages/cli/src/ui/AppContainer.test.tsx @@ -230,6 +230,7 @@ describe('AppContainer State Management', () => { const mockedUseKeypress = useKeypress as Mock; let originalStdoutIsTTY: boolean | undefined; let restoreCiEnv = () => {}; + let mockClearPendingState: Mock; const mockedRestorePromptStash = vi.mocked(restorePromptStash); beforeEach(() => { @@ -249,6 +250,7 @@ describe('AppContainer State Management', () => { capturedUIActions = null!; capturedRenderMode = 'render'; capturedThoughtExpanded = null!; + mockClearPendingState = vi.fn(); // **Provide a default return value for EVERY mocked hook.** mockedUseHistory.mockReturnValue({ @@ -338,6 +340,7 @@ describe('AppContainer State Management', () => { retryLastPrompt: vi.fn(), streamingResponseLengthRef: { current: 0 }, isReceivingContent: false, + clearPendingState: mockClearPendingState, }); mockedUseVim.mockReturnValue({ handleInput: vi.fn() }); mockedUseFolderTrust.mockReturnValue({ @@ -1259,6 +1262,7 @@ describe('AppContainer State Management', () => { capturedUIActions.handleClearScreen(); expect(clearSpy).toHaveBeenCalledTimes(1); + expect(mockClearPendingState).toHaveBeenCalledTimes(1); expect(mockStdout.write).not.toHaveBeenCalledWith( ansiEscapes.clearTerminal, ); @@ -5774,6 +5778,10 @@ describe('AppContainer State Management', () => { rewindUserItem(1, 'first prompt', 'prompt-1'), { id: 2, type: 'gemini', text: 'first response' }, ]); + expect(mockClearPendingState).toHaveBeenCalledTimes(1); + expect(mockClearPendingState.mock.invocationCallOrder[0]).toBeLessThan( + harness.loadHistory.mock.invocationCallOrder[0]!, + ); expect(harness.setText).toHaveBeenCalledWith('second prompt'); expect(harness.addItem).toHaveBeenCalledWith( expect.objectContaining({ diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index a53f248d43b..de5964922cd 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -1497,6 +1497,12 @@ export const AppContainer = (props: AppContainerProps) => { }); }, [addHistoryItem, config]); + const clearPendingStateRef = useRef<() => void>(() => {}); + const clearPendingStateFromRef = useCallback( + () => clearPendingStateRef.current(), + [], + ); + const { isResumeDialogOpen, resumeMatchedSessions, @@ -1508,6 +1514,7 @@ export const AppContainer = (props: AppContainerProps) => { settings, historyManager, startNewSession, + clearPendingState: clearPendingStateFromRef, setSessionName, remount: refreshStatic, }); @@ -1517,6 +1524,7 @@ export const AppContainer = (props: AppContainerProps) => { settings, historyManager, startNewSession, + clearPendingState: clearPendingStateFromRef, setSessionName, remount: refreshStatic, }); @@ -1735,6 +1743,7 @@ export const AppContainer = (props: AppContainerProps) => { handleBranch, openDeleteDialog, openHelpDialog, + clearPendingState: () => clearPendingStateRef.current(), }), [ openAuthDialog, @@ -2044,6 +2053,7 @@ export const AppContainer = (props: AppContainerProps) => { submitQuery, initError, pendingHistoryItems: pendingGeminiHistoryItems, + clearPendingState, thought, cancelOngoingRequest, preemptGoalTurn, @@ -2082,6 +2092,7 @@ export const AppContainer = (props: AppContainerProps) => { goalQueueRef, ); cancelOngoingRequestRef.current = cancelOngoingRequest; + clearPendingStateRef.current = clearPendingState; // Now that streamingState is available, keep isIdleRef in sync and // flush any deferred update notifications when the model finishes responding. @@ -2930,6 +2941,7 @@ export const AppContainer = (props: AppContainerProps) => { ); const handleClearScreen = useCallback(() => { + clearPendingStateRef.current(); historyManager.clearItems(); clearScreen(); remountStaticHistory(); @@ -3577,6 +3589,7 @@ export const AppContainer = (props: AppContainerProps) => { const truncatedUi = expandCollapsedHistory( originalHistory.filter((h) => h.id < userItem.id), ); + clearPendingStateRef.current(); historyManager.loadHistory(truncatedUi); refreshStatic(); diff --git a/packages/cli/src/ui/commands/restoreCommand.test.ts b/packages/cli/src/ui/commands/restoreCommand.test.ts index 3be2274748e..aa10b631783 100644 --- a/packages/cli/src/ui/commands/restoreCommand.test.ts +++ b/packages/cli/src/ui/commands/restoreCommand.test.ts @@ -54,6 +54,7 @@ describe('restoreCommand', () => { config: mockConfig, }, }); + mockContext.ui.clearPendingState = vi.fn(); }); afterEach(async () => { @@ -170,6 +171,7 @@ describe('restoreCommand', () => { expect(mockContext.ui.loadHistory).toHaveBeenCalledWith( toolCallData.history, ); + expect(mockContext.ui.clearPendingState).toHaveBeenCalledTimes(1); expect(mockSetHistory).toHaveBeenCalledWith(toolCallData.clientHistory); expect(mockRewind).toHaveBeenCalledWith(toolCallData.promptId, true); expect(mockContext.ui.addItem).toHaveBeenCalledWith( diff --git a/packages/cli/src/ui/commands/restoreCommand.ts b/packages/cli/src/ui/commands/restoreCommand.ts index 8c335c8dec9..8fc3c36f1a6 100644 --- a/packages/cli/src/ui/commands/restoreCommand.ts +++ b/packages/cli/src/ui/commands/restoreCommand.ts @@ -144,6 +144,7 @@ async function restoreAction( content: 'loadHistory function is not available.', }; } + context.ui.clearPendingState?.(); loadHistory(toolCallData.history); } diff --git a/packages/cli/src/ui/commands/types.ts b/packages/cli/src/ui/commands/types.ts index d82680e7bd4..f2444722e31 100644 --- a/packages/cli/src/ui/commands/types.ts +++ b/packages/cli/src/ui/commands/types.ts @@ -63,6 +63,8 @@ export interface CommandContext { addItem: UseHistoryManagerReturn['addItem']; /** Clears all history items and the console screen. */ clear: () => void; + /** Clears transient assistant output before replacing conversation history. */ + clearPendingState?: () => void; /** * Sets the transient debug message displayed in the application footer in debug mode. */ diff --git a/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx b/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx index 9e7ff78b5b3..fce85d88a77 100644 --- a/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx +++ b/packages/cli/src/ui/components/HistoryItemDisplay.test.tsx @@ -5,6 +5,7 @@ */ import { describe, it, expect, vi } from 'vitest'; +import { Text } from 'ink'; import { HistoryItemDisplay } from './HistoryItemDisplay.js'; import { type HistoryItem, ToolCallStatus } from '../types.js'; import { MessageType } from '../types.js'; @@ -30,6 +31,12 @@ vi.mock('./messages/ToolGroupMessage.js', () => ({ ToolGroupMessage: vi.fn(() =>
), })); +vi.mock('./TerminalImage.js', () => ({ + TerminalImage: ({ image }: { image: { mimeType: string } }) => ( + MockTerminalImage:{image.mimeType} + ), +})); + vi.mock('../hooks/useMouseEvents.js', () => ({ useMouseEvents: vi.fn(), })); @@ -94,6 +101,29 @@ describe('', () => { expect(output).toContain('◆\uFE0E Hello'); }); + it.each(['gemini', 'gemini_content'] as const)( + 'passes images from a %s history item to the assistant renderer', + (type) => { + const item: HistoryItem = { + id: 1, + type, + text: '', + images: [{ data: 'aW1hZ2U=', mimeType: 'image/png' }], + omittedImageCount: 2, + }; + const { lastFrame } = renderWithProviders( + , + ); + + expect(lastFrame()).toContain('MockTerminalImage:image/png'); + expect(lastFrame()).toContain('[+2 more images]'); + }, + ); + it('renders tool summaries without a leading spacer row', () => { const item: HistoryItem = { id: 1, diff --git a/packages/cli/src/ui/components/HistoryItemDisplay.tsx b/packages/cli/src/ui/components/HistoryItemDisplay.tsx index bb6074cab44..7ff5a9e54da 100644 --- a/packages/cli/src/ui/components/HistoryItemDisplay.tsx +++ b/packages/cli/src/ui/components/HistoryItemDisplay.tsx @@ -300,6 +300,8 @@ const HistoryItemDisplayComponent: React.FC = ({ )} = ({ {itemForDisplay.type === 'gemini_content' && ( ({ writtenKeys: new Set() })); vi.mock('../utils/terminal-image-renderer.js', () => ({ + prepareInlineTerminalImage: vi.fn(), renderTerminalImage: vi.fn(), wasKittyImageWritten: vi.fn((key: string) => writtenKeys.has(key)), markKittyImageWritten: vi.fn((key: string) => { @@ -25,7 +27,16 @@ vi.mock('../utils/terminal-image-renderer.js', () => ({ }), })); +vi.mock('ink', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + useIsScreenReaderEnabled: vi.fn(() => false), + }; +}); + const mockedRenderTerminalImage = vi.mocked(renderTerminalImage); +const mockedPrepareInlineTerminalImage = vi.mocked(prepareInlineTerminalImage); function configWithWorkspaceResult(isWithinWorkspace: boolean): Config { return { @@ -41,6 +52,11 @@ const IMAGE = { mimeType: 'image/png' as const, }; +const INLINE_IMAGE = { + data: 'iVBORw0KGgo=', + mimeType: 'image/png', +}; + const KITTY_RESULT: TerminalImageRenderResult = { kind: 'kitty', key: 'kitty-payload', @@ -77,6 +93,7 @@ describe('TerminalImage', () => { beforeEach(() => { writtenKeys.clear(); vi.clearAllMocks(); + vi.mocked(useIsScreenReaderEnabled).mockReturnValue(false); }); it('writes trusted Kitty data and renders its placeholder', async () => { @@ -183,4 +200,64 @@ describe('TerminalImage', () => { expect(secondWriteRaw).not.toHaveBeenCalled(); second.unmount(); }); + + it('renders inline image data through the shared renderer', async () => { + const writeRaw = vi.fn(); + mockedPrepareInlineTerminalImage.mockReturnValue({ + fallbackText: '[image: 1x1 png]', + result: KITTY_RESULT, + }); + + const { lastFrame } = render( + + + , + ); + + await vi.waitFor(() => { + expect(writeRaw).toHaveBeenCalledWith(KITTY_RESULT.sequence); + }); + expect(lastFrame()).toContain('placeholder'); + expect(mockedPrepareInlineTerminalImage).toHaveBeenCalledWith( + expect.objectContaining({ + data: INLINE_IMAGE.data, + mimeType: INLINE_IMAGE.mimeType, + disabled: false, + }), + ); + }); + + it('explains why an inline image renderer is unavailable', () => { + mockedPrepareInlineTerminalImage.mockReturnValue({ + fallbackText: '[image: 1x1 png]', + result: { + kind: 'unavailable', + reason: 'chafa is not installed', + }, + }); + + const { lastFrame } = render( + , + ); + + expect(lastFrame()).toContain('[image: 1x1 png]'); + expect(lastFrame()).toContain('chafa is not installed'); + }); + + it('uses the deterministic inline placeholder for screen readers', () => { + vi.mocked(useIsScreenReaderEnabled).mockReturnValue(true); + mockedPrepareInlineTerminalImage.mockReturnValue({ + fallbackText: '[image: 1x1 png]', + result: null, + }); + + const { lastFrame } = render( + , + ); + + expect(lastFrame()).toContain('[image: 1x1 png]'); + expect(mockedPrepareInlineTerminalImage).toHaveBeenCalledWith( + expect.objectContaining({ disabled: true }), + ); + }); }); diff --git a/packages/cli/src/ui/components/TerminalImage.tsx b/packages/cli/src/ui/components/TerminalImage.tsx index 36169dfcfcc..0bc33650f8d 100644 --- a/packages/cli/src/ui/components/TerminalImage.tsx +++ b/packages/cli/src/ui/components/TerminalImage.tsx @@ -6,12 +6,14 @@ import path from 'node:path'; import React from 'react'; -import { Box, Text } from 'ink'; +import { Box, Text, useIsScreenReaderEnabled } from 'ink'; import type { Config, TerminalImageDisplay } from '@qwen-code/qwen-code-core'; +import type { InlineImageData } from '../types.js'; import { MaxSizedBox } from './shared/MaxSizedBox.js'; import { useTerminalOutput } from '../contexts/TerminalOutputContext.js'; import { markKittyImageWritten, + prepareInlineTerminalImage, renderTerminalImage, wasKittyImageWritten, type TerminalImageRenderResult, @@ -22,62 +24,42 @@ import { sanitizeTerminalText, } from '../utils/textUtils.js'; -interface TerminalImageProps { - data: TerminalImageDisplay; - config: Config; +interface SharedTerminalImageProps { contentWidth: number; availableTerminalHeight?: number; } -export const TerminalImage: React.FC = ({ - data, - config, - contentWidth, - availableTerminalHeight, -}) => { +interface FileTerminalImageProps extends SharedTerminalImageProps { + data: TerminalImageDisplay; + config: Config; +} + +interface InlineTerminalImageProps extends SharedTerminalImageProps { + image: InlineImageData; +} + +type TerminalImageProps = FileTerminalImageProps | InlineTerminalImageProps; + +const RenderedTerminalImage: React.FC< + SharedTerminalImageProps & { + result: TerminalImageRenderResult; + unavailableText: string; + } +> = ({ result, unavailableText, contentWidth, availableTerminalHeight }) => { const writeRaw = useTerminalOutput(); - const filePath = path.resolve(data.filePath); - const safePath = config.getWorkspaceContext().isPathWithinWorkspace(filePath); - const result = React.useMemo( - () => - safePath - ? renderTerminalImage({ - display: { - type: 'terminal_image', - filePath, - mimeType: data.mimeType, - }, - contentWidth, - availableTerminalHeight, - }) - : null, - [availableTerminalHeight, contentWidth, data.mimeType, filePath, safePath], - ); - // The Kitty payload is written once per terminal session per render key; the - // terminal keeps the image and redraws it from the placeholder cells, so a - // remount (live row -> Static row, or a resize) must not re-transmit it. React.useEffect(() => { - if (!result || result.kind !== 'kitty') return; + if (result.kind !== 'kitty') return; if (wasKittyImageWritten(result.key)) return; markKittyImageWritten(result.key); const sequence = result.sequence; process.nextTick(() => writeRaw(sequence)); }, [result, writeRaw]); - if (!safePath) { - return ( - - Refusing to display an image outside the current workspace. - - ); - } - if (!result) return null; if (result.kind === 'unavailable') { - const fileName = sanitizeMultilineForDisplay(path.basename(filePath)); return ( - {fileName}: {sanitizeTerminalText(result.reason)} + {unavailableText} ); } @@ -112,3 +94,98 @@ export const TerminalImage: React.FC = ({ ); }; + +const FileTerminalImage: React.FC = ({ + data, + config, + contentWidth, + availableTerminalHeight, +}) => { + const filePath = path.resolve(data.filePath); + const safePath = config.getWorkspaceContext().isPathWithinWorkspace(filePath); + const result = React.useMemo( + () => + safePath + ? renderTerminalImage({ + display: { + type: 'terminal_image', + filePath, + mimeType: data.mimeType, + }, + contentWidth, + availableTerminalHeight, + }) + : null, + [availableTerminalHeight, contentWidth, data.mimeType, filePath, safePath], + ); + + if (!safePath) { + return ( + + Refusing to display an image outside the current workspace. + + ); + } + if (!result) return null; + const unavailableText = + result.kind === 'unavailable' + ? `${sanitizeMultilineForDisplay(path.basename(filePath))}: ${sanitizeTerminalText(result.reason)}` + : ''; + + return ( + + ); +}; + +const InlineTerminalImage: React.FC = ({ + image, + contentWidth, + availableTerminalHeight, +}) => { + const isScreenReaderEnabled = useIsScreenReaderEnabled(); + const prepared = React.useMemo( + () => + prepareInlineTerminalImage({ + data: image.data, + mimeType: image.mimeType, + contentWidth, + availableTerminalHeight, + disabled: isScreenReaderEnabled, + }), + [ + availableTerminalHeight, + contentWidth, + image.data, + image.mimeType, + isScreenReaderEnabled, + ], + ); + + if (!prepared.result) { + return {prepared.fallbackText}; + } + return ( + + ); +}; + +export const TerminalImage: React.FC = (props) => + 'image' in props ? ( + + ) : ( + + ); diff --git a/packages/cli/src/ui/components/messages/ConversationMessages.test.tsx b/packages/cli/src/ui/components/messages/ConversationMessages.test.tsx index 9020e8237d1..8f24edb90ea 100644 --- a/packages/cli/src/ui/components/messages/ConversationMessages.test.tsx +++ b/packages/cli/src/ui/components/messages/ConversationMessages.test.tsx @@ -5,12 +5,91 @@ */ import { render } from 'ink-testing-library'; +import { Text } from 'ink'; +import { vi } from 'vitest'; import { + AssistantMessage, + AssistantMessageContent, ThinkMessage, ThinkMessageContent, toggleKeyHint, } from './ConversationMessages.js'; +vi.mock('../TerminalImage.js', () => ({ + TerminalImage: ({ + image, + availableTerminalHeight, + }: { + image: { mimeType: string }; + availableTerminalHeight?: number; + }) => ( + + MockTerminalImage:{image.mimeType}:height= + {availableTerminalHeight ?? 'undef'} + + ), +})); + +describe('', () => { + it('routes assistant images through TerminalImage', () => { + const { lastFrame } = render( + , + ); + + expect(lastFrame()).toContain('MockTerminalImage:image/png'); + }); + + it('renders the number of omitted images', () => { + const { lastFrame } = render( + , + ); + + expect(lastFrame()).toContain('[+2 more images]'); + }); + + it('shares the pending height budget across assistant images', () => { + const { lastFrame } = render( + , + ); + + expect(lastFrame()).toContain('MockTerminalImage:image/png:height=6'); + }); +}); + +describe('', () => { + it('routes continuation images through TerminalImage', () => { + const { lastFrame } = render( + , + ); + + expect(lastFrame()).toContain('MockTerminalImage:image/png'); + }); +}); + describe('', () => { const defaultProps = { text: 'Analyzing the code structure', diff --git a/packages/cli/src/ui/components/messages/ConversationMessages.tsx b/packages/cli/src/ui/components/messages/ConversationMessages.tsx index 949a3e4d34b..fa81f473d24 100644 --- a/packages/cli/src/ui/components/messages/ConversationMessages.tsx +++ b/packages/cli/src/ui/components/messages/ConversationMessages.tsx @@ -22,6 +22,9 @@ import { ErrorBoundary } from '../shared/ErrorBoundary.js'; import { ICON } from '../../constants.js'; import { sanitizeTerminalText } from '../../utils/textUtils.js'; import { formatDuration } from '../../utils/displayUtils.js'; +import type { InlineImageData } from '../../types.js'; +import { TerminalImage } from '../TerminalImage.js'; +import { formatInlineImageOverflow } from '../../utils/inline-image-parts.js'; const debugLogger = createDebugLogger('THINK_RENDER'); @@ -40,6 +43,8 @@ interface UserShellMessageProps { interface AssistantMessageProps { text: string; + images?: InlineImageData[]; + omittedImageCount?: number; isPending: boolean; availableTerminalHeight?: number; contentWidth: number; @@ -48,6 +53,8 @@ interface AssistantMessageProps { interface AssistantMessageContentProps { text: string; + images?: InlineImageData[]; + omittedImageCount?: number; isPending: boolean; availableTerminalHeight?: number; contentWidth: number; @@ -89,6 +96,8 @@ interface PrefixedTextMessageProps { interface PrefixedMarkdownMessageProps { text: string; + images?: InlineImageData[]; + omittedImageCount?: number; prefix: string; prefixColor: string; isPending: boolean; @@ -101,6 +110,8 @@ interface PrefixedMarkdownMessageProps { interface ContinuationMarkdownMessageProps { text: string; + images?: InlineImageData[]; + omittedImageCount?: number; isPending: boolean; availableTerminalHeight?: number; contentWidth: number; @@ -148,6 +159,8 @@ const PrefixedTextMessage: React.FC = ({ const PrefixedMarkdownMessage: React.FC = ({ text, + images, + omittedImageCount, prefix, prefixColor, isPending, @@ -158,6 +171,10 @@ const PrefixedMarkdownMessage: React.FC = ({ sourceCopyIndexOffsets, }) => { const prefixWidth = getPrefixWidth(prefix); + const imageHeightBudget = + availableTerminalHeight !== undefined && images?.length + ? Math.max(1, Math.floor(availableTerminalHeight / (images.length + 1))) + : availableTerminalHeight; return ( @@ -167,14 +184,27 @@ const PrefixedMarkdownMessage: React.FC = ({ - + {text.length > 0 && ( + + )} + {images?.map((image, index) => ( + + ))} + {omittedImageCount !== undefined && omittedImageCount > 0 && ( + {formatInlineImageOverflow(omittedImageCount)} + )} ); @@ -184,6 +214,8 @@ const ContinuationMarkdownMessage: React.FC< ContinuationMarkdownMessageProps > = ({ text, + images, + omittedImageCount, isPending, availableTerminalHeight, contentWidth, @@ -192,17 +224,34 @@ const ContinuationMarkdownMessage: React.FC< sourceCopyIndexOffsets, }) => { const prefixWidth = getPrefixWidth(basePrefix); + const imageHeightBudget = + availableTerminalHeight !== undefined && images?.length + ? Math.max(1, Math.floor(availableTerminalHeight / (images.length + 1))) + : availableTerminalHeight; return ( - + {text.length > 0 && ( + + )} + {images?.map((image, index) => ( + + ))} + {omittedImageCount !== undefined && omittedImageCount > 0 && ( + {formatInlineImageOverflow(omittedImageCount)} + )} ); }; @@ -236,6 +285,8 @@ export const UserShellMessage: React.FC = ({ text }) => { export const AssistantMessage: React.FC = ({ text, + images, + omittedImageCount, isPending, availableTerminalHeight, contentWidth, @@ -243,6 +294,8 @@ export const AssistantMessage: React.FC = ({ }) => ( = ({ text, + images, + omittedImageCount, isPending, availableTerminalHeight, contentWidth, @@ -264,6 +319,8 @@ export const AssistantMessageContent: React.FC< }) => ( ', () => { expect(frame).not.toContain('MockTool'); }); + it('renders image-bearing collapsible tools individually', () => { + const toolCalls = [ + createToolCall({ + callId: 'image-read', + name: 'ReadFile', + description: 'chart.png', + images: [{ data: 'aW1hZ2U=', mimeType: 'image/png' }], + }), + ]; + const { lastFrame } = renderWithProviders( + , + ); + + expect(lastFrame()).toContain('MockTool[image-read]'); + }); + + it('renders an overflow-only collapsible tool individually', () => { + const toolCalls = [ + createToolCall({ + callId: 'overflow-read', + name: 'ReadFile', + description: 'many charts', + omittedImageCount: 2, + }), + ]; + const { lastFrame } = renderWithProviders( + , + ); + + expect(lastFrame()).toContain('MockTool[overflow-read]'); + }); + it('renders mixed group with summary + individual tools', () => { const toolCalls = [ createToolCall({ callId: 'r1', name: 'ReadFile', description: 'a.ts' }), diff --git a/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx b/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx index 09488a37622..f3f4fceb26c 100644 --- a/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx +++ b/packages/cli/src/ui/components/messages/ToolGroupMessage.tsx @@ -45,6 +45,10 @@ function isRunningAgent( ); } +function hasInlineImageOutput(tool: IndividualToolCallDisplay): boolean { + return Boolean(tool.images?.length || tool.omittedImageCount); +} + /** * Predicate: tool entry whose `resultDisplay` is an `AgentResultDisplay` * (i.e. a `task_execution` subagent invocation), regardless of status. @@ -388,13 +392,17 @@ export const ToolGroupMessage: React.FC = ({ ? [] : inlineToolCalls.filter( (t) => - isCollapsibleTool(t.name) && t.status !== ToolCallStatus.Canceled, + isCollapsibleTool(t.name) && + t.status !== ToolCallStatus.Canceled && + !hasInlineImageOutput(t), ); const nonCollapsibleTools = forceExpandAll ? inlineToolCalls : inlineToolCalls.filter( (t) => - !isCollapsibleTool(t.name) || t.status === ToolCallStatus.Canceled, + !isCollapsibleTool(t.name) || + t.status === ToolCallStatus.Canceled || + hasInlineImageOutput(t), ); // Memory badge — shared between all-collapsible and mixed paths. @@ -448,7 +456,10 @@ export const ToolGroupMessage: React.FC = ({ let countToolCallsWithResults = 0; for (const tool of nonCollapsibleTools) { - if (tool.resultDisplay !== undefined && tool.resultDisplay !== '') { + if ( + (tool.resultDisplay !== undefined && tool.resultDisplay !== '') || + hasInlineImageOutput(tool) + ) { countToolCallsWithResults++; } } diff --git a/packages/cli/src/ui/components/messages/ToolMessage.test.tsx b/packages/cli/src/ui/components/messages/ToolMessage.test.tsx index 7b13aa36adf..3c5bbac5a73 100644 --- a/packages/cli/src/ui/components/messages/ToolMessage.test.tsx +++ b/packages/cli/src/ui/components/messages/ToolMessage.test.tsx @@ -73,11 +73,17 @@ vi.mock('../AnsiOutput.js', () => ({ vi.mock('../TerminalImage.js', () => ({ TerminalImage: ({ data, + image, + availableTerminalHeight, }: { - data: { filePath: string; mimeType: string }; + data?: { filePath: string; mimeType: string }; + image?: { mimeType: string }; + availableTerminalHeight?: number; }) => ( - MockTerminalImage:{data.filePath}:{data.mimeType} + {image + ? `MockTerminalImage:${image.mimeType}:height=${availableTerminalHeight ?? 'undef'}` + : `MockTerminalImage:${data?.filePath}:${data?.mimeType}:height=${availableTerminalHeight ?? 'undef'}`} ), })); @@ -187,6 +193,43 @@ describe('', () => { expect(output).not.toContain('MockMarkdown:Test result'); // collapsed }); + it('renders inline images returned by a tool', () => { + const { lastFrame } = renderWithContext( + , + StreamingState.Idle, + ); + + expect(lastFrame()).toContain('MockTerminalImage:image/png'); + }); + + it('renders the number of omitted inline images', () => { + const { lastFrame } = renderWithContext( + , + StreamingState.Idle, + ); + + expect(lastFrame()).toContain('[+2 more images]'); + }); + + it('shares the tool height budget across inline images', () => { + const { lastFrame } = renderWithContext( + , + StreamingState.Responding, + ); + + expect(lastFrame()).toContain('MockTerminalImage:image/png:height=4'); + }); + it('always shows the vision bridge disclosure for a completed read', () => { const { lastFrame } = renderWithContext( = ({ name, description, resultDisplay, + images, + omittedImageCount, visionBridgeNotice, detailedDisplay, status, @@ -775,6 +778,10 @@ export const ToolMessage: React.FC = ({ MIN_LINES_SHOWN + 1, // enforce minimum lines shown ) : undefined; + const inlineImageHeight = + availableHeight !== undefined && images?.length + ? Math.max(1, Math.floor(availableHeight / (images.length + 1))) + : availableHeight; // Cap inline shell output. Applies to both the streaming ANSI display and // the completed string display (shell.ts emits the final result as a plain // string via `returnDisplayMessage = result.output`). ShellStatsBar surfaces @@ -973,6 +980,26 @@ export const ToolMessage: React.FC = ({ )} + {((images?.length ?? 0) > 0 || + (omittedImageCount !== undefined && omittedImageCount > 0)) && ( + + {images?.map((image, index) => ( + + ))} + {omittedImageCount !== undefined && omittedImageCount > 0 && ( + {formatInlineImageOverflow(omittedImageCount)} + )} + + )} {isThisShellFocused && config && ( { const mockOpenMemoryDialog = vi.fn(); const mockOpenModelDialog = vi.fn(); const mockSetQuittingMessages = vi.fn(); + const mockClearPendingState = vi.fn(); const mockConfig = makeFakeConfig({}); mockConfig.getChatRecordingService = vi.fn().mockReturnValue({ @@ -181,6 +182,7 @@ describe('useSlashCommandProcessor', () => { openMcpDialog: vi.fn(), openHooksDialog: vi.fn(), openRewindSelector: vi.fn(), + clearPendingState: mockClearPendingState, }); beforeEach(() => { @@ -2632,6 +2634,18 @@ describe('useSlashCommandProcessor', () => { }); describe('ui.clear and /btw dialog', () => { + it('should discard pending Gemini state when ui.clear is called', async () => { + const result = setupProcessorHook(); + await waitFor(() => expect(result.current.commandContext).toBeDefined()); + + act(() => { + result.current.commandContext.ui.clear(); + }); + + expect(mockClearPendingState).toHaveBeenCalledTimes(1); + expect(mockClearItems).toHaveBeenCalledTimes(1); + }); + it('should dismiss an active btw dialog when ui.clear is called', async () => { const result = setupProcessorHook(); await waitFor(() => expect(result.current.commandContext).toBeDefined()); diff --git a/packages/cli/src/ui/hooks/slashCommandProcessor.ts b/packages/cli/src/ui/hooks/slashCommandProcessor.ts index 68922cba960..73eb734893f 100644 --- a/packages/cli/src/ui/hooks/slashCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/slashCommandProcessor.ts @@ -163,6 +163,7 @@ export interface SlashCommandProcessorActions { openRewindSelector: () => void; openDiffDialog: () => void; openHelpDialog: () => void; + clearPendingState: () => void; } /** @@ -479,11 +480,13 @@ export const useSlashCommandProcessor = ( addItem, clear: () => { cancelBtw(); + actions.clearPendingState(); clearItems(); clearScreen(); refreshStatic(); setSessionName?.(null); }, + clearPendingState: actions.clearPendingState, loadHistory, refreshStatic, setDebugMessage: actions.setDebugMessage, diff --git a/packages/cli/src/ui/hooks/useBranchCommand.test.ts b/packages/cli/src/ui/hooks/useBranchCommand.test.ts index e53797ceb91..cee5f7d3683 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.test.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.test.ts @@ -23,6 +23,7 @@ describe('useBranchCommand', () => { let startNewSessionConfig: ReturnType; let getGoalRuntimeReady: ReturnType; let startNewSessionUI: ReturnType; + let clearPendingState: ReturnType; let findSessionTitlesByPrefix: ReturnType; let clearItems: ReturnType; let loadHistory: ReturnType; @@ -54,6 +55,7 @@ describe('useBranchCommand', () => { settings: mockSettings, historyManager: { clearItems, loadHistory, addItem }, startNewSession: startNewSessionUI, + clearPendingState, setSessionName, remount, }); @@ -92,6 +94,7 @@ describe('useBranchCommand', () => { startNewSessionConfig = vi.fn(); getGoalRuntimeReady = vi.fn().mockResolvedValue({}); startNewSessionUI = vi.fn(); + clearPendingState = vi.fn(); clearItems = vi.fn(); loadHistory = vi.fn(); setSessionName = vi.fn(); @@ -164,6 +167,10 @@ describe('useBranchCommand', () => { expect(monitorRegistry.reset).toHaveBeenCalledOnce(); expect(backgroundShellRegistry.reset).toHaveBeenCalledOnce(); expect(workflowRunRegistry.reset).toHaveBeenCalledOnce(); + expect(clearPendingState).toHaveBeenCalledOnce(); + expect(clearPendingState.mock.invocationCallOrder[0]).toBeLessThan( + loadHistory.mock.invocationCallOrder[0]!, + ); expect(startNewSessionUI.mock.invocationCallOrder[0]).toBeLessThan( backgroundTaskRegistry.reset.mock.invocationCallOrder[0]!, ); diff --git a/packages/cli/src/ui/hooks/useBranchCommand.ts b/packages/cli/src/ui/hooks/useBranchCommand.ts index 7fd09d608d2..2ac2d6343e7 100644 --- a/packages/cli/src/ui/hooks/useBranchCommand.ts +++ b/packages/cli/src/ui/hooks/useBranchCommand.ts @@ -68,6 +68,7 @@ export interface UseBranchCommandOptions { 'clearItems' | 'loadHistory' | 'addItem' >; startNewSession: (sessionId: string) => void; + clearPendingState?: () => void; setSessionName?: (name: string | null) => void; remount?: () => void; } @@ -93,8 +94,14 @@ export interface UseBranchCommandResult { export function useBranchCommand( options: UseBranchCommandOptions, ): UseBranchCommandResult { - const { config, historyManager, startNewSession, setSessionName, remount } = - options; + const { + config, + historyManager, + startNewSession, + clearPendingState, + setSessionName, + remount, + } = options; const handleBranch = useCallback( async (name?: string) => { @@ -208,6 +215,7 @@ export function useBranchCommand( collapsePreviewCount, ); startNewSession(newSessionId); + clearPendingState?.(); historyManager.clearItems(); historyManager.loadHistory(uiHistoryItems); uiSwapped = true; @@ -294,6 +302,7 @@ export function useBranchCommand( config, historyManager, startNewSession, + clearPendingState, setSessionName, remount, options.settings.merged.ui?.history?.collapseOnResume, diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index c13624f1356..2a30df2d486 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx +++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx @@ -44,6 +44,10 @@ import type { HistoryItem, SlashCommandProcessorResult } from '../types.js'; import { MessageType, StreamingState, ToolCallStatus } from '../types.js'; import type { LoadedSettings } from '../../config/settings.js'; import { findLastSafeSplitPoint } from '../utils/markdownUtilities.js'; +import { + MAX_INLINE_IMAGE_ENCODED_LENGTH, + MAX_INLINE_IMAGES_PER_ITEM, +} from '../utils/inline-image-parts.js'; import type { DirectUserAdmission, QueuedGoalTurn } from './useMessageQueue.js'; // --- MOCKS --- @@ -6097,254 +6101,1221 @@ describe('useGeminiStream', () => { }); }); - it('does not render leading blank content chunks as an empty assistant item', async () => { + it('preserves text and inline image ordering in streamed content', async () => { vi.useFakeTimers(); - let releaseNextChunk!: () => void; - const waitForNextChunk = new Promise((resolve) => { - releaseNextChunk = resolve; - }); let releaseStream!: () => void; const holdStream = new Promise((resolve) => { releaseStream = resolve; }); - vi.mocked(findLastSafeSplitPoint).mockImplementation((s: string) => - s.startsWith('\n\n') ? 2 : s.length, + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: image }, + { text: 'after' }, + ], + }; + await holdStream; + })(), ); - const mockStream = (async function* () { - yield { - type: ServerGeminiEventType.Content, - value: '\n\n', - }; - await waitForNextChunk; - yield { - type: ServerGeminiEventType.Content, - value: '哈哈', - }; - await holdStream; - })(); - mockSendMessageStream.mockReturnValue(mockStream); - const { result } = renderTestHook(); - act(() => { - void result.current.submitQuery('test query'); + void result.current.submitQuery('show a chart'); }); - await act(async () => { await Promise.resolve(); await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(60); }); - await act(async () => { - vi.advanceTimersByTime(60); - }); - - expect(result.current.pendingHistoryItems).toEqual([]); + const committedAssistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(committedAssistantItems).toEqual([]); + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'after' }, + ]); + act(() => result.current.cancelOngoingRequest()); + expect( + mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ), + ).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'after' }, + ]); await act(async () => { - releaseNextChunk(); - await Promise.resolve(); - await Promise.resolve(); + releaseStream(); }); + }); + + it('commits mixed content in order before scheduling a tool call', async () => { + const image = { + data: 'dG9vbC1ib3VuZGFyeQ==', + mimeType: 'image/png', + }; + const toolCall = { + callId: 'tool-after-image', + name: 'read_file', + args: { path: '/tmp/example.ts' }, + isClientInitiated: false, + prompt_id: 'prompt-tool-boundary', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: image }, + { text: 'after' }, + ], + }; + yield { + type: ServerGeminiEventType.ToolCallRequest, + value: toolCall, + }; + })(), + ); + const { result } = renderTestHook(); await act(async () => { - vi.advanceTimersByTime(60); + await result.current.submitQuery('read after showing a chart'); }); - expect(result.current.pendingHistoryItems).toEqual([ - expect.objectContaining({ - type: 'gemini', - text: '哈哈', - }), + const assistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(assistantItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'after' }, ]); - - act(() => { - result.current.cancelOngoingRequest(); - }); - - await act(async () => { - releaseStream(); - }); + expect(mockScheduleToolCalls).toHaveBeenCalledTimes(1); + expect(mockScheduleToolCalls.mock.calls[0][0]).toEqual([toolCall]); }); - it('buffers streamed thoughts until the throttle interval elapses', async () => { + it('does not overwrite an image with whitespace before the next image', async () => { vi.useFakeTimers(); let releaseStream!: () => void; const holdStream = new Promise((resolve) => { releaseStream = resolve; }); - - const mockStream = (async function* () { - yield { - type: ServerGeminiEventType.Thought, - value: { description: 'Think' }, - }; - yield { - type: ServerGeminiEventType.Thought, - value: { description: 'ing' }, - }; - await holdStream; - })(); - mockSendMessageStream.mockReturnValue(mockStream); + const firstImage = { + data: 'Zmlyc3Q=', + mimeType: 'image/png', + }; + const secondImage = { + data: 'c2Vjb25k', + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: '\n', + parts: [ + { inlineData: firstImage }, + { text: '\n' }, + { inlineData: secondImage }, + ], + }; + await holdStream; + })(), + ); const { result } = renderTestHook(); - act(() => { - void result.current.submitQuery('test query'); + void result.current.submitQuery('show two charts'); }); - await act(async () => { await Promise.resolve(); await Promise.resolve(); - // Flush the macrotask yield (setImmediate) added after addItem() await vi.advanceTimersByTimeAsync(0); - }); - - expect(mockSendMessageStream).toHaveBeenCalledTimes(1); - expect(result.current.pendingHistoryItems).toEqual([]); - - await act(async () => { - vi.advanceTimersByTime(60); + await vi.advanceTimersByTimeAsync(60); }); expect(result.current.pendingHistoryItems).toEqual([ expect.objectContaining({ - type: 'gemini_thought', - durationMs: expect.any(Number), + type: 'gemini', + text: '', + images: [firstImage], }), + { type: 'gemini_content', text: '', images: [secondImage] }, ]); - expect(result.current.thought).toEqual({ description: 'Thinking' }); - - act(() => { - result.current.cancelOngoingRequest(); - }); + act(() => result.current.cancelOngoingRequest()); await act(async () => { releaseStream(); }); }); - it('splits oversized streamed thoughts so the pending item stays bounded', async () => { + it('caps inline images for one assistant output and reports the overflow', async () => { vi.useFakeTimers(); - const splitLimit = 16_384; - const tailLength = 123; - const longThought = 'a'.repeat(splitLimit * 2 + tailLength); - vi.mocked(findLastSafeSplitPoint).mockImplementation( - (s: string, max?: number) => - max !== undefined && s.length > max ? max : s.length, - ); - let releaseStream!: () => void; const holdStream = new Promise((resolve) => { releaseStream = resolve; }); - - const mockStream = (async function* () { - yield { - type: ServerGeminiEventType.Thought, - value: { description: longThought }, - }; - await holdStream; - })(); - mockSendMessageStream.mockReturnValue(mockStream); + const images = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM + 2 }, + (_, index) => ({ + data: Buffer.from(`assistant-image-${index}`).toString('base64'), + mimeType: 'image/png', + }), + ); + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: images.map((inlineData) => ({ inlineData })), + }; + await holdStream; + })(), + ); const { result } = renderTestHook(); - act(() => { - void result.current.submitQuery('test query'); + void result.current.submitQuery('show many charts'); }); - await act(async () => { await Promise.resolve(); await Promise.resolve(); - // Flush the macrotask yield (setImmediate) added after addItem() await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(60); }); - await act(async () => { - vi.advanceTimersByTime(60); - }); - - const thoughtItems = mockAddItem.mock.calls - .map(([item]) => item as HistoryItem) - .filter( - (item) => - item.type === 'gemini_thought' || - item.type === 'gemini_thought_content', - ); - expect(thoughtItems).toEqual([ - expect.objectContaining({ - type: 'gemini_thought', - text: 'a'.repeat(splitLimit), - durationMs: expect.any(Number), - }), - expect.objectContaining({ - type: 'gemini_thought_content', - text: 'a'.repeat(splitLimit), - }), - ]); - expect(result.current.pendingHistoryItems).toEqual([ - expect.objectContaining({ - type: 'gemini_thought_content', - text: 'a'.repeat(tailLength), - }), - ]); - expect(result.current.thought?.description).toHaveLength(4_096); - - act(() => { - result.current.cancelOngoingRequest(); + const assistantItems = result.current.pendingHistoryItems.filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(assistantItems.flatMap((item) => item.images ?? [])).toEqual( + images.slice(0, MAX_INLINE_IMAGES_PER_ITEM), + ); + expect(assistantItems).toHaveLength(MAX_INLINE_IMAGES_PER_ITEM + 1); + expect(assistantItems.at(-1)).toMatchObject({ + text: '', + omittedImageCount: 2, }); + act(() => result.current.cancelOngoingRequest()); await act(async () => { releaseStream(); }); }); - it('repairs the fence when an oversized thought is split inside a code block', async () => { - vi.useFakeTimers(); - - const splitLimit = 16_384; - vi.mocked(findLastSafeSplitPoint).mockImplementation( - (s: string, max?: number) => - max !== undefined && s.length > max ? max : s.length, + it('applies the inline image cap across multiple content events', async () => { + const images = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM + 2 }, + (_, index) => ({ + data: Buffer.from(`multi-event-image-${index}`).toString('base64'), + mimeType: 'image/png', + }), + ); + mockSendMessageStream.mockReturnValue( + (async function* () { + for (const inlineData of images) { + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: [{ inlineData }], + }; + } + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), ); - - // A reasoning stream whose fenced code block spans the char-cap boundary. - const codeBody = Array.from( - { length: 2000 }, - (_, i) => `const x${i} = ${i};`, - ).join('\n'); - const longThought = '```ts\n' + codeBody; - expect(longThought.length).toBeGreaterThan(splitLimit); - - let releaseStream!: () => void; - const holdStream = new Promise((resolve) => { - releaseStream = resolve; - }); - const mockStream = (async function* () { - yield { - type: ServerGeminiEventType.Thought, - value: { description: longThought }, - }; - await holdStream; - })(); - mockSendMessageStream.mockReturnValue(mockStream); const { result } = renderTestHook(); - act(() => { - void result.current.submitQuery('test query'); - }); - await act(async () => { - await Promise.resolve(); - await Promise.resolve(); - await vi.advanceTimersByTimeAsync(0); - }); await act(async () => { - vi.advanceTimersByTime(60); + await result.current.submitQuery('show many streamed charts'); + }); + + const assistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(assistantItems.flatMap((item) => item.images ?? [])).toEqual( + images.slice(0, MAX_INLINE_IMAGES_PER_ITEM), + ); + expect(assistantItems.at(-1)).toMatchObject({ + text: '', + omittedImageCount: 2, + }); + }); + + it('resets the inline image cap after a fresh retry', async () => { + const failedImages = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM }, + (_, index) => ({ + data: Buffer.from(`retry-image-${index}`).toString('base64'), + mimeType: 'image/png', + }), + ); + const replacementImage = { + data: Buffer.from('retry-replacement').toString('base64'), + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: failedImages.map((inlineData) => ({ inlineData })), + }; + yield { + type: ServerGeminiEventType.Retry, + isContinuation: false, + }; + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: [{ inlineData: replacementImage }], + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('retry the charts'); + }); + + const assistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(assistantItems.flatMap((item) => item.images ?? [])).toEqual([ + replacementImage, + ]); + expect(assistantItems.some((item) => item.omittedImageCount)).toBe(false); + }); + + it('resets the inline image cap after model fallback', async () => { + const failedImages = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM }, + (_, index) => ({ + data: Buffer.from(`fallback-image-${index}`).toString('base64'), + mimeType: 'image/png', + }), + ); + const replacementImage = { + data: Buffer.from('fallback-replacement').toString('base64'), + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: failedImages.map((inlineData) => ({ inlineData })), + }; + yield { + type: ServerGeminiEventType.ModelFallback, + fromModel: 'primary-model', + toModel: 'fallback-model', + fallbackIndex: 1, + }; + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: [{ inlineData: replacementImage }], + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('fallback the charts'); + }); + + const assistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(assistantItems.flatMap((item) => item.images ?? [])).toEqual([ + replacementImage, + ]); + expect(assistantItems.some((item) => item.omittedImageCount)).toBe(false); + }); + + it('resets the inline image cap at a finished response boundary', async () => { + const firstOutputImages = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM }, + (_, index) => ({ + data: Buffer.from(`first-output-${index}`).toString('base64'), + mimeType: 'image/png', + }), + ); + const nextOutputImage = { + data: Buffer.from('next-output').toString('base64'), + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: firstOutputImages.map((inlineData) => ({ inlineData })), + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: [{ inlineData: nextOutputImage }], + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('show charts across responses'); + }); + + const assistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(assistantItems.flatMap((item) => item.images ?? [])).toEqual([ + ...firstOutputImages, + nextOutputImage, + ]); + expect(assistantItems.some((item) => item.omittedImageCount)).toBe(false); + expect(assistantItems.at(-1)).toMatchObject({ + type: 'gemini', + images: [nextOutputImage], + }); + }); + + it('does not retain an oversized inline image payload in UI state', async () => { + vi.useFakeTimers(); + + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + const oversizedData = 'A'.repeat(MAX_INLINE_IMAGE_ENCODED_LENGTH + 1); + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: '', + parts: [ + { + inlineData: { + data: oversizedData, + mimeType: 'image/png', + }, + }, + ], + }; + await holdStream; + })(), + ); + + const { result } = renderTestHook(); + act(() => { + void result.current.submitQuery('show oversized chart'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(60); + }); + + const assistantItems = result.current.pendingHistoryItems.filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(assistantItems).toEqual([]); + expect(JSON.stringify(result.current.pendingHistoryItems)).not.toContain( + oversizedData, + ); + + act(() => result.current.cancelOngoingRequest()); + await act(async () => { + releaseStream(); + }); + }); + + it('keeps an image committable when the stream pauses after whitespace', async () => { + vi.useFakeTimers(); + + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: '\n', + parts: [{ inlineData: image }, { text: '\n' }], + }; + await holdStream; + })(), + ); + + const { result } = renderTestHook(); + act(() => { + void result.current.submitQuery('show a chart'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(60); + }); + + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ + type: 'gemini', + text: '', + images: [image], + }), + ]); + + act(() => result.current.cancelOngoingRequest()); + expect( + mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ), + ).toEqual([ + expect.objectContaining({ + type: 'gemini', + text: '', + images: [image], + }), + ]); + await act(async () => { + releaseStream(); + }); + }); + + it('discards every staged mixed-content run on a fresh retry', async () => { + vi.useFakeTimers(); + + let emitRetry!: () => void; + const waitForRetry = new Promise((resolve) => { + emitRetry = resolve; + }); + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: image }, + { text: 'after' }, + ], + }; + await waitForRetry; + yield { + type: ServerGeminiEventType.Retry, + isContinuation: false, + }; + yield { + type: ServerGeminiEventType.Content, + value: 'replacement', + }; + await holdStream; + })(), + ); + + const { result } = renderTestHook(); + act(() => { + void result.current.submitQuery('show a chart'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(60); + }); + + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'after' }, + ]); + + await act(async () => { + emitRetry(); + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(60); + }); + + const committedAssistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(committedAssistantItems).toEqual([]); + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'replacement' }), + ]); + + act(() => result.current.cancelOngoingRequest()); + await act(async () => { + releaseStream(); + }); + }); + + it('preserves staged mixed-content runs on a continuation retry', async () => { + vi.useFakeTimers(); + + let emitRetry!: () => void; + const waitForRetry = new Promise((resolve) => { + emitRetry = resolve; + }); + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: image }, + { text: 'after' }, + ], + }; + await waitForRetry; + yield { + type: ServerGeminiEventType.Retry, + isContinuation: true, + }; + yield { + type: ServerGeminiEventType.Content, + value: ' continued', + }; + await holdStream; + })(), + ); + + const { result } = renderTestHook(); + act(() => { + void result.current.submitQuery('show a chart'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(60); + }); + + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'after' }, + ]); + + await act(async () => { + emitRetry(); + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(60); + }); + + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'after continued' }, + ]); + + act(() => result.current.cancelOngoingRequest()); + await act(async () => { + releaseStream(); + }); + }); + + it('discards every staged mixed-content run on model fallback', async () => { + vi.useFakeTimers(); + + let emitFallback!: () => void; + const waitForFallback = new Promise((resolve) => { + emitFallback = resolve; + }); + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: image }, + { text: 'after' }, + ], + }; + await waitForFallback; + yield { + type: ServerGeminiEventType.ModelFallback, + fromModel: 'primary-model', + toModel: 'fallback-model', + fallbackIndex: 1, + }; + await holdStream; + })(), + ); + + const { result } = renderTestHook(); + act(() => { + void result.current.submitQuery('show a chart'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(60); + }); + + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'after' }, + ]); + + await act(async () => { + emitFallback(); + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(60); + }); + + expect( + mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ), + ).toEqual([]); + expect(result.current.pendingHistoryItems).toEqual([]); + expect(mockAddItem).toHaveBeenCalledWith( + { + type: 'notification', + text: 'Model primary-model unavailable, falling back to fallback-model', + }, + expect.any(Number), + ); + + act(() => result.current.cancelOngoingRequest()); + await act(async () => { + releaseStream(); + }); + }); + + it('discards staged mixed content before an explicit retry after a thrown stream', async () => { + const failedImage = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream + .mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: failedImage }, + { text: 'after' }, + ], + }; + throw new Error('stream failed'); + })(), + ) + .mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'replacement', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('show a chart'); + }); + + expect(result.current.pendingHistoryItems).toEqual( + expect.arrayContaining([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [failedImage] }, + { type: 'gemini_content', text: 'after' }, + ]), + ); + + await act(async () => { + await result.current.submitQuery('show a chart', SendMessageType.Retry); + }); + + const committedAssistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(committedAssistantItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'replacement' }), + ]); + expect(result.current.pendingHistoryItems).not.toEqual( + expect.arrayContaining([ + expect.objectContaining({ images: [failedImage] }), + ]), + ); + }); + + it('commits staged mixed content before a new user turn after a thrown stream', async () => { + const failedImage = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream + .mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: failedImage }, + { text: 'after' }, + ], + }; + throw new Error('stream failed'); + })(), + ) + .mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'next answer', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('first question'); + }); + await act(async () => { + await result.current.submitQuery('second question'); + }); + + const relevantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => + item.type === 'user' || + item.type === 'gemini' || + item.type === 'gemini_content' || + item.type === 'error', + ); + expect(relevantItems).toEqual([ + expect.objectContaining({ type: 'user', text: 'first question' }), + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [failedImage] }, + { type: 'gemini_content', text: 'after' }, + expect.objectContaining({ type: 'error' }), + expect.objectContaining({ type: 'user', text: 'second question' }), + expect.objectContaining({ type: 'gemini', text: 'next answer' }), + ]); + }); + + it('does not restore staged mixed content after pending state is cleared', async () => { + const failedImage = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream + .mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: failedImage }, + { text: 'after' }, + ], + }; + throw new Error('stream failed'); + })(), + ) + .mockReturnValueOnce( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'new answer', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('first question'); + }); + expect(result.current.pendingHistoryItems).toEqual( + expect.arrayContaining([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [failedImage] }, + { type: 'gemini_content', text: 'after' }, + expect.objectContaining({ type: 'error' }), + ]), + ); + + act(() => { + result.current.clearPendingState(); + }); + expect(result.current.pendingHistoryItems).toEqual([]); + + await act(async () => { + await result.current.submitQuery('second question'); + }); + + const assistantItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + expect(assistantItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'new answer' }), + ]); + }); + + it('does not render leading blank content chunks as an empty assistant item', async () => { + vi.useFakeTimers(); + + let releaseNextChunk!: () => void; + const waitForNextChunk = new Promise((resolve) => { + releaseNextChunk = resolve; + }); + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + vi.mocked(findLastSafeSplitPoint).mockImplementation((s: string) => + s.startsWith('\n\n') ? 2 : s.length, + ); + + const mockStream = (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: '\n\n', + }; + await waitForNextChunk; + yield { + type: ServerGeminiEventType.Content, + value: '哈哈', + }; + await holdStream; + })(); + mockSendMessageStream.mockReturnValue(mockStream); + + const { result } = renderTestHook(); + + act(() => { + void result.current.submitQuery('test query'); + }); + + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + }); + + await act(async () => { + vi.advanceTimersByTime(60); + }); + + expect(result.current.pendingHistoryItems).toEqual([]); + + await act(async () => { + releaseNextChunk(); + await Promise.resolve(); + await Promise.resolve(); + }); + + await act(async () => { + vi.advanceTimersByTime(60); + }); + + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ + type: 'gemini', + text: '哈哈', + }), + ]); + + act(() => { + result.current.cancelOngoingRequest(); + }); + + await act(async () => { + releaseStream(); + }); + }); + + it('buffers streamed thoughts until the throttle interval elapses', async () => { + vi.useFakeTimers(); + + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + + const mockStream = (async function* () { + yield { + type: ServerGeminiEventType.Thought, + value: { description: 'Think' }, + }; + yield { + type: ServerGeminiEventType.Thought, + value: { description: 'ing' }, + }; + await holdStream; + })(); + mockSendMessageStream.mockReturnValue(mockStream); + + const { result } = renderTestHook(); + + act(() => { + void result.current.submitQuery('test query'); + }); + + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + // Flush the macrotask yield (setImmediate) added after addItem() + await vi.advanceTimersByTimeAsync(0); + }); + + expect(mockSendMessageStream).toHaveBeenCalledTimes(1); + expect(result.current.pendingHistoryItems).toEqual([]); + + await act(async () => { + vi.advanceTimersByTime(60); + }); + + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ + type: 'gemini_thought', + durationMs: expect.any(Number), + }), + ]); + expect(result.current.thought).toEqual({ description: 'Thinking' }); + + act(() => { + result.current.cancelOngoingRequest(); + }); + + await act(async () => { + releaseStream(); + }); + }); + + it('splits oversized streamed thoughts so the pending item stays bounded', async () => { + vi.useFakeTimers(); + + const splitLimit = 16_384; + const tailLength = 123; + const longThought = 'a'.repeat(splitLimit * 2 + tailLength); + vi.mocked(findLastSafeSplitPoint).mockImplementation( + (s: string, max?: number) => + max !== undefined && s.length > max ? max : s.length, + ); + + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + + const mockStream = (async function* () { + yield { + type: ServerGeminiEventType.Thought, + value: { description: longThought }, + }; + await holdStream; + })(); + mockSendMessageStream.mockReturnValue(mockStream); + + const { result } = renderTestHook(); + + act(() => { + void result.current.submitQuery('test query'); + }); + + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + // Flush the macrotask yield (setImmediate) added after addItem() + await vi.advanceTimersByTimeAsync(0); + }); + + await act(async () => { + vi.advanceTimersByTime(60); + }); + + const thoughtItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => + item.type === 'gemini_thought' || + item.type === 'gemini_thought_content', + ); + expect(thoughtItems).toEqual([ + expect.objectContaining({ + type: 'gemini_thought', + text: 'a'.repeat(splitLimit), + durationMs: expect.any(Number), + }), + expect.objectContaining({ + type: 'gemini_thought_content', + text: 'a'.repeat(splitLimit), + }), + ]); + expect(result.current.pendingHistoryItems).toEqual([ + expect.objectContaining({ + type: 'gemini_thought_content', + text: 'a'.repeat(tailLength), + }), + ]); + expect(result.current.thought?.description).toHaveLength(4_096); + + act(() => { + result.current.cancelOngoingRequest(); + }); + + await act(async () => { + releaseStream(); + }); + }); + + it('repairs the fence when an oversized thought is split inside a code block', async () => { + vi.useFakeTimers(); + + const splitLimit = 16_384; + vi.mocked(findLastSafeSplitPoint).mockImplementation( + (s: string, max?: number) => + max !== undefined && s.length > max ? max : s.length, + ); + + // A reasoning stream whose fenced code block spans the char-cap boundary. + const codeBody = Array.from( + { length: 2000 }, + (_, i) => `const x${i} = ${i};`, + ).join('\n'); + const longThought = '```ts\n' + codeBody; + expect(longThought.length).toBeGreaterThan(splitLimit); + + let releaseStream!: () => void; + const holdStream = new Promise((resolve) => { + releaseStream = resolve; + }); + const mockStream = (async function* () { + yield { + type: ServerGeminiEventType.Thought, + value: { description: longThought }, + }; + await holdStream; + })(); + mockSendMessageStream.mockReturnValue(mockStream); + + const { result } = renderTestHook(); + act(() => { + void result.current.submitQuery('test query'); + }); + await act(async () => { + await Promise.resolve(); + await Promise.resolve(); + await vi.advanceTimersByTimeAsync(0); + }); + await act(async () => { + vi.advanceTimersByTime(60); }); const thoughtItems = mockAddItem.mock.calls @@ -9095,14 +10066,200 @@ describe('useGeminiStream', () => { }); }); + describe('Citation event', () => { + it('starts a fresh assistant item after a shown citation', async () => { + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'Hello world', + }; + yield { + type: ServerGeminiEventType.Citation, + value: 'Citation text', + }; + yield { + type: ServerGeminiEventType.Content, + value: ' more', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('test shown citation'); + }); + + const outputItems = mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => + item.type === 'gemini' || + item.type === 'gemini_content' || + item.type === MessageType.INFO, + ); + expect(outputItems).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'Hello world' }), + expect.objectContaining({ + type: MessageType.INFO, + text: 'Citation text', + }), + expect.objectContaining({ type: 'gemini', text: ' more' }), + ]); + }); + + it('preserves streamed text across hidden citation events', async () => { + const settingsWithCitationsHidden = { + ...mockLoadedSettings, + merged: { + ...mockLoadedSettings.merged, + ui: { + ...mockLoadedSettings.merged.ui, + showCitations: false, + }, + }, + } as LoadedSettings; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'Hello world', + }; + yield { + type: ServerGeminiEventType.Citation, + value: 'Citation text', + }; + yield { + type: ServerGeminiEventType.Content, + value: ' more', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderHook(() => + useGeminiStream( + new MockedGeminiClientClass(mockConfig), + [], + mockAddItem, + mockConfig, + true, + settingsWithCitationsHidden, + mockOnDebugMessage, + mockHandleSlashCommand, + false, + () => 'vscode' as EditorType, + () => {}, + () => Promise.resolve(), + false, + () => {}, + () => {}, + () => {}, + () => {}, + 80, + 24, + ), + ); + + await act(async () => { + await result.current.submitQuery('test hidden citation'); + }); + + expect( + mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter((item) => item.type === 'gemini'), + ).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'Hello world more' }), + ]); + expect(mockAddItem).not.toHaveBeenCalledWith( + expect.objectContaining({ type: MessageType.INFO }), + expect.any(Number), + ); + }); + }); + + describe('ChatCompressed event', () => { + it('starts a fresh prefixed text item after the status row', async () => { + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'before compression', + }; + yield { + type: ServerGeminiEventType.ChatCompressed, + value: { + originalTokenCount: 100, + newTokenCount: 50, + }, + }; + yield { + type: ServerGeminiEventType.Content, + value: 'after compression', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('test compression boundary'); + }); + + expect( + mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => + item.type === 'gemini' || + item.type === 'gemini_content' || + item.type === 'info', + ), + ).toEqual([ + expect.objectContaining({ + type: 'gemini', + text: 'before compression', + }), + expect.objectContaining({ + type: 'info', + text: expect.stringContaining('compressed context'), + }), + expect.objectContaining({ + type: 'gemini', + text: 'after compression', + }), + ]); + }); + }); + describe('handleFinishedEvent', () => { - it('should add info message for MAX_TOKENS finish reason', async () => { + it('commits mixed assistant output before a MAX_TOKENS warning', async () => { + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; // Setup mock to return a stream with MAX_TOKENS finish reason mockSendMessageStream.mockReturnValue( (async function* () { yield { type: ServerGeminiEventType.Content, value: 'This is a truncated response...', + parts: [ + { text: 'This is ' }, + { inlineData: image }, + { text: 'a truncated response...' }, + ], }; yield { type: ServerGeminiEventType.Finished, @@ -9140,16 +10297,91 @@ describe('useGeminiStream', () => { await result.current.submitQuery('Generate long text'); }); - // Check that the info message was added - await waitFor(() => { - expect(mockAddItem).toHaveBeenCalledWith( - { - type: 'info', - text: '⚠ Response truncated due to token limits.', + expect( + mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => + item.type === 'gemini' || + item.type === 'gemini_content' || + item.type === 'info', + ), + ).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'This is ' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'a truncated response...' }, + { + type: 'info', + text: '⚠ Response truncated due to token limits.', + }, + ]); + }); + + it.each([ + { + name: 'maximum-turns notice', + event: { type: ServerGeminiEventType.MaxSessionTurns }, + expected: { + type: 'info', + text: expect.stringContaining('maximum number of turns'), + }, + }, + { + name: 'session-token-limit error', + event: { + type: ServerGeminiEventType.SessionTokenLimitExceeded, + value: { + currentTokens: 200, + limit: 100, + message: 'limit reached', }, - expect.any(Number), - ); + }, + expected: { + type: 'error', + text: expect.stringContaining('Session token limit exceeded'), + }, + }, + ])('commits mixed assistant output before a $name', async (testCase) => { + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + mockSendMessageStream.mockReturnValue( + (async function* () { + yield { + type: ServerGeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { inlineData: image }, + { text: 'after' }, + ], + }; + yield testCase.event; + })(), + ); + + const { result } = renderTestHook(); + await act(async () => { + await result.current.submitQuery('test terminal boundary'); }); + + expect( + mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => + item.type === 'gemini' || + item.type === 'gemini_content' || + item.type === 'info' || + item.type === 'error', + ), + ).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'before' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'after' }, + testCase.expected, + ]); }); it('should not add message for STOP finish reason', async () => { @@ -11882,12 +13114,21 @@ describe('useGeminiStream', () => { }); describe('HookSystemMessage Event', () => { - it('commits buffered content before a displayed Goal state', async () => { + it('commits staged inline content and restarts after a displayed Goal state', async () => { + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; mockSendMessageStream.mockReturnValue( (async function* () { yield { type: ServerGeminiEventType.Content, value: 'Final Goal output', + parts: [ + { text: 'Final ' }, + { inlineData: image }, + { text: 'Goal output' }, + ], }; yield { type: ServerGeminiEventType.GoalState, @@ -11908,20 +13149,36 @@ describe('useGeminiStream', () => { }, }, }; + yield { + type: ServerGeminiEventType.Content, + value: ' continued', + }; + yield { + type: ServerGeminiEventType.Finished, + value: { reason: 'STOP', usageMetadata: undefined }, + }; })(), ); const { result } = renderTestHook(); await act(async () => { await result.current.submitQuery('finish the Goal'); }); - const contentIndex = mockAddItem.mock.calls.findIndex( - ([item]) => item.type === 'gemini' && item.text === 'Final Goal output', - ); - const goalIndex = mockAddItem.mock.calls.findIndex( - ([item]) => item.type === 'goal_state' && item.cause === 'complete', - ); - expect(contentIndex).toBeGreaterThanOrEqual(0); - expect(goalIndex).toBeGreaterThan(contentIndex); + expect( + mockAddItem.mock.calls + .map(([item]) => item as HistoryItem) + .filter( + (item) => + item.type === 'gemini' || + item.type === 'gemini_content' || + item.type === 'goal_state', + ), + ).toEqual([ + expect.objectContaining({ type: 'gemini', text: 'Final ' }), + { type: 'gemini_content', text: '', images: [image] }, + { type: 'gemini_content', text: 'Goal output' }, + expect.objectContaining({ type: 'goal_state', cause: 'complete' }), + expect.objectContaining({ type: 'gemini', text: ' continued' }), + ]); }); it('should handle HookSystemMessage event and add stop_hook_system_message history item', async () => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 6245513d2ec..346cb2c7cb8 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -75,6 +75,7 @@ import type { HistoryItemWithoutId, HistoryItemToolGroup, HistoryItemGemini, + InlineImageData, SlashCommandProcessorResult, } from '../types.js'; import { StreamingState, MessageType, ToolCallStatus } from '../types.js'; @@ -124,6 +125,10 @@ import { } from './useMessageQueue.js'; import { classifyApiError } from '../../utils/classify-api-error.js'; import { cleanupReviewWorktreeLeases } from '../../services/review-worktree-lease.js'; +import { + getInlineImageData, + MAX_INLINE_IMAGES_PER_ITEM, +} from '../utils/inline-image-parts.js'; const debugLogger = createDebugLogger('GEMINI_STREAM'); @@ -367,6 +372,7 @@ const LOADING_THOUGHT_DESCRIPTION_MAX_CHARS = 4_096; type BufferedStreamEvent = | { kind: 'content'; value: string } + | { kind: 'image'; value: InlineImageData } | { kind: 'thought'; value: ThoughtSummary }; function showCitations(settings: LoadedSettings): boolean { @@ -693,6 +699,44 @@ export const useGeminiStream = ( const auxiliaryAbortRefsRef = useRef>(new Set()); const [pendingHistoryItem, pendingHistoryItemRef, setPendingHistoryItem] = useStateAndRef(null); + // Mixed assistant output needs multiple live rows to preserve + // text/image/text ordering. Keep completed runs in the dynamic region until + // the response reaches a normal commit boundary so a fresh retry or model + // fallback can still discard the entire failed attempt. + const [ + pendingAssistantItems, + pendingAssistantItemsRef, + setPendingAssistantItems, + ] = useStateAndRef([]); + const commitPendingAssistantItems = useCallback( + (userMessageTimestamp: number) => { + const items = pendingAssistantItemsRef.current; + if (items.length === 0) { + return; + } + for (const item of items) { + commitItem(item, userMessageTimestamp); + } + setPendingAssistantItems([]); + }, + [commitItem, pendingAssistantItemsRef, setPendingAssistantItems], + ); + const commitItemInOrder = useCallback( + (item: HistoryItemWithoutId, userMessageTimestamp: number): number => { + commitPendingAssistantItems(userMessageTimestamp); + return commitItem(item, userMessageTimestamp); + }, + [commitItem, commitPendingAssistantItems], + ); + const stagePendingAssistantItem = useCallback((): boolean => { + const item = pendingHistoryItemRef.current; + if (item?.type !== 'gemini' && item?.type !== 'gemini_content') { + return false; + } + setPendingAssistantItems((items) => [...items, item]); + setPendingHistoryItem(null); + return true; + }, [pendingHistoryItemRef, setPendingAssistantItems, setPendingHistoryItem]); // Streamed model reasoning for the current turn. Rendered (height-limited) // above the answer while thinking, then committed to history as a // collapsible `gemini_thought` block when the answer/tool/turn begins. @@ -709,6 +753,15 @@ export const useGeminiStream = ( pendingRetryCountdownItemRef, setPendingRetryCountdownItem, ] = useStateAndRef(null); + const clearPendingState = useCallback(() => { + setPendingAssistantItems([]); + setPendingHistoryItem(null); + setPendingRetryErrorItem(null); + }, [ + setPendingAssistantItems, + setPendingHistoryItem, + setPendingRetryErrorItem, + ]); const retryCountdownTimerRef = useRef | null>( null, ); @@ -1027,7 +1080,7 @@ export const useGeminiStream = ( logApiCancel(config, cancellationEvent); if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, Date.now()); + commitItemInOrder(pendingHistoryItemRef.current, Date.now()); } addItem( { @@ -1078,7 +1131,7 @@ export const useGeminiStream = ( }, [ streamingState, addItem, - commitItem, + commitItemInOrder, setPendingHistoryItem, onCancelSubmit, pendingHistoryItemRef, @@ -1421,6 +1474,7 @@ export const useGeminiStream = ( eventValue: ContentEvent['value'], currentGeminiMessageBuffer: string, userMessageTimestamp: number, + startAsContinuation = false, ): string => { if (turnCancelledRef.current) { // Prevents additional output after a user initiated cancel. @@ -1435,6 +1489,17 @@ export const useGeminiStream = ( // React history by the time AppContainer's guard runs). turnSawContentEventRef.current = true; let newGeminiMessageBuffer = currentGeminiMessageBuffer + eventValue; + const pendingItem = pendingHistoryItemRef.current; + if ( + (pendingItem?.type === 'gemini' || + pendingItem?.type === 'gemini_content') && + (pendingItem.images?.length || pendingItem.omittedImageCount) + ) { + if (newGeminiMessageBuffer.trim().length === 0) { + return newGeminiMessageBuffer; + } + stagePendingAssistantItem(); + } if ( pendingHistoryItemRef.current?.type !== 'gemini' && pendingHistoryItemRef.current?.type !== 'gemini_content' @@ -1443,20 +1508,28 @@ export const useGeminiStream = ( return newGeminiMessageBuffer; } if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder( + pendingHistoryItemRef.current, + userMessageTimestamp, + ); } - setPendingHistoryItem({ - type: 'gemini', - text: '', - timestamp: Date.now(), - }); + setPendingHistoryItem( + startAsContinuation + ? { type: 'gemini_content', text: '' } + : { type: 'gemini', text: '', timestamp: Date.now() }, + ); newGeminiMessageBuffer = stripLeadingBlankLines(newGeminiMessageBuffer); } // Split large messages for better rendering performance. Ideally, // we should maximize the amount of output sent to . - let nextPendingType = pendingHistoryItemRef.current?.type as - | 'gemini' - | 'gemini_content'; + let nextPendingType: 'gemini' | 'gemini_content' = + pendingHistoryItemRef.current?.type === 'gemini_content' + ? 'gemini_content' + : pendingHistoryItemRef.current?.type === 'gemini' + ? 'gemini' + : startAsContinuation + ? 'gemini_content' + : 'gemini'; while (newGeminiMessageBuffer.length > STREAM_PENDING_ITEM_MAX_CHARS) { const splitPoint = findLastSafeSplitPoint( newGeminiMessageBuffer, @@ -1481,7 +1554,7 @@ export const useGeminiStream = ( newGeminiMessageBuffer, safeSplitPoint, ); - commitItem( + commitItemInOrder( { type: nextPendingType, text: beforeText, @@ -1595,7 +1668,7 @@ export const useGeminiStream = ( newGeminiMessageBuffer, splitPoint, ); - commitItem( + commitItemInOrder( { type: nextPendingType, text: beforeText, @@ -1621,9 +1694,10 @@ export const useGeminiStream = ( return newGeminiMessageBuffer; }, [ - commitItem, + commitItemInOrder, pendingHistoryItemRef, setPendingHistoryItem, + stagePendingAssistantItem, terminalWidth, terminalHeight, availableTerminalHeightRef, @@ -1811,7 +1885,10 @@ export const useGeminiStream = ( }; addItem(pendingItem, userMessageTimestamp); } else { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder( + pendingHistoryItemRef.current, + userMessageTimestamp, + ); } setPendingHistoryItem(null); } @@ -1826,7 +1903,7 @@ export const useGeminiStream = ( [ addItem, commitPendingThought, - commitItem, + commitItemInOrder, pendingHistoryItemRef, setPendingHistoryItem, setThought, @@ -1848,7 +1925,7 @@ export const useGeminiStream = ( // Persist any streamed reasoning (collapsed) above the error. commitPendingThought(userMessageTimestamp); if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder(pendingHistoryItemRef.current, userMessageTimestamp); setPendingHistoryItem(null); } // Only show Ctrl+Y hint if not already showing an auto-retry countdown @@ -1892,7 +1969,7 @@ export const useGeminiStream = ( }, [ commitPendingThought, - commitItem, + commitItemInOrder, pendingHistoryItemRef, setPendingHistoryItem, setPendingRetryErrorItem, @@ -1909,14 +1986,14 @@ export const useGeminiStream = ( } if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder(pendingHistoryItemRef.current, userMessageTimestamp); setPendingHistoryItem(null); } addItem({ type: MessageType.INFO, text }, userMessageTimestamp); }, [ addItem, - commitItem, + commitItemInOrder, pendingHistoryItemRef, setPendingHistoryItem, settings, @@ -1987,7 +2064,7 @@ export const useGeminiStream = ( ) => { autonomousLoopTickResolverRef.current?.resetCache(); if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder(pendingHistoryItemRef.current, userMessageTimestamp); setPendingHistoryItem(null); } const activeModel = modelOverrideRef.current ?? config.getModel(); @@ -2011,7 +2088,13 @@ export const useGeminiStream = ( Date.now(), ); }, - [addItem, commitItem, config, pendingHistoryItemRef, setPendingHistoryItem], + [ + addItem, + commitItemInOrder, + config, + pendingHistoryItemRef, + setPendingHistoryItem, + ], ); const handleMaxSessionTurnsEvent = useCallback( @@ -2084,7 +2167,7 @@ export const useGeminiStream = ( userMessageTimestamp: number, ) => { if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder(pendingHistoryItemRef.current, userMessageTimestamp); setPendingHistoryItem(null); } addItem( @@ -2096,7 +2179,7 @@ export const useGeminiStream = ( userMessageTimestamp, ); }, - [addItem, commitItem, pendingHistoryItemRef, setPendingHistoryItem], + [addItem, commitItemInOrder, pendingHistoryItemRef, setPendingHistoryItem], ); const handleStopHookLoopEvent = useCallback( @@ -2109,7 +2192,7 @@ export const useGeminiStream = ( userMessageTimestamp: number, ) => { if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder(pendingHistoryItemRef.current, userMessageTimestamp); setPendingHistoryItem(null); } addItem( @@ -2122,7 +2205,7 @@ export const useGeminiStream = ( userMessageTimestamp, ); }, - [addItem, commitItem, pendingHistoryItemRef, setPendingHistoryItem], + [addItem, commitItemInOrder, pendingHistoryItemRef, setPendingHistoryItem], ); const processGeminiStreamEvents = useCallback( @@ -2136,6 +2219,19 @@ export const useGeminiStream = ( let geminiMessageBuffer = ''; let thoughtBuffer = ''; let scheduledToolContinuation = false; + let assistantOutputStarted = + pendingHistoryItemRef.current?.type === 'gemini' || + pendingHistoryItemRef.current?.type === 'gemini_content'; + let assistantInlineImageCount = [ + ...pendingAssistantItemsRef.current, + pendingHistoryItemRef.current, + ].reduce( + (count, item) => + item?.type === 'gemini' || item?.type === 'gemini_content' + ? count + (item.images?.length ?? 0) + : count, + 0, + ); const toolCallRequests: ToolCallRequestInfo[] = []; const bufferedEvents: BufferedStreamEvent[] = []; let flushTimer: ReturnType | null = null; @@ -2176,7 +2272,67 @@ export const useGeminiStream = ( contentParts.join(''), geminiMessageBuffer, userMessageTimestamp, + assistantOutputStarted, ); + if (contentParts.some((part) => part.trim().length > 0)) { + assistantOutputStarted = true; + } + continue; + } + + if (nextEvent.kind === 'image') { + if (turnCancelledRef.current) { + continue; + } + setIsReceivingContent(true); + turnSawContentEventRef.current = true; + const pendingItem = pendingHistoryItemRef.current; + const isOverflowOnlyItem = + (pendingItem?.type === 'gemini' || + pendingItem?.type === 'gemini_content') && + pendingItem.text.length === 0 && + !pendingItem.images?.length && + Boolean(pendingItem.omittedImageCount); + const shouldDisplayImage = + assistantInlineImageCount < MAX_INLINE_IMAGES_PER_ITEM; + + if (!shouldDisplayImage && isOverflowOnlyItem) { + setPendingHistoryItem({ + ...pendingItem, + omittedImageCount: (pendingItem.omittedImageCount ?? 0) + 1, + }); + geminiMessageBuffer = ''; + assistantOutputStarted = true; + continue; + } + + if (pendingHistoryItemRef.current) { + if (!stagePendingAssistantItem()) { + commitItemInOrder( + pendingHistoryItemRef.current, + userMessageTimestamp, + ); + setPendingHistoryItem(null); + } + } + geminiMessageBuffer = ''; + if (shouldDisplayImage) { + setPendingHistoryItem({ + type: assistantOutputStarted ? 'gemini_content' : 'gemini', + text: '', + images: [nextEvent.value], + ...(!assistantOutputStarted ? { timestamp: Date.now() } : {}), + }); + assistantInlineImageCount++; + } else { + setPendingHistoryItem({ + type: assistantOutputStarted ? 'gemini_content' : 'gemini', + text: '', + omittedImageCount: 1, + ...(!assistantOutputStarted ? { timestamp: Date.now() } : {}), + }); + } + assistantOutputStarted = true; continue; } @@ -2238,7 +2394,7 @@ export const useGeminiStream = ( scheduleBufferedStreamFlush(); } break; - case ServerGeminiEventType.Content: + case ServerGeminiEventType.Content: { // Thinking is done once the answer starts streaming; reset the // title status. On the thinking→answer transition, flush any // buffered reasoning so the full thought is captured, then commit @@ -2253,9 +2409,24 @@ export const useGeminiStream = ( thoughtBuffer = ''; } setThought((prev) => (prev ? null : prev)); - bufferedEvents.push({ kind: 'content', value: event.value }); + const displayParts = event.parts ?? [{ text: event.value }]; + for (const part of displayParts) { + if ('text' in part) { + if (part.text.length > 0) { + bufferedEvents.push({ kind: 'content', value: part.text }); + } + } else { + const image = getInlineImageData({ + inlineData: part.inlineData, + }); + if (image) { + bufferedEvents.push({ kind: 'image', value: image }); + } + } + } scheduleBufferedStreamFlush(); break; + } case ServerGeminiEventType.ToolCallRequest: // Thinking is done once a tool call is issued; flush buffered // reasoning then commit it to history (collapsed) above the tool @@ -2296,6 +2467,8 @@ export const useGeminiStream = ( case ServerGeminiEventType.ChatCompressed: flushBufferedStreamEvents(); handleChatCompressionEvent(event.value, userMessageTimestamp); + geminiMessageBuffer = ''; + assistantOutputStarted = false; break; case ServerGeminiEventType.ToolCallConfirmation: case ServerGeminiEventType.ToolCallResponse: @@ -2303,21 +2476,35 @@ export const useGeminiStream = ( break; case ServerGeminiEventType.MaxSessionTurns: flushBufferedStreamEvents(); + if (pendingHistoryItemRef.current) { + commitItemInOrder( + pendingHistoryItemRef.current, + userMessageTimestamp, + ); + setPendingHistoryItem(null); + } handleMaxSessionTurnsEvent(); + geminiMessageBuffer = ''; + assistantOutputStarted = false; break; case ServerGeminiEventType.SessionTokenLimitExceeded: flushBufferedStreamEvents(); + if (pendingHistoryItemRef.current) { + commitItemInOrder( + pendingHistoryItemRef.current, + userMessageTimestamp, + ); + setPendingHistoryItem(null); + } handleSessionTokenLimitExceededEvent(event.value); + geminiMessageBuffer = ''; + assistantOutputStarted = false; break; case ServerGeminiEventType.Finished: flushBufferedStreamEvents(); // A thinking-only turn (no content/tool) still commits its // reasoning so it persists collapsed in history. commitPendingThought(userMessageTimestamp); - handleFinishedEvent( - event as ServerGeminiFinishedEvent, - userMessageTimestamp, - ); // Seal off this turn's UI state before the parent re-enters // sendMessageStream for a continuation (Stop-hook block at // client.ts:1378 or next-speaker auto-continue at 1444). Both @@ -2327,16 +2514,29 @@ export const useGeminiStream = ( // as "t" → "te" → "tes" cumulative rendering even though each // turn is persisted as a clean, separate assistant message. if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder( + pendingHistoryItemRef.current, + userMessageTimestamp, + ); setPendingHistoryItem(null); } geminiMessageBuffer = ''; thoughtBuffer = ''; + assistantOutputStarted = false; + assistantInlineImageCount = 0; setThought(null); + handleFinishedEvent( + event as ServerGeminiFinishedEvent, + userMessageTimestamp, + ); break; case ServerGeminiEventType.Citation: flushBufferedStreamEvents(); handleCitationEvent(event.value, userMessageTimestamp); + if (showCitations(settings)) { + geminiMessageBuffer = ''; + assistantOutputStarted = false; + } break; case ServerGeminiEventType.LoopDetected: flushBufferedStreamEvents(); @@ -2354,6 +2554,7 @@ export const useGeminiStream = ( // losing the partial text we meant to preserve. if (!event.isContinuation) { discardBufferedStreamEvents(); + setPendingAssistantItems([]); if (pendingHistoryItemRef.current) { setPendingHistoryItem(null); } @@ -2361,6 +2562,8 @@ export const useGeminiStream = ( thoughtBuffer = ''; setThought(null); geminiMessageBuffer = ''; + assistantOutputStarted = false; + assistantInlineImageCount = 0; } else { flushBufferedStreamEvents(); } @@ -2384,6 +2587,7 @@ export const useGeminiStream = ( // switching to the next fallback model. Discard partial content // from the failed attempt and show a notification. discardBufferedStreamEvents(); + setPendingAssistantItems([]); if (pendingHistoryItemRef.current) { setPendingHistoryItem(null); } @@ -2391,6 +2595,8 @@ export const useGeminiStream = ( thoughtBuffer = ''; setThought(null); geminiMessageBuffer = ''; + assistantOutputStarted = false; + assistantInlineImageCount = 0; toolCallRequests.length = 0; clearRetryCountdown(); const fromModel = @@ -2410,7 +2616,10 @@ export const useGeminiStream = ( // Display system message from Stop hooks with "Stop says:" prefix // First commit any pending AI response to ensure correct ordering if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder( + pendingHistoryItemRef.current, + userMessageTimestamp, + ); setPendingHistoryItem(null); } addItem( @@ -2420,6 +2629,8 @@ export const useGeminiStream = ( } as HistoryItemWithoutId, userMessageTimestamp, ); + geminiMessageBuffer = ''; + assistantOutputStarted = false; break; case ServerGeminiEventType.UserPromptSubmitBlocked: flushBufferedStreamEvents(); @@ -2427,10 +2638,14 @@ export const useGeminiStream = ( event.value, userMessageTimestamp, ); + geminiMessageBuffer = ''; + assistantOutputStarted = false; break; case ServerGeminiEventType.StopHookLoop: flushBufferedStreamEvents(); handleStopHookLoopEvent(event.value, userMessageTimestamp); + geminiMessageBuffer = ''; + assistantOutputStarted = false; break; case ServerGeminiEventType.ActiveGoal: break; @@ -2438,7 +2653,7 @@ export const useGeminiStream = ( if (event.cause && shouldDisplayGoalStateCause(event.cause)) { flushBufferedStreamEvents(); if (pendingHistoryItemRef.current) { - commitItem( + commitItemInOrder( pendingHistoryItemRef.current, userMessageTimestamp, ); @@ -2452,6 +2667,8 @@ export const useGeminiStream = ( }, userMessageTimestamp, ); + geminiMessageBuffer = ''; + assistantOutputStarted = false; } break; default: { @@ -2585,18 +2802,22 @@ export const useGeminiStream = ( handleMaxSessionTurnsEvent, handleSessionTokenLimitExceededEvent, handleCitationEvent, + settings, startRetryCountdown, clearRetryCountdown, setThought, commitPendingThought, pendingHistoryItemRef, + pendingAssistantItemsRef, pendingThoughtItemRef, setPendingHistoryItem, handleUserPromptSubmitBlockedEvent, handleStopHookLoopEvent, bindGoalTurn, addItem, - commitItem, + commitItemInOrder, + stagePendingAssistantItem, + setPendingAssistantItems, dualOutput, ], ); @@ -2954,6 +3175,34 @@ export const useGeminiStream = ( const userMessageTimestamp = Date.now(); + // A thrown stream can leave partial assistant runs in the dynamic + // region. An explicit Ctrl+Y retry is a fresh attempt, matching a core + // non-continuation Retry event, so discard every run from the failed + // attempt before the replacement stream starts. A different top-level + // turn preserves what the user already saw, but must commit it before + // prepareQueryForGemini appends the next user item. + if (submitType === SendMessageType.Retry) { + setPendingAssistantItems([]); + const pendingItem = pendingHistoryItemRef.current; + if ( + pendingItem?.type === 'gemini' || + pendingItem?.type === 'gemini_content' + ) { + setPendingHistoryItem(null); + } + } else if (!isTurnContinuation && !allowConcurrentBtwDuringResponse) { + const pendingItem = pendingHistoryItemRef.current; + if ( + pendingItem?.type === 'gemini' || + pendingItem?.type === 'gemini_content' + ) { + commitItemInOrder(pendingItem, userMessageTimestamp); + setPendingHistoryItem(null); + } else { + commitPendingAssistantItems(userMessageTimestamp); + } + } + // Reset quota error flag when starting a new query (not a continuation). // Notifications (background agent/shell/monitor completions) are system // events, not new user turns: they must not clear the user's model @@ -3303,7 +3552,10 @@ export const useGeminiStream = ( } if (pendingHistoryItemRef.current) { - commitItem(pendingHistoryItemRef.current, userMessageTimestamp); + commitItemInOrder( + pendingHistoryItemRef.current, + userMessageTimestamp, + ); setPendingHistoryItem(null); } @@ -3509,7 +3761,9 @@ export const useGeminiStream = ( processGeminiStreamEvents, pendingHistoryItemRef, addItem, - commitItem, + commitPendingAssistantItems, + commitItemInOrder, + setPendingAssistantItems, setPendingHistoryItem, setInitError, geminiClient, @@ -4358,6 +4612,7 @@ export const useGeminiStream = ( [ // Reasoning renders above the streaming answer. pendingThoughtItem, + ...pendingAssistantItems, pendingHistoryItem, pendingRetryErrorItem, pendingRetryCountdownItem, @@ -4365,6 +4620,7 @@ export const useGeminiStream = ( ].filter((i) => i !== undefined && i !== null), [ pendingThoughtItem, + pendingAssistantItems, pendingHistoryItem, pendingRetryErrorItem, pendingRetryCountdownItem, @@ -4928,6 +5184,7 @@ export const useGeminiStream = ( submitQuery, initError, pendingHistoryItems, + clearPendingState, thought, cancelOngoingRequest, preemptGoalTurn, diff --git a/packages/cli/src/ui/hooks/useHistoryManager.test.ts b/packages/cli/src/ui/hooks/useHistoryManager.test.ts index 382115a50bd..b88e5738ee4 100644 --- a/packages/cli/src/ui/hooks/useHistoryManager.test.ts +++ b/packages/cli/src/ui/hooks/useHistoryManager.test.ts @@ -6,7 +6,11 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { renderHook, act } from '@testing-library/react'; -import { useHistory, UI_COMPACT_CLEARED_MESSAGE } from './useHistoryManager.js'; +import { + useHistory, + UI_COMPACT_CLEARED_MESSAGE, + UI_COMPACT_CLEARED_IMAGE_MESSAGE, +} from './useHistoryManager.js'; import type { UseHistoryManagerReturn } from './useHistoryManager.js'; import type { HistoryItemWithoutId, HistoryItemToolGroup } from '../types.js'; import { ToolCallStatus } from '../types.js'; @@ -370,6 +374,136 @@ describe('useHistoryManager', () => { expect(recentTool.detailedDisplay).toBe('full secret file content here'); }); + it('clears image payloads from old tool results', () => { + const { result } = renderHook(() => useHistory()); + const ts = Date.now(); + + for (let i = 0; i < 25; i++) { + act(() => { + result.current.addItem( + { + type: 'tool_group', + tools: [ + { + callId: String(i), + name: 'screenshot', + description: '', + resultDisplay: undefined, + images: [{ data: 'aW1hZ2U=', mimeType: 'image/png' }], + omittedImageCount: 2, + status: ToolCallStatus.Success, + confirmationDetails: undefined, + }, + ], + } as unknown as HistoryItemWithoutId, + ts + i, + ); + }); + } + + act(() => { + result.current.compactOldItems(); + }); + + const oldestTool = ( + result.current.history[0] as unknown as HistoryItemToolGroup + ).tools[0]; + expect(oldestTool.resultDisplay).toBe(UI_COMPACT_CLEARED_MESSAGE); + expect(oldestTool.images).toBeUndefined(); + expect(oldestTool.omittedImageCount).toBeUndefined(); + + const recentTool = ( + result.current.history[24] as unknown as HistoryItemToolGroup + ).tools[0]; + expect(recentTool.images).toHaveLength(1); + expect(recentTool.omittedImageCount).toBe(2); + }); + + it('clears old assistant image payloads while keeping recent images', () => { + const { result } = renderHook(() => useHistory()); + const ts = Date.now(); + + for (let i = 0; i < 25; i++) { + act(() => { + result.current.addItem( + { + type: i === 0 ? 'gemini' : 'gemini_content', + text: i === 0 ? 'Generated chart' : '', + images: [{ data: `aW1hZ2Ut${i}`, mimeType: 'image/png' }], + omittedImageCount: 2, + }, + ts + i, + ); + }); + } + + act(() => { + result.current.compactOldItems(); + }); + + const oldestItem = result.current.history[0]; + expect(oldestItem).toMatchObject({ + type: 'gemini', + text: `Generated chart\n\n${UI_COMPACT_CLEARED_IMAGE_MESSAGE}`, + }); + expect( + oldestItem.type === 'gemini' ? oldestItem.images : undefined, + ).toBeUndefined(); + expect( + oldestItem.type === 'gemini' ? oldestItem.omittedImageCount : undefined, + ).toBeUndefined(); + + const recentItem = result.current.history[24]; + expect( + recentItem.type === 'gemini_content' ? recentItem.images : undefined, + ).toHaveLength(1); + expect( + recentItem.type === 'gemini_content' + ? recentItem.omittedImageCount + : undefined, + ).toBe(2); + }); + + it('compacts old assistant image overflow markers without payloads', () => { + const { result } = renderHook(() => useHistory()); + const ts = Date.now(); + + for (let i = 0; i < 25; i++) { + act(() => { + result.current.addItem( + { + type: 'gemini_content', + text: '', + omittedImageCount: 2, + }, + ts + i, + ); + }); + } + + act(() => { + result.current.compactOldItems(); + }); + + const oldestItem = result.current.history[0]; + expect(oldestItem).toMatchObject({ + type: 'gemini_content', + text: UI_COMPACT_CLEARED_IMAGE_MESSAGE, + }); + expect( + oldestItem.type === 'gemini_content' + ? oldestItem.omittedImageCount + : undefined, + ).toBeUndefined(); + + const recentItem = result.current.history[24]; + expect( + recentItem.type === 'gemini_content' + ? recentItem.omittedImageCount + : undefined, + ).toBe(2); + }); + it('clears a tool that carries detailedDisplay but no resultDisplay (defensive)', () => { const { result } = renderHook(() => useHistory()); const ts = Date.now(); diff --git a/packages/cli/src/ui/hooks/useHistoryManager.ts b/packages/cli/src/ui/hooks/useHistoryManager.ts index fc5b41df276..9aed5ba379d 100644 --- a/packages/cli/src/ui/hooks/useHistoryManager.ts +++ b/packages/cli/src/ui/hooks/useHistoryManager.ts @@ -17,6 +17,8 @@ type HistoryItemUpdater = ( ) => Partial; export const UI_COMPACT_CLEARED_MESSAGE = '[Old tool result content cleared]'; +export const UI_COMPACT_CLEARED_IMAGE_MESSAGE = + '[Old assistant image content cleared]'; const UI_COMPACT_KEEP_RECENT = 20; export interface UseHistoryManagerReturn { @@ -148,22 +150,30 @@ export function useHistory(): UseHistoryManagerReturn { let thoughtRemoved = 0; let toolGroupsCompacted = 0; + let assistantImageItemsCompacted = 0; let totalThoughts = 0; let totalToolGroupsWithOutput = 0; + let totalAssistantItemsWithImages = 0; for (const item of prev) { if ( item.type === 'gemini_thought' || item.type === 'gemini_thought_content' ) { totalThoughts++; + } else if ( + (item.type === 'gemini' || item.type === 'gemini_content') && + (item.images?.length || item.omittedImageCount) + ) { + totalAssistantItemsWithImages++; } else if ( item.type === 'tool_group' && item.tools.some( (t) => (t.resultDisplay != null && t.resultDisplay !== UI_COMPACT_CLEARED_MESSAGE) || - t.detailedDisplay != null, + t.detailedDisplay != null || + Boolean(t.images?.length || t.omittedImageCount), ) ) { totalToolGroupsWithOutput++; @@ -177,8 +187,13 @@ export function useHistory(): UseHistoryManagerReturn { 0, totalToolGroupsWithOutput - UI_COMPACT_KEEP_RECENT, ); + const assistantImageItemsToCompact = Math.max( + 0, + totalAssistantItemsWithImages - UI_COMPACT_KEEP_RECENT, + ); let thoughtsDropped = 0; let toolGroupsSeen = 0; + let assistantImageItemsSeen = 0; const next = prev .filter((item) => { @@ -195,16 +210,34 @@ export function useHistory(): UseHistoryManagerReturn { return true; }) .map((item) => { + if ( + (item.type === 'gemini' || item.type === 'gemini_content') && + (item.images?.length || item.omittedImageCount) + ) { + assistantImageItemsSeen++; + if (assistantImageItemsSeen <= assistantImageItemsToCompact) { + assistantImageItemsCompacted++; + return { + ...item, + text: item.text + ? `${item.text}\n\n${UI_COMPACT_CLEARED_IMAGE_MESSAGE}` + : UI_COMPACT_CLEARED_IMAGE_MESSAGE, + images: undefined, + omittedImageCount: undefined, + }; + } + } if (item.type !== 'tool_group') return item; // Check for any non-null resultDisplay (covers string, FileDiff, - // AnsiOutputDisplay, AgentResultDisplay, etc.) OR a lingering - // `detailedDisplay` — so a tool that somehow carries only the raw - // transcript detail still triggers compaction and gets cleared. + // AnsiOutputDisplay, AgentResultDisplay, etc.), a lingering + // `detailedDisplay`, or image payloads. Every retained output form + // must participate in the same keep-recent limit. const hasOldOutput = item.tools.some( (t) => (t.resultDisplay != null && t.resultDisplay !== UI_COMPACT_CLEARED_MESSAGE) || - t.detailedDisplay != null, + t.detailedDisplay != null || + Boolean(t.images?.length || t.omittedImageCount), ); if (!hasOldOutput) return item; toolGroupsSeen++; @@ -216,19 +249,23 @@ export function useHistory(): UseHistoryManagerReturn { if ( (t.resultDisplay != null && t.resultDisplay !== UI_COMPACT_CLEARED_MESSAGE) || - t.detailedDisplay != null + t.detailedDisplay != null || + t.images?.length || + t.omittedImageCount ) { // Also drop `detailedDisplay` (the raw functionResponse text // kept for the Ctrl+O full-detail transcript): clearing only // `resultDisplay` would let a post-compaction transcript reopen // re-surface the supposedly cleared read/search/list output, // defeating the memory/privacy compaction. The `detailedDisplay` - // arm keeps the guard robust even if a tool ever carries the - // raw detail without a `resultDisplay`. + // and `images` arms keep the guard robust when a tool carries + // raw detail or media without a `resultDisplay`. return { ...t, resultDisplay: UI_COMPACT_CLEARED_MESSAGE, detailedDisplay: undefined, + images: undefined, + omittedImageCount: undefined, }; } return t; @@ -236,17 +273,26 @@ export function useHistory(): UseHistoryManagerReturn { }; }); - if (thoughtRemoved > 0 || toolGroupsCompacted > 0) { + if ( + thoughtRemoved > 0 || + toolGroupsCompacted > 0 || + assistantImageItemsCompacted > 0 + ) { if (debugLogger.isEnabled()) { debugLogger.debug( `[COMPACT_UI_HISTORY] removed ${thoughtRemoved} thought item(s), ` + + `compacted ${assistantImageItemsCompacted} assistant image item(s), ` + `compacted ${toolGroupsCompacted} tool group(s), ` + `historyLength ${prev.length} -> ${next.length}, ` + `memory=${(process.memoryUsage().heapUsed / 1024 / 1024).toFixed(1)}MB`, ); } } - return thoughtRemoved > 0 || toolGroupsCompacted > 0 ? next : prev; + return thoughtRemoved > 0 || + toolGroupsCompacted > 0 || + assistantImageItemsCompacted > 0 + ? next + : prev; }); }, []); diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.test.tsx b/packages/cli/src/ui/hooks/useReactToolScheduler.test.tsx index f68ee1b1148..0c947c2762e 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.test.tsx +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.test.tsx @@ -5,13 +5,19 @@ */ import { describe, it, expect } from 'vitest'; +import type { Part } from '@google/genai'; import { mapToDisplay, type TrackedToolCall } from './useReactToolScheduler.js'; +import { MAX_INLINE_IMAGES_PER_ITEM } from '../utils/inline-image-parts.js'; // Build a minimal successful tracked tool call with the fields mapToDisplay's // success branch reads. `displayName` drives the collapsible gate. -const makeSuccess = (displayName: string): TrackedToolCall => +const makeCompleted = ( + status: 'success' | 'error' | 'cancelled', + displayName: string, + responseMedia: Part[] = [], +): TrackedToolCall => ({ - status: 'success', + status, request: { callId: 'call-1', name: 'read_file', args: {} }, tool: { displayName, isOutputMarkdown: false }, invocation: { getDescription: () => 'reading' }, @@ -23,12 +29,18 @@ const makeSuccess = (displayName: string): TrackedToolCall => id: 'call-1', name: 'read_file', response: { output: 'FULL FILE CONTENT' }, + ...(responseMedia.length > 0 ? { parts: responseMedia } : {}), }, }, ], }, }) as unknown as TrackedToolCall; +const makeSuccess = ( + displayName: string, + responseMedia: Part[] = [], +): TrackedToolCall => makeCompleted('success', displayName, responseMedia); + describe('mapToDisplay — detailedDisplay (§4.9 live path)', () => { it('extracts detailedDisplay for a collapsible (read/search/list) tool', () => { const group = mapToDisplay(makeSuccess('Read File')); @@ -45,4 +57,50 @@ describe('mapToDisplay — detailedDisplay (§4.9 live path)', () => { const group = mapToDisplay(makeSuccess('Edit')); expect(group.tools[0].detailedDisplay).toBeUndefined(); }); + + it.each(['success', 'error', 'cancelled'] as const)( + 'extracts nested inline images from %s tool response parts', + (status) => { + const group = mapToDisplay( + makeCompleted(status, 'Read File', [ + { + inlineData: { + data: 'dG9vbC1pbWFnZQ==', + mimeType: 'image/png', + displayName: 'result.png', + }, + }, + ]), + ); + + expect(group.tools[0].images).toEqual([ + { + data: 'dG9vbC1pbWFnZQ==', + mimeType: 'image/png', + }, + ]); + }, + ); + + it('caps tool images and reports the overflow count', () => { + const images = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM + 2 }, + (_, index) => ({ + inlineData: { + data: Buffer.from(`tool-image-${index}`).toString('base64'), + mimeType: 'image/png', + }, + }), + ); + + const tool = mapToDisplay(makeCompleted('success', 'Read File', images)) + .tools[0]; + + expect(tool.images).toEqual( + images + .slice(0, MAX_INLINE_IMAGES_PER_ITEM) + .map((part) => part.inlineData), + ); + expect(tool.omittedImageCount).toBe(2); + }); }); diff --git a/packages/cli/src/ui/hooks/useReactToolScheduler.ts b/packages/cli/src/ui/hooks/useReactToolScheduler.ts index bd822e1f9c9..f834f43ace5 100644 --- a/packages/cli/src/ui/hooks/useReactToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useReactToolScheduler.ts @@ -38,6 +38,7 @@ import type { } from '../types.js'; import { ToolCallStatus } from '../types.js'; import { isCollapsibleTool } from '../components/messages/CompactToolGroupDisplay.js'; +import { collectInlineImages } from '../utils/inline-image-parts.js'; const debugLogger = createDebugLogger('REACT_TOOL_SCHEDULER'); @@ -390,8 +391,15 @@ export function mapToDisplay( : undefined, }; + const inlineImageCollection = + trackedCall.status === 'success' || + trackedCall.status === 'error' || + trackedCall.status === 'cancelled' + ? collectInlineImages(trackedCall.response.responseParts) + : null; + switch (trackedCall.status) { - case 'success': + case 'success': { return { ...baseDisplayProperties, status: mapCoreStatusToDisplayStatus(trackedCall.status), @@ -415,9 +423,19 @@ export function mapToDisplay( detailedDisplay: isCollapsibleTool(displayName) ? getToolResponseDisplayText(trackedCall.response.responseParts) : undefined, + ...(inlineImageCollection?.images.length + ? { images: inlineImageCollection.images } + : {}), + ...(inlineImageCollection?.omittedImageCount + ? { + omittedImageCount: inlineImageCollection.omittedImageCount, + } + : {}), confirmationDetails: undefined, }; + } case 'error': + case 'cancelled': { return { ...baseDisplayProperties, status: mapCoreStatusToDisplayStatus(trackedCall.status), @@ -429,22 +447,17 @@ export function mapToDisplay( visionBridgeNotice: trackedCall.response.visionBridgeNotice, } : {}), - confirmationDetails: undefined, - }; - case 'cancelled': - return { - ...baseDisplayProperties, - status: mapCoreStatusToDisplayStatus(trackedCall.status), - resultDisplay: compactToolResultDisplayForHistory( - trackedCall.response.resultDisplay, - ), - ...(trackedCall.response.visionBridgeNotice !== undefined + ...(inlineImageCollection?.images.length + ? { images: inlineImageCollection.images } + : {}), + ...(inlineImageCollection?.omittedImageCount ? { - visionBridgeNotice: trackedCall.response.visionBridgeNotice, + omittedImageCount: inlineImageCollection.omittedImageCount, } : {}), confirmationDetails: undefined, }; + } case 'awaiting_approval': return { ...baseDisplayProperties, diff --git a/packages/cli/src/ui/hooks/useResumeCommand.test.ts b/packages/cli/src/ui/hooks/useResumeCommand.test.ts index 1ad5dfbca1d..6d3555c776e 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.test.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.test.ts @@ -240,6 +240,7 @@ describe('useResumeCommand', () => { loadHistory: vi.fn(), }; const startNewSession = vi.fn(); + const clearPendingState = vi.fn(); const geminiClient = { initialize: vi.fn().mockResolvedValue(undefined), }; @@ -287,6 +288,7 @@ describe('useResumeCommand', () => { settings: mockSettings, historyManager, startNewSession, + clearPendingState, }), ); @@ -325,6 +327,10 @@ describe('useResumeCommand', () => { expect(geminiClient.initialize).toHaveBeenCalledWith(); expect(historyManager.clearItems).toHaveBeenCalledTimes(1); expect(historyManager.loadHistory).toHaveBeenCalledTimes(1); + expect(clearPendingState).toHaveBeenCalledTimes(1); + expect(clearPendingState.mock.invocationCallOrder[0]).toBeLessThan( + historyManager.loadHistory.mock.invocationCallOrder[0]!, + ); expect(resetMonitorRegistry).toHaveBeenCalledTimes(1); expect(config.getGoalRuntimeReady).toHaveBeenCalledTimes(1); }); diff --git a/packages/cli/src/ui/hooks/useResumeCommand.ts b/packages/cli/src/ui/hooks/useResumeCommand.ts index 646da4a2dca..276c2f141b0 100644 --- a/packages/cli/src/ui/hooks/useResumeCommand.ts +++ b/packages/cli/src/ui/hooks/useResumeCommand.ts @@ -32,6 +32,7 @@ export interface UseResumeCommandOptions { 'addItem' | 'clearItems' | 'loadHistory' >; startNewSession: (sessionId: string) => void; + clearPendingState?: () => void; setSessionName?: (name: string | null) => void; remount?: () => void; } @@ -80,6 +81,7 @@ export function useResumeCommand( settings, historyManager, startNewSession, + clearPendingState, setSessionName, remount, } = options; @@ -177,6 +179,7 @@ export function useResumeCommand( // into the old JSONL (split-brain). startNewSession(sessionId); setSessionName?.(customTitle ?? null); + clearPendingState?.(); clearItems(); loadHistory(uiHistoryItems); if (recoveredBackgroundAgentsNotice) { @@ -242,6 +245,7 @@ export function useResumeCommand( clearItems, loadHistory, startNewSession, + clearPendingState, setSessionName, remount, settings.merged.ui?.history?.collapseOnResume, diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 29baace5db0..469efda1485 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -54,6 +54,11 @@ export enum ToolCallStatus { Error = 'Error', } +export interface InlineImageData { + data: string; + mimeType: string; +} + export interface ToolCallEvent { type: 'tool_call'; status: ToolCallStatus; @@ -79,6 +84,10 @@ export interface IndividualToolCallDisplay { * is only a count. Undefined → fall back to the summary. */ detailedDisplay?: string; + /** Inline images carried by this tool's persisted response parts. */ + images?: InlineImageData[]; + /** Images hidden after the per-row rendering limit. */ + omittedImageCount?: number; status: ToolCallStatus; confirmationDetails: ToolCallConfirmationDetails | undefined; renderOutputAsMarkdown?: boolean; @@ -138,12 +147,16 @@ export type HistoryItemUser = HistoryItemBase & { export type HistoryItemGemini = HistoryItemBase & { type: 'gemini'; text: string; + images?: InlineImageData[]; + omittedImageCount?: number; timestamp?: number; }; export type HistoryItemGeminiContent = HistoryItemBase & { type: 'gemini_content'; text: string; + images?: InlineImageData[]; + omittedImageCount?: number; }; export type HistoryItemGeminiThought = HistoryItemBase & { diff --git a/packages/cli/src/ui/utils/inline-image-parts.test.ts b/packages/cli/src/ui/utils/inline-image-parts.test.ts new file mode 100644 index 00000000000..b92b3d570c1 --- /dev/null +++ b/packages/cli/src/ui/utils/inline-image-parts.test.ts @@ -0,0 +1,139 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { + collectInlineImages, + extractInlineContentRuns, + getInlineImageData, + MAX_INLINE_IMAGE_ENCODED_LENGTH, + MAX_INLINE_IMAGES_PER_ITEM, +} from './inline-image-parts.js'; + +describe('collectInlineImages', () => { + it('extracts an image from a top-level tool response part', () => { + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + displayName: 'chart.png', + }; + + expect(collectInlineImages([{ inlineData: image }])).toEqual({ + images: [{ data: image.data, mimeType: image.mimeType }], + omittedImageCount: 0, + }); + }); + + it('extracts an image from nested function response parts', () => { + const image = { + data: 'bmVzdGVkLWltYWdl', + mimeType: 'image/webp', + }; + + expect( + collectInlineImages([ + { + functionResponse: { + id: 'call-1', + name: 'generate_image', + response: { output: 'done' }, + parts: [{ inlineData: image }], + }, + }, + ]), + ).toEqual({ images: [image], omittedImageCount: 0 }); + }); + + it('ignores non-image inline data', () => { + expect( + collectInlineImages([ + { + inlineData: { + data: 'bm90LWFuLWltYWdl', + mimeType: 'text/plain', + }, + }, + ]), + ).toEqual({ images: [], omittedImageCount: 0 }); + }); + + it('caps images and reports how many were omitted', () => { + const images = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM + 2 }, + (_, index) => ({ + data: Buffer.from(`image-${index}`).toString('base64'), + mimeType: 'image/png', + }), + ); + + expect( + collectInlineImages(images.map((inlineData) => ({ inlineData }))), + ).toEqual({ + images: images.slice(0, MAX_INLINE_IMAGES_PER_ITEM), + omittedImageCount: 2, + }); + }); +}); + +describe('getInlineImageData', () => { + it('rejects payloads above the renderer encoded-length limit', () => { + expect( + getInlineImageData({ + inlineData: { + data: 'A'.repeat(MAX_INLINE_IMAGE_ENCODED_LENGTH + 1), + mimeType: 'image/png', + }, + }), + ).toBeNull(); + }); +}); + +describe('extractInlineContentRuns', () => { + it('preserves text-image-text order and skips thought parts', () => { + const image = { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }; + + expect( + extractInlineContentRuns([ + { text: 'before' }, + { text: 'hidden reasoning', thought: true }, + { inlineData: image }, + { text: 'after' }, + ]), + ).toEqual([ + { kind: 'text', text: 'before' }, + { kind: 'image', image }, + { kind: 'text', text: 'after' }, + ]); + }); + + it('caps images and preserves an overflow marker at the first omission', () => { + const images = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM + 2 }, + (_, index) => ({ + data: Buffer.from(`image-${index}`).toString('base64'), + mimeType: 'image/png', + }), + ); + + expect( + extractInlineContentRuns([ + { text: 'before' }, + ...images.map((inlineData) => ({ inlineData })), + { text: 'after' }, + ]), + ).toEqual([ + { kind: 'text', text: 'before' }, + ...images + .slice(0, MAX_INLINE_IMAGES_PER_ITEM) + .map((image) => ({ kind: 'image' as const, image })), + { kind: 'omitted_images', count: 2 }, + { kind: 'text', text: 'after' }, + ]); + }); +}); diff --git a/packages/cli/src/ui/utils/inline-image-parts.ts b/packages/cli/src/ui/utils/inline-image-parts.ts new file mode 100644 index 00000000000..45c045b3962 --- /dev/null +++ b/packages/cli/src/ui/utils/inline-image-parts.ts @@ -0,0 +1,119 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Part } from '@google/genai'; +import { MAX_TERMINAL_IMAGE_BYTES } from '@qwen-code/qwen-code-core'; +import type { InlineImageData } from '../types.js'; + +export const MAX_INLINE_IMAGES_PER_ITEM = 4; +export const MAX_INLINE_IMAGE_ENCODED_LENGTH = + Math.ceil((MAX_TERMINAL_IMAGE_BYTES * 4) / 3) + 4; + +export type InlineContentRun = + | { kind: 'text'; text: string } + | { kind: 'image'; image: InlineImageData } + | { kind: 'omitted_images'; count: number }; + +export interface InlineImageCollection { + images: InlineImageData[]; + omittedImageCount: number; +} + +export function formatInlineImageOverflow(count: number): string { + return `[+${count} more ${count === 1 ? 'image' : 'images'}]`; +} + +export function getInlineImageData(part: Part): InlineImageData | null { + const inlineData = part.inlineData; + if ( + !inlineData?.mimeType?.trim().toLowerCase().startsWith('image/') || + typeof inlineData.data !== 'string' || + inlineData.data.length === 0 || + inlineData.data.length > MAX_INLINE_IMAGE_ENCODED_LENGTH + ) { + return null; + } + + return { + data: inlineData.data, + mimeType: inlineData.mimeType, + }; +} + +export function collectInlineImages( + parts: Part[] | undefined, +): InlineImageCollection { + if (!parts) { + return { images: [], omittedImageCount: 0 }; + } + + const images: InlineImageData[] = []; + let omittedImageCount = 0; + const collectImage = (part: Part) => { + const image = getInlineImageData(part); + if (!image) { + return; + } + if (images.length < MAX_INLINE_IMAGES_PER_ITEM) { + images.push(image); + } else { + omittedImageCount++; + } + }; + + for (const part of parts) { + collectImage(part); + + for (const nested of part.functionResponse?.parts ?? []) { + collectImage(nested as Part); + } + } + return { images, omittedImageCount }; +} + +export function extractInlineContentRuns( + parts: Part[] | undefined, + textSeparator = '', +): InlineContentRun[] { + if (!parts) { + return []; + } + + const runs: InlineContentRun[] = []; + let textParts: string[] = []; + let displayedImageCount = 0; + let overflowRun: Extract< + InlineContentRun, + { kind: 'omitted_images' } + > | null = null; + const flushText = () => { + if (textParts.length === 0) return; + runs.push({ kind: 'text', text: textParts.join(textSeparator) }); + textParts = []; + }; + + for (const part of parts) { + if (part.thought) continue; + if (part.text) { + textParts.push(part.text); + } + const image = getInlineImageData(part); + if (image) { + flushText(); + if (displayedImageCount < MAX_INLINE_IMAGES_PER_ITEM) { + runs.push({ kind: 'image', image }); + displayedImageCount++; + } else if (overflowRun) { + overflowRun.count++; + } else { + overflowRun = { kind: 'omitted_images', count: 1 }; + runs.push(overflowRun); + } + } + } + flushText(); + return runs; +} diff --git a/packages/cli/src/ui/utils/resumeHistoryUtils.test.ts b/packages/cli/src/ui/utils/resumeHistoryUtils.test.ts index 0a6bbdcd993..90557deb34c 100644 --- a/packages/cli/src/ui/utils/resumeHistoryUtils.test.ts +++ b/packages/cli/src/ui/utils/resumeHistoryUtils.test.ts @@ -21,6 +21,7 @@ import type { } from '@qwen-code/qwen-code-core'; import type { Part } from '@google/genai'; import type { HistoryItem } from '../types.js'; +import { MAX_INLINE_IMAGES_PER_ITEM } from './inline-image-parts.js'; const makeConfig = (tools: Record) => ({ @@ -805,6 +806,166 @@ describe('resumeHistoryUtils', () => { expect(items[0]).not.toHaveProperty('sentToModel'); }); + // The current Core recorder flattens assistant output before persistence. + // This fixture covers the parser for records written by a compatible writer. + it('parses persisted assistant text and images in their original order', () => { + const conversation = { + messages: [ + { + type: 'assistant', + timestamp: '2026-01-15T19:00:00.000Z', + message: { + parts: [ + { text: 'before' } as Part, + { + inlineData: { + data: 'aW1hZ2U=', + mimeType: 'image/png', + displayName: 'chart.png', + }, + } as Part, + { text: 'after' } as Part, + ], + }, + }, + ], + } as unknown as ConversationRecord; + + const items = buildResumedHistoryItems( + { conversation } as ResumedSessionData, + makeConfig({}), + 100, + ); + + expect(items).toEqual([ + { + id: 101, + type: 'gemini', + text: 'before', + timestamp: new Date('2026-01-15T19:00:00.000Z').getTime(), + }, + { + id: 102, + type: 'gemini_content', + text: '', + images: [ + { + data: 'aW1hZ2U=', + mimeType: 'image/png', + }, + ], + }, + { id: 103, type: 'gemini_content', text: 'after' }, + ]); + }); + + it('caps persisted assistant images and retains the overflow count', () => { + const images = Array.from( + { length: MAX_INLINE_IMAGES_PER_ITEM + 2 }, + (_, index) => ({ + data: Buffer.from(`restored-image-${index}`).toString('base64'), + mimeType: 'image/png', + }), + ); + const conversation = { + messages: [ + { + type: 'assistant', + timestamp: '2026-01-15T19:00:00.000Z', + message: { + parts: images.map((inlineData) => ({ inlineData })), + }, + }, + ], + } as unknown as ConversationRecord; + + const items = buildResumedHistoryItems( + { conversation } as ResumedSessionData, + makeConfig({}), + ).filter( + (item) => item.type === 'gemini' || item.type === 'gemini_content', + ); + + expect(items.flatMap((item) => item.images ?? [])).toEqual( + images.slice(0, MAX_INLINE_IMAGES_PER_ITEM), + ); + expect(items.at(-1)).toMatchObject({ + type: 'gemini_content', + text: '', + omittedImageCount: 2, + }); + }); + + it('restores images nested in persisted tool response parts', () => { + const conversation = { + messages: [ + { + type: 'assistant', + message: { + parts: [ + { + functionCall: { + id: 'call-image', + name: 'replace', + args: {}, + }, + } as unknown as Part, + ], + }, + }, + { + type: 'tool_result', + toolCallResult: { + callId: 'call-image', + resultDisplay: 'Generated chart', + status: 'success', + responseParts: [ + { + functionResponse: { + id: 'call-image', + name: 'replace', + response: { output: 'Generated chart' }, + parts: [ + { + inlineData: { + data: 'dG9vbC1pbWFnZQ==', + mimeType: 'image/webp', + }, + }, + ], + }, + }, + ], + }, + }, + ], + } as unknown as ConversationRecord; + + const items = buildResumedHistoryItems( + { conversation } as ResumedSessionData, + makeConfig({ replace: mockTool }), + 200, + ); + + expect(items).toEqual([ + { + id: 201, + type: 'tool_group', + tools: [ + expect.objectContaining({ + callId: 'call-image', + images: [ + { + data: 'dG9vbC1pbWFnZQ==', + mimeType: 'image/webp', + }, + ], + }), + ], + }, + ]); + }); + describe('detailedDisplay (§4.9 Ctrl+O full detail on resume)', () => { type ToolGroupItem = Extract; const firstTool = (items: HistoryItem[]) => diff --git a/packages/cli/src/ui/utils/resumeHistoryUtils.ts b/packages/cli/src/ui/utils/resumeHistoryUtils.ts index 385898bf9de..6ed5d934d48 100644 --- a/packages/cli/src/ui/utils/resumeHistoryUtils.ts +++ b/packages/cli/src/ui/utils/resumeHistoryUtils.ts @@ -27,6 +27,7 @@ import type { HistoryItemInfo, HistoryItemWithoutId, IndividualToolCallDisplay, + InlineImageData, } from '../types.js'; import { ToolCallStatus, MessageType } from '../types.js'; import { t } from '../../i18n/index.js'; @@ -36,6 +37,10 @@ import { indexGapsByChild, } from './history-gap-notice.js'; import { shouldDisplayGoalStateCause } from './goal-runtime.js'; +import { + collectInlineImages, + extractInlineContentRuns, +} from './inline-image-parts.js'; /** * Projects a plain user record to its display text. @@ -217,6 +222,8 @@ function convertToHistoryItems( resultDisplay: ToolResultDisplay | undefined; visionBridgeNotice?: string; detailedDisplay?: string; + images?: InlineImageData[]; + omittedImageCount?: number; status: ToolCallStatus; confirmationDetails: undefined; }> = []; @@ -439,8 +446,7 @@ function convertToHistoryItems( // verbatim because there is no live loading area in that view. const thoughtText = !config ? extractThoughtTextFromParts(parts) : ''; - // Extract text content (non-function-call, non-thought) - const text = extractTextFromParts(parts); + const displayRuns = extractInlineContentRuns(parts, '\n'); // Extract function calls const functionCalls = extractFunctionCalls(parts); @@ -458,9 +464,8 @@ function convertToHistoryItems( items.push({ type: 'gemini_thought', text: thoughtText }); } - // If there's text content, add it as a gemini message - if (text) { - // Flush any pending tool group before text + if (displayRuns.length > 0) { + // Flush any pending tool group before assistant output. if (currentToolGroup.length > 0) { items.push({ type: 'tool_group', @@ -468,11 +473,30 @@ function convertToHistoryItems( }); currentToolGroup = []; } - items.push({ - type: 'gemini', - text, - timestamp: new Date(record.timestamp).getTime(), - }); + for (const [index, run] of displayRuns.entries()) { + const type = index === 0 ? 'gemini' : 'gemini_content'; + const timestamp = + index === 0 + ? { timestamp: new Date(record.timestamp).getTime() } + : {}; + if (run.kind === 'text') { + items.push({ type, text: run.text, ...timestamp }); + } else if (run.kind === 'image') { + items.push({ + type, + text: '', + images: [run.image], + ...timestamp, + }); + } else { + items.push({ + type, + text: '', + omittedImageCount: run.count, + ...timestamp, + }); + } + } } // Track function calls for pairing with results @@ -500,6 +524,9 @@ function convertToHistoryItems( const callId = record.toolCallResult.callId; const toolCall = currentToolGroup.find((t) => t.callId === callId); if (toolCall) { + const responseParts = + (record.toolCallResult.responseParts as Part[] | undefined) ?? + (record.message?.parts as Part[] | undefined); // Preserve the resultDisplay as-is - it can be a string or structured object const rawDisplay = record.toolCallResult.resultDisplay; toolCall.resultDisplay = rawDisplay; @@ -515,6 +542,14 @@ function convertToHistoryItems( rawStatus === 'error' ? ToolCallStatus.Error : ToolCallStatus.Success; + const { images, omittedImageCount } = + collectInlineImages(responseParts); + if (images.length > 0) { + toolCall.images = images; + } + if (omittedImageCount > 0) { + toolCall.omittedImageCount = omittedImageCount; + } // Full detail for the Ctrl+O transcript (§4.9): the complete // functionResponse parts are persisted on the tool_result record // (only resultDisplay is sanitized), so resume yields full detail @@ -529,10 +564,8 @@ function convertToHistoryItems( toolCall.status === ToolCallStatus.Success && isCollapsibleTool(toolCall.name) ) { - toolCall.detailedDisplay = getToolResponseDisplayText( - (record.toolCallResult.responseParts as Part[] | undefined) ?? - (record.message?.parts as Part[] | undefined), - ); + toolCall.detailedDisplay = + getToolResponseDisplayText(responseParts); } } pendingToolCalls.delete(callId || ''); diff --git a/packages/cli/src/ui/utils/terminal-image-renderer.test.ts b/packages/cli/src/ui/utils/terminal-image-renderer.test.ts index 6065386c558..fc993c35ee2 100644 --- a/packages/cli/src/ui/utils/terminal-image-renderer.test.ts +++ b/packages/cli/src/ui/utils/terminal-image-renderer.test.ts @@ -11,12 +11,15 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { containsCmdShellMetacharacters, getTerminalImageRenderSupport, + MAX_INLINE_IMAGE_PIXELS, markKittyImageWritten, + prepareInlineTerminalImage, renderTerminalImage, supportsKittyImageProtocol, TRANSMITTED_KEY_LIMIT, wasKittyImageWritten, } from './terminal-image-renderer.js'; +import { MAX_INLINE_IMAGE_ENCODED_LENGTH } from './inline-image-parts.js'; const PNG_1X1 = Buffer.from( 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+/p9sAAAAASUVORK5CYII=', @@ -86,6 +89,102 @@ describe('terminalImageRenderer', () => { expect(result.placeholder.lines[0]).toContain('\u{10EEEE}'); }); + it('renders bounded inline PNG data through the Kitty path', () => { + const prepared = prepareInlineTerminalImage({ + data: PNG_1X1.toString('base64'), + mimeType: 'image/png', + contentWidth: 24, + availableTerminalHeight: 12, + env: { TERM: 'xterm-kitty' }, + stdoutIsTTY: true, + }); + + expect(prepared.fallbackText).toBe('[image: 1x1 png]'); + expect(prepared.result?.kind).toBe('kitty'); + if (prepared.result?.kind !== 'kitty') return; + expect(prepared.result.sequence).toContain('\u001b_Ga=T,f=100'); + expect(prepared.result.placeholder.lines).toHaveLength(1); + }); + + it('validates inline PNG payloads before rendering', () => { + for (const testCase of [ + { data: 'AB==', mimeType: 'image/png', fallback: '[image: png]' }, + { + data: Buffer.from('not a png').toString('base64'), + mimeType: 'image/png', + fallback: '[image: png]', + }, + { + data: PNG_1X1.toString('base64'), + mimeType: 'image/jpeg', + fallback: '[image: jpeg]', + }, + ]) { + expect( + prepareInlineTerminalImage({ + data: testCase.data, + mimeType: testCase.mimeType, + contentWidth: 24, + env: { TERM: 'xterm-kitty' }, + stdoutIsTTY: true, + }), + ).toEqual({ fallbackText: testCase.fallback, result: null }); + } + }); + + it('rejects inline payloads above the shared image limit before decoding', () => { + const oversizedBase64 = 'A'.repeat(MAX_INLINE_IMAGE_ENCODED_LENGTH + 1); + + expect( + prepareInlineTerminalImage({ + data: oversizedBase64, + mimeType: 'image/png', + contentWidth: 24, + env: { TERM: 'xterm-kitty' }, + stdoutIsTTY: true, + }), + ).toEqual({ fallbackText: '[image: png]', result: null }); + }); + + it('rejects inline PNG dimensions above the shared image limit', () => { + expect( + prepareInlineTerminalImage({ + data: pngWithSize(1_000_001, 1).toString('base64'), + mimeType: 'image/png', + contentWidth: 24, + env: { TERM: 'xterm-kitty' }, + stdoutIsTTY: true, + }), + ).toEqual({ fallbackText: '[image: png]', result: null }); + }); + + it('rejects inline PNGs above the total pixel limit', () => { + const width = 8_001; + const height = Math.floor(MAX_INLINE_IMAGE_PIXELS / width) + 1; + expect( + prepareInlineTerminalImage({ + data: pngWithSize(width, height).toString('base64'), + mimeType: 'image/png', + contentWidth: 24, + env: { TERM: 'xterm-kitty' }, + stdoutIsTTY: true, + }), + ).toEqual({ fallbackText: '[image: png]', result: null }); + }); + + it('does not render inline image data when output is disabled', () => { + expect( + prepareInlineTerminalImage({ + data: PNG_1X1.toString('base64'), + mimeType: 'image/png', + contentWidth: 24, + env: { TERM: 'xterm-kitty' }, + stdoutIsTTY: true, + disabled: true, + }), + ).toEqual({ fallbackText: '[image: 1x1 png]', result: null }); + }); + it('exposes a stable render key on Kitty results so remounts can skip re-transmission', async () => { await fs.writeFile(imagePath, pngWithSize(1600, 800)); const options = { @@ -288,6 +387,38 @@ describe('terminalImageRenderer', () => { }, ); + it.runIf(process.platform !== 'win32')( + 'passes inline PNG bytes to chafa over stdin', + async () => { + const binDir = path.join(tempDir, 'inline-bin'); + await fs.mkdir(binDir); + const chafaPath = path.join(binDir, 'chafa'); + await fs.writeFile( + chafaPath, + '#!/usr/bin/env node\nconst chunks=[];process.stdin.on("data",(chunk)=>chunks.push(chunk));process.stdin.on("end",()=>{const data=Buffer.concat(chunks);process.stdout.write(`${process.argv.at(-1)}:${data.subarray(0,8).toString("hex")}:${process.env.TEST_RENDERER_SECRET ? "LEAKED" : "safe"}\\n`);});\n', + ); + await fs.chmod(chafaPath, 0o755); + + const prepared = prepareInlineTerminalImage({ + data: PNG_1X1.toString('base64'), + mimeType: 'image/png', + contentWidth: 20, + env: { + PATH: `${binDir}${path.delimiter}${process.env['PATH'] ?? ''}`, + TERM_PROGRAM: 'WarpTerminal', + TEST_RENDERER_SECRET: 'must-not-reach-chafa', + }, + stdoutIsTTY: true, + }); + + expect(prepared.result).toEqual({ + kind: 'ansi', + lines: ['-:89504e470d0a1a0a:safe'], + }); + expect(prepared.fallbackText).toBe('[image: 1x1 png]'); + }, + ); + it('rejects cmd.exe metacharacters before invoking a shell shim', async () => { const dangerousImagePath = path.join(tempDir, 'chart & whoami.png'); await fs.writeFile(dangerousImagePath, PNG_1X1); diff --git a/packages/cli/src/ui/utils/terminal-image-renderer.ts b/packages/cli/src/ui/utils/terminal-image-renderer.ts index b46fb9517ea..f6ff1962dff 100644 --- a/packages/cli/src/ui/utils/terminal-image-renderer.ts +++ b/packages/cli/src/ui/utils/terminal-image-renderer.ts @@ -21,6 +21,7 @@ import { shouldRunThroughShell, type KittyImagePlaceholder, } from './mermaidImageRenderer.js'; +import { MAX_INLINE_IMAGE_ENCODED_LENGTH } from './inline-image-parts.js'; const CHAFA_TIMEOUT_MS = 8000; const CHAFA_MAX_OUTPUT_BYTES = 2 * 1024 * 1024; @@ -29,6 +30,8 @@ const MAX_PREVIEW_ROWS = 24; const ESTIMATED_CELL_WIDTH_PX = 8; const ESTIMATED_CELL_HEIGHT_PX = 16; const MAX_REASON_CHARS = 200; +const MAX_INLINE_IMAGE_DIMENSION = 1_000_000; +export const MAX_INLINE_IMAGE_PIXELS = 64_000_000; // cmd.exe treats these as command separators / metacharacters. With shell:true // Node forwards arguments unquoted, so a model-chosen path containing any of // them is a command-injection vector through a .cmd/.bat chafa shim. @@ -43,6 +46,11 @@ const RENDER_CACHE_LIMIT = 40; const RENDER_CACHE_BYTE_LIMIT = 32 * 1024 * 1024; const renderCache = new Map(); let renderCacheBytes = 0; +const INLINE_DECODE_CACHE_LIMIT = 4; +const inlineDecodeCache = new Map< + string, + { png: Buffer; size: { width: number; height: number } } +>(); // A Kitty terminal keeps a transmitted image and redraws it from the placeholder // cells alone. The live-row -> Static-row move and every resize remount @@ -87,6 +95,21 @@ export interface TerminalImageRenderOptions { stdoutIsTTY?: boolean; } +export interface InlineTerminalImageRenderOptions { + data: string; + mimeType: string; + contentWidth: number; + availableTerminalHeight?: number; + env?: NodeJS.ProcessEnv; + stdoutIsTTY?: boolean; + disabled?: boolean; +} + +export interface PreparedInlineTerminalImage { + fallbackText: string; + result: TerminalImageRenderResult | null; +} + export function supportsKittyImageProtocol( env: NodeJS.ProcessEnv = process.env, stdoutIsTTY = process.stdout.isTTY, @@ -133,6 +156,50 @@ export function containsCmdShellMetacharacters(filePath: string): boolean { return CMD_SHELL_METACHARACTERS.test(filePath); } +export function prepareInlineTerminalImage({ + data, + mimeType, + contentWidth, + availableTerminalHeight, + env = process.env, + stdoutIsTTY = process.stdout.isTTY, + disabled = false, +}: InlineTerminalImageRenderOptions): PreparedInlineTerminalImage { + const format = getImageFormat(mimeType); + const emptyFallback = format ? `[image: ${format}]` : '[image]'; + if (format !== 'png') { + return { fallbackText: emptyFallback, result: null }; + } + + const decoded = getDecodedInlinePng(data); + if (!decoded) { + return { fallbackText: emptyFallback, result: null }; + } + const { png, size } = decoded; + + const fallbackText = `[image: ${size.width}x${size.height} png]`; + if (disabled) { + return { fallbackText, result: null }; + } + + const shape = fitImageToTerminal(size, contentWidth, availableTerminalHeight); + const useKitty = supportsKittyImageProtocol(env, stdoutIsTTY); + const chafaPath = useKitty ? null : findExecutable('chafa', env); + const cacheKey = createInlineRenderCacheKey(png, shape, useKitty, chafaPath); + const cached = getCachedRenderResult(cacheKey); + if (cached) { + return { fallbackText, result: cached }; + } + + const result: TerminalImageRenderResult = useKitty + ? { ...renderKitty(png, shape), key: cacheKey } + : renderWithChafa({ data: png }, shape, env, chafaPath); + if (result.kind !== 'unavailable') { + rememberRenderResult(cacheKey, result); + } + return { fallbackText, result }; +} + export function renderTerminalImage({ display, contentWidth, @@ -208,7 +275,7 @@ export function renderTerminalImage({ const result: TerminalImageRenderResult = useKitty ? { ...renderKitty(png, shape), key: cacheKey } - : renderWithChafa(display.filePath, shape, env, chafaPath); + : renderWithChafa({ filePath: display.filePath }, shape, env, chafaPath); if (result.kind !== 'unavailable') { rememberRenderResult(cacheKey, result); } @@ -254,6 +321,105 @@ function createRenderCacheKey( ].join('\0'); } +function createInlineRenderCacheKey( + png: Buffer, + shape: { widthCells: number; rows: number }, + useKitty: boolean, + chafaPath: string | null, +): string { + return [ + 'inline', + crypto.createHash('sha256').update(png).digest('hex'), + shape.widthCells, + shape.rows, + useKitty ? 'kitty' : (chafaPath ?? 'none'), + ].join('\0'); +} + +function getImageFormat(mimeType: string): string | null { + const match = /^image\/([a-z0-9][a-z0-9.+-]*)$/.exec( + mimeType.trim().toLowerCase(), + ); + return match?.[1] ?? null; +} + +function decodeInlineImage(data: string): Buffer | null { + if (data.length === 0 || data.length > MAX_INLINE_IMAGE_ENCODED_LENGTH) { + return null; + } + + const normalized = data.replace(/\s/g, ''); + if ( + normalized.length === 0 || + normalized.length % 4 === 1 || + !/^[A-Za-z0-9+/]*={0,2}$/.test(normalized) + ) { + return null; + } + + const decoded = Buffer.from(normalized, 'base64'); + if ( + decoded.length === 0 || + decoded.length > MAX_TERMINAL_IMAGE_BYTES || + decoded.toString('base64').replace(/=+$/, '') !== + normalized.replace(/=+$/, '') + ) { + return null; + } + return decoded; +} + +function getDecodedInlinePng( + data: string, +): { png: Buffer; size: { width: number; height: number } } | null { + const cached = inlineDecodeCache.get(data); + if (cached) { + inlineDecodeCache.delete(data); + inlineDecodeCache.set(data, cached); + return cached; + } + + const png = decodeInlineImage(data); + if (!png) return null; + const size = readValidatedInlinePngSize(png); + if (!size) return null; + + const decoded = { png, size }; + inlineDecodeCache.set(data, decoded); + while (inlineDecodeCache.size > INLINE_DECODE_CACHE_LIMIT) { + const oldest = inlineDecodeCache.keys().next().value; + if (oldest === undefined) break; + inlineDecodeCache.delete(oldest); + } + return decoded; +} + +function readValidatedInlinePngSize( + png: Buffer, +): { width: number; height: number } | null { + if ( + png.length < 24 || + png.readUInt32BE(8) !== 13 || + png.subarray(12, 16).toString('ascii') !== 'IHDR' + ) { + return null; + } + const size = readPngSize(png); + if ( + !size || + !Number.isInteger(size.width) || + !Number.isInteger(size.height) || + size.width <= 0 || + size.height <= 0 || + size.width > MAX_INLINE_IMAGE_DIMENSION || + size.height > MAX_INLINE_IMAGE_DIMENSION || + size.width * size.height > MAX_INLINE_IMAGE_PIXELS + ) { + return null; + } + return size; +} + function getCachedRenderResult( key: string, ): TerminalImageRenderResult | undefined { @@ -385,8 +551,10 @@ function renderKitty( }; } +type ChafaImageSource = { filePath: string } | { data: Buffer }; + function renderWithChafa( - filePath: string, + source: ChafaImageSource, shape: { widthCells: number; rows: number }, env: NodeJS.ProcessEnv, chafaPath: string | null, @@ -402,7 +570,11 @@ function renderWithChafa( }; } const useShell = shouldRunThroughShell(chafaPath); - if (useShell && containsCmdShellMetacharacters(filePath)) { + if ( + useShell && + 'filePath' in source && + containsCmdShellMetacharacters(source.filePath) + ) { return { kind: 'unavailable', reason: @@ -418,7 +590,7 @@ function renderWithChafa( '--format=symbols', '--symbols=block', `--size=${shape.widthCells}x${shape.rows}`, - filePath, + 'filePath' in source ? source.filePath : '-', ], { encoding: 'utf8', @@ -426,6 +598,7 @@ function renderWithChafa( shell: useShell, maxBuffer: CHAFA_MAX_OUTPUT_BYTES, timeout: CHAFA_TIMEOUT_MS, + ...('data' in source ? { input: source.data } : {}), }, ); const lines = stdout.split(/\r?\n/).filter((line) => line.length > 0); diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts index ce7f21d953e..65277a506a6 100644 --- a/packages/core/src/core/turn.test.ts +++ b/packages/core/src/core/turn.test.ts @@ -199,6 +199,97 @@ describe('Turn', () => { ]); }); + it('should preserve ordered image parts in content events', async () => { + const mockResponseStream = (async function* () { + yield { + type: StreamEventType.CHUNK, + value: { + candidates: [ + { + content: { + parts: [ + { text: 'before' }, + { + inlineData: { + data: 'aW1hZ2U=', + mimeType: 'image/png', + displayName: 'chart.png', + }, + }, + { thought: true, text: 'hidden' }, + { text: 'after' }, + ], + }, + }, + ], + } as GenerateContentResponse, + }; + yield { + type: StreamEventType.CHUNK, + value: { + candidates: [ + { + content: { + parts: [ + { + inlineData: { + data: 'c2Vjb25k', + mimeType: 'image/webp', + }, + }, + ], + }, + }, + ], + } as GenerateContentResponse, + }; + })(); + mockSendMessageStream.mockResolvedValue(mockResponseStream); + + const events = []; + for await (const event of turn.run( + 'test-model', + [{ text: 'Hi' }], + new AbortController().signal, + )) { + events.push(event); + } + + expect(events).toEqual([ + { + type: GeminiEventType.Thought, + value: { subject: '', description: 'hidden' }, + }, + { + type: GeminiEventType.Content, + value: 'beforeafter', + parts: [ + { text: 'before' }, + { + inlineData: { + data: 'aW1hZ2U=', + mimeType: 'image/png', + displayName: 'chart.png', + }, + }, + { text: 'after' }, + ], + }, + { + type: GeminiEventType.Content, + value: '', + parts: [ + { + inlineData: { + data: 'c2Vjb25k', + mimeType: 'image/webp', + }, + }, + ], + }, + ]); + }); + it('should emit Thought events when a thought part is present', async () => { const mockResponseStream = (async function* () { yield { diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 4439846df03..88ef3663de2 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -285,9 +285,21 @@ export interface ServerToolCallConfirmationDetails { details: ToolCallConfirmationDetails; } +export type ServerGeminiContentPart = + | { text: string } + | { + inlineData: { + data: string; + mimeType: string; + displayName?: string; + }; + }; + export type ServerGeminiContentEvent = { type: GeminiEventType.Content; value: string; + /** Ordered display parts, present only when the chunk contains an image. */ + parts?: ServerGeminiContentPart[]; }; export type ServerGeminiThoughtEvent = { @@ -458,6 +470,40 @@ export type ServerGeminiStreamEvent = | ServerGeminiSessionTokenLimitExceededEvent | ServerGeminiRetryEvent; +function getDisplayContentParts( + response: GenerateContentResponse, +): ServerGeminiContentPart[] { + const parts = response.candidates?.[0]?.content?.parts ?? []; + const displayParts: ServerGeminiContentPart[] = []; + + for (const part of parts) { + if (part.thought) { + continue; + } + if (typeof part.text === 'string' && part.text.length > 0) { + displayParts.push({ text: part.text }); + } + const inlineData = part.inlineData; + if ( + inlineData?.mimeType?.trim().toLowerCase().startsWith('image/') && + typeof inlineData.data === 'string' && + inlineData.data.length > 0 + ) { + displayParts.push({ + inlineData: { + data: inlineData.data, + mimeType: inlineData.mimeType, + ...(typeof inlineData.displayName === 'string' + ? { displayName: inlineData.displayName } + : {}), + }, + }); + } + } + + return displayParts; +} + // A turn manages the agentic loop turn within the server context. export class Turn { readonly pendingToolCalls: ToolCallRequestInfo[] = []; @@ -563,9 +609,15 @@ export class Turn { }; } - const text = getResponseText(resp); - if (text) { - yield { type: GeminiEventType.Content, value: text }; + const text = getResponseText(resp) ?? ''; + const displayParts = getDisplayContentParts(resp); + const hasImage = displayParts.some((part) => 'inlineData' in part); + if (text || hasImage) { + yield { + type: GeminiEventType.Content, + value: text, + ...(hasImage ? { parts: displayParts } : {}), + }; } // Handle function calls (requesting tool execution)