From e4a87fa3ce7b1aec9dff143d945523333a6a2705 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 17 Jun 2026 13:19:28 +0200 Subject: [PATCH] feat(coding-agent): added matplotlib rendering and image persistence across session reload - Added Matplotlib figure PNG rendering and display tracking in Python runner to emit PNG output immediately when figures are displayed via display(fig). - Extended session persistence to externalize oversized image payloads in both content and details.images, enabling tool result images to survive session reload. - Enhanced session loader to resolve image data payloads and blob references across content and details.images during session reconstruction. - Added image cache invalidation in TUI image component when image protocol, cell dimensions, or Kitty Unicode placeholder mode changes. - Added comprehensive test coverage for Matplotlib display, image persistence across reload, and TUI image rendering with protocol and dimension changes. --- packages/coding-agent/CHANGELOG.md | 5 + packages/coding-agent/src/eval/py/runner.py | 44 ++++ .../src/session/session-loader.ts | 53 ++--- .../src/session/session-persistence.ts | 38 +++- .../core/python-runner.integration.test.ts | 71 +++++++ .../utils/render-initial-messages.test.ts | 199 +++++++++++++++++- .../signature-persistence.test.ts | 67 +++++- .../test/session-persistence-images.test.ts | 71 +++++++ packages/tui/CHANGELOG.md | 5 +- packages/tui/src/components/image.ts | 25 ++- packages/tui/test/image-budget.test.ts | 95 +++++++++ packages/tui/test/image-render.test.ts | 109 +++++++++- 12 files changed, 723 insertions(+), 59 deletions(-) create mode 100644 packages/coding-agent/test/session-persistence-images.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 2adc1f7cb..0ef4ef32c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -17,6 +17,9 @@ ### Fixed +- Fixed Matplotlib figure display to emit PNG output immediately when `display(fig)` is called, even if the figure is closed before the end-of-cell flush +- Fixed persisted tool-result image payloads in `details.images` to externalize and resolve through the session blob store, so generated-image details survive resume without stale blob refs or truncation +- Fixed duplicate Matplotlib image output by skipping the automatic end-of-request figure flush for figures that were already displayed through `display(fig)` - Fixed google-antigravity image generation and web search requests to fail over to the alternate antigravity endpoint on 429/server/network failures instead of stopping at the first endpoint - Fixed context usage breakdown to use a completed assistant usage anchor from the current turn instead of a pending prompt snapshot so totals no longer overcount when a large in-turn tool step returns usage - Fixed side-channel turns and advisor requests to keep using credential resolvers during retries, so Google `Resource exhausted` 429s can rotate to the next account instead of surfacing a terminal error banner @@ -28,6 +31,8 @@ - Fixed the ssh tool rejecting valid Windows identity files before invoking OpenSSH by skipping Unix mode-bit key validation on native Windows ([#2850](https://github.com/can1357/oh-my-pi/issues/2850)). - Fixed `web_search`/`omp q` aborting before any provider ran when the global Settings singleton was not initialized; `executeSearch` now reads `providers.antigravityEndpoint` once and tolerates an uninitialized settings store instead of throwing - Fixed the new `git.enabled` and `images.describeForTextModels` settings declaring section groups (`Git`, `Vision`) that were not registered in `TAB_GROUPS`, so they now render in their intended settings-panel sections +- Fixed Python `display(fig)` for Matplotlib figures to emit PNG output immediately, even when user code closes the figure before the end-of-cell flush. +- Fixed persisted tool-result image payloads stored in `details.images` to externalize and resolve through the session blob store, so generated-image details survive resume without stale blob refs or truncation. ## [16.0.4] - 2026-06-17 diff --git a/packages/coding-agent/src/eval/py/runner.py b/packages/coding-agent/src/eval/py/runner.py index 253c57c46..9db1513e5 100644 --- a/packages/coding-agent/src/eval/py/runner.py +++ b/packages/coding-agent/src/eval/py/runner.py @@ -190,6 +190,11 @@ class _RunnerState: _CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar("omp_current_rid", default=None) +_CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS: contextvars.ContextVar[set[int] | None] = contextvars.ContextVar( + "omp_displayed_matplotlib_figure_ids", + default=None, +) + _STATE = _RunnerState() @@ -670,6 +675,36 @@ _REPR_MIMES = [ ] +def _is_matplotlib_figure(value: Any) -> bool: + figure_module = sys.modules.get("matplotlib.figure") + figure_cls = getattr(figure_module, "Figure", None) + if isinstance(figure_cls, type) and isinstance(value, figure_cls): + return True + + value_type = type(value) + return value_type.__module__ == "matplotlib.figure" and value_type.__name__ == "Figure" + + +def _matplotlib_figure_png(value: Any) -> str | None: + if not _is_matplotlib_figure(value): + return None + + savefig = getattr(value, "savefig", None) + if not callable(savefig): + return None + + try: + buf = io.BytesIO() + savefig(buf, format="png", bbox_inches="tight") + except Exception: + return None + + displayed_ids = _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.get() + if displayed_ids is not None: + displayed_ids.add(id(value)) + return base64.b64encode(buf.getvalue()).decode("ascii") + + def _coerce_image_bytes(value: Any) -> str: if isinstance(value, (bytes, bytearray)): return base64.b64encode(bytes(value)).decode("ascii") @@ -685,6 +720,10 @@ def _mime_bundle(value: Any) -> dict: accessors, and always provides ``text/plain``. """ bundle: dict[str, Any] = {} + matplotlib_png = _matplotlib_figure_png(value) + if matplotlib_png is not None: + bundle["image/png"] = matplotlib_png + mimebundle = getattr(value, "_repr_mimebundle_", None) if callable(mimebundle): @@ -758,6 +797,9 @@ def _flush_matplotlib_figures() -> None: for num in fignums: try: fig = plt.figure(num) + if id(fig) in (_CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.get() or set()): + plt.close(fig) + continue buf = io.BytesIO() fig.savefig(buf, format="png", bbox_inches="tight") data = base64.b64encode(buf.getvalue()).decode("ascii") @@ -978,6 +1020,7 @@ def _start_parent_watchdog() -> None: async def _handle_request_async(req: dict) -> None: rid = str(req.get("id")) token = _CURRENT_RID.set(rid) + displayed_matplotlib_token = _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.set(set()) _STATE.capture_rid = rid _STATE.user_ns["__omp_run_id__"] = rid _STATE.cancel_requested = False @@ -1046,6 +1089,7 @@ async def _handle_request_async(req: dict) -> None: _STATE.capture_rid = None _flush_stream_proxies(rid) _CURRENT_RID.reset(token) + _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.reset(displayed_matplotlib_token) def _emit_error(rid: str, exc: BaseException) -> None: diff --git a/packages/coding-agent/src/session/session-loader.ts b/packages/coding-agent/src/session/session-loader.ts index 702eb16cf..d75755007 100644 --- a/packages/coding-agent/src/session/session-loader.ts +++ b/packages/coding-agent/src/session/session-loader.ts @@ -4,7 +4,7 @@ import { BlobStore, isBlobRef, resolveImageData, resolveImageDataUrl } from "./b import { buildSessionContext } from "./session-context"; import type { FileEntry, SessionEntry, SessionHeader } from "./session-entries"; import { migrateToCurrentVersion } from "./session-migrations"; -import { isImageBlock } from "./session-persistence"; +import { isImageBlock, isImageDataPayload } from "./session-persistence"; import { FileSessionStorage, type SessionStorage } from "./session-storage"; /** Exported for compaction.test.ts */ @@ -44,9 +44,21 @@ function hasImageUrl(value: unknown): value is { image_url: string } { return typeof value === "object" && value !== null && "image_url" in value && typeof value.image_url === "string"; } -async function resolvePersistedImageUrlRefs(value: unknown, blobStore: BlobStore): Promise { +function shouldResolveImagePayload(value: unknown, key: string | undefined): value is { data: string } { + if (!isImageDataPayload(value) || !isBlobRef(value.data)) return false; + return (key === "content" && isImageBlock(value)) || key === "images"; +} + +async function resolvePersistedBlobRefs(value: unknown, blobStore: BlobStore, key?: string): Promise { + if (shouldResolveImagePayload(value, key)) { + value.data = await resolveImageData(blobStore, value.data); + return; + } + if (Array.isArray(value)) { - await Promise.all(value.map(item => resolvePersistedImageUrlRefs(item, blobStore))); + await Promise.all( + value.map(item => resolvePersistedBlobRefs(item, blobStore, key)), + ); return; } @@ -56,38 +68,15 @@ async function resolvePersistedImageUrlRefs(value: unknown, blobStore: BlobStore value.image_url = await resolveImageDataUrl(blobStore, value.image_url); } - await Promise.all(Object.values(value).map(item => resolvePersistedImageUrlRefs(item, blobStore))); + await Promise.all( + Object.entries(value).map(([childKey, item]) => resolvePersistedBlobRefs(item, blobStore, childKey)), + ); } export async function resolveBlobRefsInEntries(entries: FileEntry[], blobStore: BlobStore): Promise { - const promises: Promise[] = []; - - for (const entry of entries) { - if (entry.type === "session") continue; - - let contentArray: unknown[] | undefined; - if (entry.type === "message" && "content" in entry.message && Array.isArray(entry.message.content)) { - contentArray = entry.message.content; - } else if (entry.type === "custom_message" && Array.isArray(entry.content)) { - contentArray = entry.content; - } - - if (contentArray) { - for (const block of contentArray) { - if (isImageBlock(block) && isBlobRef(block.data)) { - promises.push( - resolveImageData(blobStore, block.data).then(resolved => { - block.data = resolved; - }), - ); - } - } - } - - promises.push(resolvePersistedImageUrlRefs(entry, blobStore)); - } - - await Promise.all(promises); + await Promise.all( + entries.filter(entry => entry.type !== "session").map(entry => resolvePersistedBlobRefs(entry, blobStore)), + ); } /** diff --git a/packages/coding-agent/src/session/session-persistence.ts b/packages/coding-agent/src/session/session-persistence.ts index c274c083d..3a20b829a 100644 --- a/packages/coding-agent/src/session/session-persistence.ts +++ b/packages/coding-agent/src/session/session-persistence.ts @@ -36,10 +36,33 @@ export function isImageBlock(value: unknown): value is { type: "image"; data: st ); } +function isImageMimeType(value: unknown): value is string { + return typeof value === "string" && value.toLowerCase().startsWith("image/"); +} + +export function isImageDataPayload(value: unknown): value is { data: string; mimeType?: string } { + return ( + typeof value === "object" && + value !== null && + "data" in value && + typeof (value as { data?: string }).data === "string" && + (isImageBlock(value) || ("mimeType" in value && isImageMimeType((value as { mimeType?: unknown }).mimeType))) + ); +} + +function shouldExternalizeImagePayload( + value: unknown, + key: string | undefined, +): value is { data: string; mimeType?: string } { + if (!isImageDataPayload(value)) return false; + if (isBlobRef(value.data) || value.data.length < BLOB_EXTERNALIZE_THRESHOLD) return false; + return (key === TEXT_CONTENT_KEY && isImageBlock(value)) || key === "images"; +} + /** * Recursively truncate large strings in an object for session persistence. * - Truncates any oversized string fields (key-agnostic) - * - Replaces oversized image blocks with text notices + * - Externalizes oversized image payloads to blob refs * - Updates lineCount when content is truncated * - Returns original object if no changes needed (structural sharing) * @@ -50,6 +73,9 @@ export function isImageBlock(value: unknown): value is { type: "image"; data: st */ function truncateForPersistence(obj: unknown, blobStore: BlobStore, key?: string): unknown { if (obj === null || obj === undefined) return obj; + if (shouldExternalizeImagePayload(obj, key)) { + return { ...obj, data: externalizeImageDataSync(blobStore, obj.data, obj.mimeType) }; + } if (typeof obj === "string") { if (key === "image_url" && isImageDataUrl(obj)) { @@ -72,16 +98,6 @@ function truncateForPersistence(obj: unknown, blobStore: BlobStore, key?: string const result: unknown[] = new Array(obj.length); for (let i = 0; i < obj.length; i++) { const item = obj[i]; - if ( - key === TEXT_CONTENT_KEY && - isImageBlock(item) && - !isBlobRef(item.data) && - item.data.length >= BLOB_EXTERNALIZE_THRESHOLD - ) { - changed = true; - result[i] = { ...item, data: externalizeImageDataSync(blobStore, item.data, item.mimeType) }; - continue; - } const newItem = truncateForPersistence(item, blobStore, key); if (newItem !== item) changed = true; result[i] = newItem; diff --git a/packages/coding-agent/test/core/python-runner.integration.test.ts b/packages/coding-agent/test/core/python-runner.integration.test.ts index fb82cfc33..182379494 100644 --- a/packages/coding-agent/test/core/python-runner.integration.test.ts +++ b/packages/coding-agent/test/core/python-runner.integration.test.ts @@ -6,11 +6,37 @@ */ import { afterEach, describe, expect, it } from "bun:test"; import * as path from "node:path"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { disposeAllKernelSessions, executePythonWithKernel } from "@oh-my-pi/pi-coding-agent/eval/py/executor"; import { PythonKernel } from "@oh-my-pi/pi-coding-agent/eval/py/kernel"; +import { filterEnv, resolvePythonRuntime } from "@oh-my-pi/pi-coding-agent/eval/py/runtime"; import { TempDir } from "@oh-my-pi/pi-utils"; const SHOULD_RUN = Bun.env.PI_PYTHON_INTEGRATION === "1"; +const MATPLOTLIB_TEST_CWD = process.cwd(); + +async function hasMatplotlib(cwd: string): Promise { + if (!SHOULD_RUN) return false; + try { + const { env } = (await Settings.init()).getShellConfig(); + const runtime = resolvePythonRuntime(cwd, filterEnv(env)); + const spawnEnv: Record = {}; + for (const [key, value] of Object.entries(runtime.env)) { + if (typeof value === "string") spawnEnv[key] = value; + } + const result = Bun.spawnSync([runtime.pythonPath, "-c", "import matplotlib"], { + cwd, + env: spawnEnv, + stdout: "ignore", + stderr: "ignore", + }); + return result.exitCode === 0; + } catch { + return false; + } +} + +const HAS_MATPLOTLIB = await hasMatplotlib(MATPLOTLIB_TEST_CWD); describe.skipIf(!SHOULD_RUN)("python runner subprocess", () => { afterEach(async () => { @@ -108,6 +134,51 @@ describe.skipIf(!SHOULD_RUN)("python runner subprocess", () => { } }); + it.skipIf(!HAS_MATPLOTLIB)("captures display(fig) as a PNG before the figure is closed", async () => { + const kernel = await PythonKernel.start({ cwd: MATPLOTLIB_TEST_CWD }); + try { + const result = await executePythonWithKernel( + kernel, + [ + "import matplotlib.pyplot as plt", + "fig, ax = plt.subplots()", + "ax.plot([0, 1], [0, 1])", + "display(fig)", + "plt.close(fig)", + ].join("\n"), + ); + + expect(result.exitCode).toBe(0); + const images = result.displayOutputs.filter(output => output.type === "image"); + expect(images).toHaveLength(1); + expect(images[0]).toMatchObject({ mimeType: "image/png" }); + expect(images[0]?.data).not.toContain("blob:"); + expect(result.output).toContain(" { + const kernel = await PythonKernel.start({ cwd: MATPLOTLIB_TEST_CWD }); + try { + const result = await executePythonWithKernel( + kernel, + [ + "import matplotlib.pyplot as plt", + "fig, ax = plt.subplots()", + "ax.plot([0, 1], [1, 0])", + "display(fig)", + ].join("\n"), + ); + + expect(result.exitCode).toBe(0); + expect(result.displayOutputs.filter(output => output.type === "image")).toHaveLength(1); + } finally { + await kernel.shutdown(); + } + }); + it("translates %pwd magic to the user namespace", async () => { using tempDir = TempDir.createSync("@python-runner-magic-"); const kernel = await PythonKernel.start({ cwd: tempDir.path() }); diff --git a/packages/coding-agent/test/modes/utils/render-initial-messages.test.ts b/packages/coding-agent/test/modes/utils/render-initial-messages.test.ts index d0267a175..a0e356560 100644 --- a/packages/coding-agent/test/modes/utils/render-initial-messages.test.ts +++ b/packages/coding-agent/test/modes/utils/render-initial-messages.test.ts @@ -12,16 +12,30 @@ * scrollback-clearing repaint (`clearTerminalHistory`). */ -import { beforeAll, describe, expect, it, type Mock, vi } from "bun:test"; +import { afterEach, beforeAll, describe, expect, it, type Mock, vi } from "bun:test"; +import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import type { AssistantMessage, ImageContent, Usage } from "@oh-my-pi/pi-ai"; +import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; import { UiHelpers } from "@oh-my-pi/pi-coding-agent/modes/utils/ui-helpers"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import type { SessionContext } from "@oh-my-pi/pi-coding-agent/session/session-context"; +import { TempDir } from "@oh-my-pi/pi-utils"; +import { type Component, Container, Image, ImageProtocol, setTerminalImageProtocol, TERMINAL } from "@oh-my-pi/pi-tui"; beforeAll(() => { initTheme(); }); +const originalImageProtocol = TERMINAL.imageProtocol; + +afterEach(() => { + resetSettingsForTest(); + setTerminalImageProtocol(originalImageProtocol); + vi.restoreAllMocks(); +}); + function makeEmptyContext(): SessionContext { return { messages: [], @@ -74,6 +88,97 @@ function makeCtx(): { return { ctx, transcriptSpy, llmContextSpy, renderSessionContextSpy }; } +const emptyUsage: Usage = { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, +}; + +const pngImage: ImageContent = { + type: "image", + data: "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==", + mimeType: "image/png", +}; + +function assistantToolCall(id: string, name: string, args: Record): AssistantMessage { + return { + role: "assistant", + content: [{ type: "toolCall", id, name, arguments: args }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet", + usage: emptyUsage, + stopReason: "toolUse", + timestamp: 1, + }; +} + +function transcriptWith(messages: AgentMessage[]): SessionContext { + return { ...makeEmptyContext(), messages }; +} + +function countImageComponents(component: Component): number { + const own = component instanceof Image ? 1 : 0; + const children = (component as { children?: unknown }).children; + if (!Array.isArray(children)) return own; + return own + children.reduce((count, child) => count + countImageComponents(child as Component), 0); +} + +function hasImageComponent(component: Component): boolean { + return countImageComponents(component) > 0; +} + +function makeRenderCtx(transcript: SessionContext): { ctx: InteractiveModeContext; chatContainer: Container } { + const chatContainer = new Container(); + let helpers: UiHelpers; + const ctx = { + chatContainer, + pendingMessagesContainer: new Container(), + pendingBashComponents: [], + pendingPythonComponents: [], + pendingTools: new Map(), + statusLine: { invalidate: vi.fn() }, + updateEditorBorderColor: vi.fn(), + updateEditorTopBorder: vi.fn(), + ui: { requestRender: vi.fn(), imageBudget: undefined }, + resetTranscript: () => chatContainer.clear(), + settings: { get: () => false }, + toolOutputExpanded: false, + hideThinkingBlock: false, + focusedAgentId: undefined, + editor: { addToHistory: vi.fn() }, + viewSession: { + buildTranscriptSessionContext: () => transcript, + getToolByName: () => undefined, + extensionRunner: undefined, + sessionManager: { + getEntries: vi.fn(() => []), + getCwd: vi.fn(() => "/tmp"), + }, + }, + sessionManager: { + getEntries: vi.fn(() => []), + getCwd: vi.fn(() => "/tmp"), + putBlobSync: vi.fn(() => ({ + hash: "hash", + path: "/tmp/hash", + displayPath: "/tmp/hash.png", + ref: "blob:sha256:hash", + })), + }, + addMessageToChat: (message: AgentMessage, options?: { populateHistory?: boolean }) => + helpers.addMessageToChat(message, options), + renderSessionContext: (context: SessionContext, options?: { updateFooter?: boolean; populateHistory?: boolean }) => + helpers.renderSessionContext(context, options), + showStatus: vi.fn(), + } as unknown as InteractiveModeContext; + helpers = new UiHelpers(ctx); + return { ctx, chatContainer }; +} + describe("UiHelpers.renderInitialMessages — transcript source", () => { it("renders the display transcript, never the LLM context", () => { const { ctx, transcriptSpy, llmContextSpy, renderSessionContextSpy } = makeCtx(); @@ -107,3 +212,95 @@ describe("UiHelpers.renderInitialMessages — clearTerminalHistory", () => { expect(clearedCall).toBeUndefined(); }); }); + +describe("UiHelpers.renderInitialMessages — image replay", () => { + it("restores read tool image blocks onto the rebuilt assistant transcript", async () => { + await Settings.init({ inMemory: true, overrides: { "terminal.showImages": true } }); + setTerminalImageProtocol(ImageProtocol.Sixel); + const transcript = transcriptWith([ + assistantToolCall("read-image", "read", { path: "sample.png" }), + { + role: "toolResult", + toolCallId: "read-image", + toolName: "read", + content: [{ type: "text", text: "Read image: sample.png" }, pngImage], + isError: false, + timestamp: 2, + }, + ]); + const { ctx, chatContainer } = makeRenderCtx(transcript); + + new UiHelpers(ctx).renderInitialMessages(); + + expect(hasImageComponent(chatContainer)).toBe(true); + expect(Bun.stripANSI(chatContainer.render(100).join("\n"))).toContain("Read sample.png"); + }); + + it("restores eval display image blocks onto rebuilt tool output", async () => { + await Settings.init({ inMemory: true, overrides: { "terminal.showImages": true } }); + setTerminalImageProtocol(ImageProtocol.Sixel); + const transcript = transcriptWith([ + assistantToolCall("eval-image", "eval", { cells: [{ language: "py", code: "display(image)" }] }), + { + role: "toolResult", + toolCallId: "eval-image", + toolName: "eval", + content: [{ type: "text", text: "(displayed 1 image; no text output)" }, pngImage], + details: { + language: "python", + cells: [{ index: 0, code: "display(image)", output: "display image 1: 1x1", status: "complete" }], + }, + isError: false, + timestamp: 2, + }, + ]); + + const { ctx, chatContainer } = makeRenderCtx(transcript); + + new UiHelpers(ctx).renderInitialMessages(); + + expect(hasImageComponent(chatContainer)).toBe(true); + expect(Bun.stripANSI(chatContainer.render(100).join("\n"))).toContain("display image 1: 1x1"); + }); + + it("replays reopened session image blocks through the cold-start rebuild path", async () => { + await Settings.init({ inMemory: true, overrides: { "terminal.showImages": true } }); + setTerminalImageProtocol(ImageProtocol.Sixel); + using tempDir = TempDir.createSync("@pi-render-initial-image-replay-"); + const session = SessionManager.create(tempDir.path(), tempDir.path()); + session.appendMessage(assistantToolCall("read-reopened", "read", { path: "reopened.png" })); + session.appendMessage({ + role: "toolResult", + toolCallId: "read-reopened", + toolName: "read", + content: [{ type: "text", text: "Read image: reopened.png" }, pngImage], + isError: false, + timestamp: 2, + }); + session.appendMessage(assistantToolCall("eval-reopened", "eval", { cells: [{ language: "py", code: "display(image)" }] })); + session.appendMessage({ + role: "toolResult", + toolCallId: "eval-reopened", + toolName: "eval", + content: [{ type: "text", text: "(displayed 1 image; no text output)" }, pngImage], + details: { + language: "python", + cells: [{ index: 0, code: "display(image)", output: "display image 1: 1x1", status: "complete" }], + }, + isError: false, + timestamp: 4, + }); + await session.flush(); + const sessionFile = session.getSessionFile(); + if (!sessionFile) throw new Error("Expected persisted session file"); + const reloaded = await SessionManager.open(sessionFile); + const transcript = reloaded.buildSessionContext({ transcript: true }); + const { ctx, chatContainer } = makeRenderCtx(transcript); + + new UiHelpers(ctx).renderInitialMessages({ clearTerminalHistory: true }); + + expect(countImageComponents(chatContainer)).toBe(2); + expect(Bun.stripANSI(chatContainer.render(100).join("\n"))).toContain("Read reopened.png"); + expect(ctx.ui.requestRender).toHaveBeenCalledWith(true, { clearScrollback: true }); + }); +}); diff --git a/packages/coding-agent/test/session-manager/signature-persistence.test.ts b/packages/coding-agent/test/session-manager/signature-persistence.test.ts index 5494e6080..09a0fe3e9 100644 --- a/packages/coding-agent/test/session-manager/signature-persistence.test.ts +++ b/packages/coding-agent/test/session-manager/signature-persistence.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from "bun:test"; import * as fs from "node:fs/promises"; import * as path from "node:path"; -import type { AssistantMessage } from "@oh-my-pi/pi-ai"; +import type { AssistantMessage, ImageContent } from "@oh-my-pi/pi-ai"; import type { SessionMessageEntry } from "@oh-my-pi/pi-coding-agent/session/session-entries"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { getBlobsDir, TempDir } from "@oh-my-pi/pi-utils"; @@ -134,6 +134,71 @@ describe("SessionManager signature persistence", () => { }); }); + it("externalizes and restores tool result image blocks across reload", async () => { + using tempDir = TempDir.createSync("@pi-session-tool-image-persistence-"); + const session = SessionManager.create(tempDir.path(), tempDir.path()); + const contentImage: ImageContent = { + type: "image", + data: Buffer.from("read-image-payload".repeat(100)).toString("base64"), + mimeType: "image/png", + }; + const detailImage: ImageContent = { + type: "image", + data: Buffer.from("eval-detail-image-payload".repeat(100)).toString("base64"), + mimeType: "image/png", + }; + + session.appendMessage({ + role: "assistant", + content: [{ type: "toolCall", id: "tool_image", name: "eval", arguments: {} }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet", + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "toolUse", + timestamp: 1, + } satisfies AssistantMessage); + session.appendMessage({ + role: "toolResult", + toolCallId: "tool_image", + toolName: "eval", + content: [{ type: "text", text: "displayed image" }, contentImage], + details: { images: [detailImage] }, + isError: false, + timestamp: 2, + }); + await session.flush(); + + const sessionFile = session.getSessionFile(); + if (!sessionFile) throw new Error("Expected persisted session file"); + const rawSession = await fs.readFile(sessionFile, "utf8"); + expect(rawSession).not.toContain(contentImage.data); + expect(rawSession).not.toContain(detailImage.data); + + const contentHash = new Bun.SHA256().update(Buffer.from(contentImage.data, "base64")).digest("hex"); + const detailHash = new Bun.SHA256().update(Buffer.from(detailImage.data, "base64")).digest("hex"); + await expect(fs.readFile(path.join(getBlobsDir(), contentHash))).resolves.toBeDefined(); + await expect(fs.readFile(path.join(getBlobsDir(), detailHash))).resolves.toBeDefined(); + + const reloaded = await SessionManager.open(sessionFile); + const reloadedToolEntry = reloaded + .getEntries() + .find(entry => entry.type === "message" && entry.message.role === "toolResult"); + if (reloadedToolEntry?.type !== "message" || reloadedToolEntry.message.role !== "toolResult") { + throw new Error("Expected tool result message"); + } + + expect(reloadedToolEntry.message.content).toEqual([{ type: "text", text: "displayed image" }, contentImage]); + expect((reloadedToolEntry.message.details as { images?: ImageContent[] }).images).toEqual([detailImage]); + }); + it("rehydrates assistant replay metadata in memory without rewriting the session file", async () => { using tempDir = TempDir.createSync("@pi-session-rehydrate-persistence-"); const session = SessionManager.create(tempDir.path(), tempDir.path()); diff --git a/packages/coding-agent/test/session-persistence-images.test.ts b/packages/coding-agent/test/session-persistence-images.test.ts new file mode 100644 index 000000000..60b123b89 --- /dev/null +++ b/packages/coding-agent/test/session-persistence-images.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, it } from "bun:test"; +import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import type { ImageContent, TextContent } from "@oh-my-pi/pi-ai"; +import { BlobStore, isBlobRef } from "@oh-my-pi/pi-coding-agent/session/blob-store"; +import type { FileEntry, SessionMessageEntry } from "@oh-my-pi/pi-coding-agent/session/session-entries"; +import { resolveBlobRefsInEntries } from "@oh-my-pi/pi-coding-agent/session/session-loader"; +import { prepareEntryForPersistence } from "@oh-my-pi/pi-coding-agent/session/session-persistence"; +import { TempDir } from "@oh-my-pi/pi-utils"; + +type ImagePayload = { data: string; mimeType: string; type?: "image" }; +type ToolResultMessage = Extract; +type ToolResultEntry = Omit & { message: ToolResultMessage }; + +const text = (value: string): TextContent => ({ type: "text", text: value }); +const png = (data: string): ImageContent => ({ type: "image", data, mimeType: "image/png" }); +const payload = (data: string): ImagePayload => ({ data, mimeType: "image/png" }); + +function messageEntry(message: ToolResultMessage): ToolResultEntry { + return { + type: "message", + id: "entry-1", + parentId: null, + timestamp: new Date(0).toISOString(), + message, + }; +} + +describe("session image persistence", () => { + it("externalizes and resolves content images and tool detail image payloads", async () => { + using tempDir = TempDir.createSync("@session-image-persistence-"); + const blobStore = new BlobStore(tempDir.path()); + const contentImageData = Buffer.alloc(1500, 1).toString("base64"); + const generatedImageData = Buffer.alloc(1500, 2).toString("base64"); + const typedDetailImageData = Buffer.alloc(1500, 3).toString("base64"); + + const original = messageEntry({ + role: "toolResult", + toolCallId: "tc1", + toolName: "generate_image", + content: [text("generated"), png(contentImageData)], + details: { + images: [payload(generatedImageData), png(typedDetailImageData)], + }, + isError: false, + timestamp: Date.now(), + }); + + const persisted = prepareEntryForPersistence(original, blobStore) as ToolResultEntry; + const persistedContentImage = persisted.message.content.find( + (block): block is ImageContent => block.type === "image", + ); + const persistedDetails = persisted.message.details as { images: ImagePayload[] }; + + expect(persistedContentImage).toBeDefined(); + expect(isBlobRef(persistedContentImage?.data ?? "")).toBe(true); + expect(persistedDetails.images).toHaveLength(2); + expect(persistedDetails.images.every(image => isBlobRef(image.data))).toBe(true); + + const loaded: FileEntry[] = [structuredClone(persisted)]; + await resolveBlobRefsInEntries(loaded, blobStore); + const resolved = loaded[0] as ToolResultEntry; + const resolvedContentImage = resolved.message.content.find( + (block): block is ImageContent => block.type === "image", + ); + const resolvedDetails = resolved.message.details as { images: ImagePayload[] }; + + expect(resolvedContentImage?.data).toBe(contentImageData); + expect(resolvedDetails.images[0]?.data).toBe(generatedImageData); + expect(resolvedDetails.images[1]?.data).toBe(typedDetailImageData); + }); +}); diff --git a/packages/tui/CHANGELOG.md b/packages/tui/CHANGELOG.md index ab2be1293..8d31c5756 100644 --- a/packages/tui/CHANGELOG.md +++ b/packages/tui/CHANGELOG.md @@ -1,14 +1,15 @@ # Changelog ## [Unreleased] - ### Changed - Coalesced byte-adjacent SGR sequences in emitted lines into a single `CSI … m`. The component tree styles each span as `text`, so adjacent spans emit runs of back-to-back SGR sequences (e.g. a `CSI 39 m` fg-reset immediately followed by the next span's `CSI 38;2;r;g;b m`); merging the run is behavior-preserving because SGR parameters apply left-to-right regardless of framing. On a real transcript this drops ~30-40% of all SGR sequences, cutting the per-frame byte volume and SGR-dispatch count a slow terminal engine (e.g. xterm.js/WebGL under a large viewport) must process. Each emitted sequence is capped at 16 parameter tokens so a long adjacent run is split across several valid CSIs instead of overflowing a terminal's parameter buffer (xterm.js caps at 32 and silently truncates, corrupting colors). A run is never extended past a parameter list that ends in an incomplete semicolon-form extended color (`38/48/58;2` missing a channel or `;5` missing the index), so a following code can't be absorbed as the missing component. Disable with `PI_NO_SGR_COALESCE=1`. ### Fixed +- Fixed image cache invalidation when terminal image protocol, Kitty placeholder mode, or cell dimensions change, preventing stale rendered output - Fixed direct inline-image placements leaving the cursor inside the reserved image block, which let following chat rows overwrite the middle of rendered screenshots ([#2863](https://github.com/can1357/oh-my-pi/issues/2863)). +- Fixed inline-image replay after startup or resume fallback paints by invalidating cached image rows when the terminal image protocol, Kitty placeholder mode, or cell dimensions change. ## [16.0.3] - 2026-06-16 @@ -1699,4 +1700,4 @@ Initial release under @oh-my-pi scope. See previous releases at [badlogic/pi-mon ### Fixed -- **Readline-style Ctrl+W**: Now skips trailing whitespace before deleting the preceding word, matching standard readline behavior. ([#306](https://github.com/badlogic/pi-mono/pull/306) by [@kim0](https://github.com/kim0)) +- **Readline-style Ctrl+W**: Now skips trailing whitespace before deleting the preceding word, matching standard readline behavior. ([#306](https://github.com/badlogic/pi-mono/pull/306) by [@kim0](https://github.com/kim0)) \ No newline at end of file diff --git a/packages/tui/src/components/image.ts b/packages/tui/src/components/image.ts index 15e39ca6c..b9de70a02 100644 --- a/packages/tui/src/components/image.ts +++ b/packages/tui/src/components/image.ts @@ -1,4 +1,6 @@ +import { getKittyGraphics } from "../kitty-graphics"; import { + getCellDimensions, getImageDimensions, type ImageDimensions, imageFallback, @@ -284,6 +286,10 @@ export class Image implements Component { #cachedLines?: string[]; #cachedWidth?: number; #cachedSuppressed = false; + #cachedImageProtocol: typeof TERMINAL.imageProtocol = null; + #cachedCellWidthPx = 0; + #cachedCellHeightPx = 0; + #cachedKittyUnicodePlaceholders = false; // Tallest graphic placement this image has rendered. The text fallback // pads itself to this height so a budget demotion never shrinks the block // (its rows may already be committed to native scrollback). @@ -311,14 +317,25 @@ export class Image implements Component { } render(width: number): readonly string[] { - const hasProtocol = TERMINAL.imageProtocol != null; + const imageProtocol = TERMINAL.imageProtocol; + const hasProtocol = imageProtocol != null; + const cellDimensions = getCellDimensions(); + const kittyUnicodePlaceholders = getKittyGraphics().unicodePlaceholders; // observe() must run on every pass — even a cache hit — so the image keeps // its display-order slot in the budget. Only graphics-capable frames count // toward (and are demoted by) the budget; without a protocol every image is // already text. const suppressed = hasProtocol && this.#budget !== undefined ? this.#budget.observe(this.#imageId ?? 0) : false; - if (this.#cachedLines && this.#cachedWidth === width && this.#cachedSuppressed === suppressed) { + if ( + this.#cachedLines && + this.#cachedWidth === width && + this.#cachedSuppressed === suppressed && + this.#cachedImageProtocol === imageProtocol && + this.#cachedCellWidthPx === cellDimensions.widthPx && + this.#cachedCellHeightPx === cellDimensions.heightPx && + this.#cachedKittyUnicodePlaceholders === kittyUnicodePlaceholders + ) { return this.#cachedLines; } @@ -370,6 +387,10 @@ export class Image implements Component { this.#cachedLines = lines; this.#cachedWidth = width; this.#cachedSuppressed = suppressed; + this.#cachedImageProtocol = imageProtocol; + this.#cachedCellWidthPx = cellDimensions.widthPx; + this.#cachedCellHeightPx = cellDimensions.heightPx; + this.#cachedKittyUnicodePlaceholders = kittyUnicodePlaceholders; return lines; } diff --git a/packages/tui/test/image-budget.test.ts b/packages/tui/test/image-budget.test.ts index 8044faf66..33ddab97b 100644 --- a/packages/tui/test/image-budget.test.ts +++ b/packages/tui/test/image-budget.test.ts @@ -1,6 +1,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import { TUI } from "@oh-my-pi/pi-tui"; import { Image, ImageBudget } from "@oh-my-pi/pi-tui/components/image"; +import { Text } from "@oh-my-pi/pi-tui/components/text"; import { encodeKittyVirtualPlacement, getKittyGraphics, @@ -247,6 +248,59 @@ describe("Image budget integration", () => { expect([...budget.takeTransmits()]).toEqual([]); }); + it("moves back up before multi-row direct Kitty placements and restores the cursor below them", () => { + const budget = new ImageBudget(3, () => {}); + const id = budget.acquireId("k"); + const image = new Image( + BASE64_ONE_PIXEL_PNG, + "image/png", + { fallbackColor: t => t }, + { maxWidthCells: 4, maxHeightCells: 4, budget, imageKey: "k" }, + { widthPx: 40, heightPx: 40 }, + ); + + budget.beginPass(); + const lines = image.render(20); + budget.endPass(); + + const last = lines.at(-1) ?? ""; + expect(lines).toHaveLength(4); + expect(lines.slice(0, -1)).toEqual(["\x1b[0m", "\x1b[0m", "\x1b[0m"]); + expect(last.startsWith("\x1b[3A")).toBe(true); + expect(last).toContain("\x1b_Ga=p"); + expect(last).toContain("C=1"); + expect(last).toContain(`i=${id}`); + expect(last).toContain("c=4"); + expect(last).toContain("r=4"); + expect(last.endsWith("\x1b[3B")).toBe(true); + }); + + it("does not move the cursor around single-row direct Kitty placements", () => { + const budget = new ImageBudget(3, () => {}); + const id = budget.acquireId("k"); + const image = new Image( + BASE64_ONE_PIXEL_PNG, + "image/png", + { fallbackColor: t => t }, + { maxWidthCells: 4, maxHeightCells: 1, budget, imageKey: "k" }, + ); + + budget.beginPass(); + const lines = image.render(20); + budget.endPass(); + + const last = lines.at(-1) ?? ""; + expect(lines).toHaveLength(1); + expect(last.startsWith("\x1b_Ga=p")).toBe(true); + expect(last).toContain("C=1"); + expect(last).toContain(`i=${id}`); + expect(last).toContain("r=1"); + expect(last.endsWith("\x1b\\")).toBe(true); + expect(last).not.toContain("\x1b[0A"); + expect(last).not.toContain("\x1b[0B"); + expect(last).not.toMatch(/\x1b\[\d+[AB]/); + }); + it("renders an over-budget image as its text fallback instead of graphics", () => { const budget = new ImageBudget(1, () => {}); const older = new Image( @@ -393,6 +447,47 @@ describe("TUI inline-image budget", () => { ); } + it("renders following text below a multi-row direct Kitty placement", async () => { + const originalGraphics = { ...getKittyGraphics() }; + const term = new VirtualTerminal(40, 12); + const writes: string[] = []; + const realWrite = term.write.bind(term); + vi.spyOn(term, "write").mockImplementation((data: string) => { + writes.push(data); + realWrite(data); + }); + + setKittyGraphics({ unicodePlaceholders: false }); + const tui = new TUI(term); + tui.addChild( + new Image( + BASE64_ONE_PIXEL_PNG, + "image/png", + { fallbackColor: t => t }, + { maxWidthCells: 4, maxHeightCells: 4, budget: tui.imageBudget, imageKey: "direct" }, + { widthPx: 40, heightPx: 40 }, + ), + ); + tui.addChild(new Text("after-image", 0, 0)); + + try { + tui.start(); + await settle(term); + + const output = writes.join(""); + expect(output).toContain("\x1b[3A"); + expect(output).toContain("C=1"); + expect(output).toContain("\x1b[3B"); + + const viewport = term.getViewport().map(line => line.trimEnd()); + expect(viewport.slice(0, 5)).toEqual(["", "", "", "", "after-image"]); + expect(viewport.slice(0, 4).some(line => line.includes("after-image"))).toBe(false); + } finally { + tui.stop(); + setKittyGraphics(originalGraphics); + } + }); + it("purges demoted image graphics and repaints the fallback without a destructive replay", async () => { const term = new VirtualTerminal(40, 12); const writes: string[] = []; diff --git a/packages/tui/test/image-render.test.ts b/packages/tui/test/image-render.test.ts index d8700a90e..b247b4782 100644 --- a/packages/tui/test/image-render.test.ts +++ b/packages/tui/test/image-render.test.ts @@ -1,5 +1,6 @@ import { afterEach, beforeEach, describe, expect, it } from "bun:test"; -import { Image } from "@oh-my-pi/pi-tui/components/image"; +import { Image, ImageBudget } from "@oh-my-pi/pi-tui/components/image"; +import { getKittyGraphics, setKittyGraphics } from "@oh-my-pi/pi-tui/kitty-graphics"; import { type CellDimensions, getCellDimensions, @@ -34,16 +35,19 @@ function parseITermWidth(sequence: string): string | null { describe("terminal image rendering", () => { const originalProtocol = TERMINAL.imageProtocol; let originalCellDims: CellDimensions; + const originalGraphics = { ...getKittyGraphics() }; beforeEach(() => { originalCellDims = { ...getCellDimensions() }; setCellDimensions({ widthPx: 10, heightPx: 10 }); terminal.imageProtocol = null; + setKittyGraphics({ unicodePlaceholders: false }); }); afterEach(() => { setCellDimensions(originalCellDims); terminal.imageProtocol = originalProtocol; + setKittyGraphics(originalGraphics); }); it("fits Kitty images within max width and max height while preserving aspect ratio", () => { @@ -70,6 +74,65 @@ describe("terminal image rendering", () => { expect(parseKittyParam(result?.sequence ?? "", "C")).toBe(1); }); + it("re-renders a cached fallback once an image protocol becomes available", () => { + const image = new Image( + BASE64_ONE_PIXEL_PNG, + "image/png", + { fallbackColor: text => text }, + { maxWidthCells: 10, maxHeightCells: 2 }, + SQUARE_DIMENSIONS, + ); + + expect(image.render(20).join("")).toContain("[Image:"); + + terminal.imageProtocol = ImageProtocol.Kitty; + const rerendered = image.render(20).join(""); + + expect(rerendered).toContain("\x1b_Ga=T"); + expect(rerendered).toContain("C=1"); + }); + + it("re-renders a cached image when cell dimensions change", () => { + terminal.imageProtocol = ImageProtocol.Kitty; + const image = new Image( + BASE64_ONE_PIXEL_PNG, + "image/png", + { fallbackColor: text => text }, + { maxWidthCells: 10, maxHeightCells: 10 }, + SQUARE_DIMENSIONS, + ); + + const first = image.render(20).join(""); + expect(parseKittyParam(first, "c")).toBe(10); + + setCellDimensions({ widthPx: 20, heightPx: 10 }); + const second = image.render(20).join(""); + + expect(parseKittyParam(second, "c")).toBe(5); + }); + + it("re-renders a cached Kitty image when Unicode placeholder support changes", () => { + terminal.imageProtocol = ImageProtocol.Kitty; + setKittyGraphics({ unicodePlaceholders: false }); + const budget = new ImageBudget(1, () => {}); + const image = new Image( + BASE64_ONE_PIXEL_PNG, + "image/png", + { fallbackColor: text => text }, + { budget, imageKey: "placeholder-cache", maxWidthCells: 10, maxHeightCells: 2 }, + SQUARE_DIMENSIONS, + ); + + const direct = image.render(20).join(""); + expect(direct).toContain("\x1b_Ga=p"); + + setKittyGraphics({ unicodePlaceholders: true }); + const placeholder = image.render(20).join(""); + + expect(placeholder).toContain("U=1"); + expect(placeholder).not.toBe(direct); + }); + it("uses intrinsic image size when no bounds are provided", () => { terminal.imageProtocol = ImageProtocol.Kitty; const result = renderImage(BASE64_DUMMY, SQUARE_DIMENSIONS); @@ -117,25 +180,51 @@ describe("terminal image rendering", () => { expect((result?.sequence ?? "").startsWith("\x1bP")).toBe(true); }); - it("Image component forwards maxHeightCells to terminal rendering", () => { + it("moves back up before multi-row direct Kitty output and restores the cursor below it", () => { terminal.imageProtocol = ImageProtocol.Kitty; const image = new Image( BASE64_DUMMY, "image/png", { fallbackColor: text => text }, - { maxWidthCells: 10, maxHeightCells: 2 }, + { maxWidthCells: 10, maxHeightCells: 3 }, SQUARE_DIMENSIONS, ); const lines = image.render(20); + const imageLine = lines.at(-1) ?? ""; - expect(lines[0]).toBe("\x1b[0m"); - expect(lines).toHaveLength(2); - expect(lines[1]).toContain("\x1b[1A"); - expect(lines[1]).toContain("C=1"); - expect(lines[1]).toContain("c=2"); - expect(lines[1]).toContain("r=2"); - expect(lines[1]?.endsWith("\x1b[1B")).toBe(true); + expect(lines).toHaveLength(3); + expect(lines.slice(0, -1)).toEqual(["\x1b[0m", "\x1b[0m"]); + expect(imageLine.startsWith("\x1b[2A")).toBe(true); + expect(imageLine).toContain("\x1b_Ga=T"); + expect(imageLine).toContain("C=1"); + expect(imageLine).toContain("c=3"); + expect(imageLine).toContain("r=3"); + expect(imageLine.endsWith("\x1b[2B")).toBe(true); + }); + + it("does not emit cursor movement around single-row direct Kitty output", () => { + terminal.imageProtocol = ImageProtocol.Kitty; + const image = new Image( + BASE64_DUMMY, + "image/png", + { fallbackColor: text => text }, + { maxWidthCells: 10, maxHeightCells: 1 }, + SQUARE_DIMENSIONS, + ); + + const lines = image.render(20); + const imageLine = lines.at(-1) ?? ""; + + expect(lines).toHaveLength(1); + expect(imageLine.startsWith("\x1b_Ga=T")).toBe(true); + expect(imageLine).toContain("C=1"); + expect(imageLine).toContain("c=1"); + expect(imageLine).toContain("r=1"); + expect(imageLine.endsWith("\x1b\\")).toBe(true); + expect(imageLine).not.toContain("\x1b[0A"); + expect(imageLine).not.toContain("\x1b[0B"); + expect(imageLine).not.toMatch(/\x1b\[\d+[AB]/); }); });