Repository navigation
[companion] fix(chat): show tool-result images as cached previews #83
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: companion/deferred-tool-images-base
Are you sure you want to change the base?
Changes from all commits
01f91ef
d672bd8
2a6bdc8
7865950
b8d39b5
943ea1e
c5ed579
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,53 @@ | ||
| import { readFile } from "node:fs/promises"; | ||
| import { createThumbnail, existingThumbnail, imageMimeType, mediaFilePath } from "@/lib/media-cache"; | ||
|
|
||
| export const dynamic = "force-dynamic"; | ||
|
|
||
| function image(body: Buffer, type: string, immutable = true): Response { | ||
| // A view, not a copy: these buffers come from readFile/sharp, never a SharedArrayBuffer. | ||
| return new Response(new Uint8Array(body.buffer as ArrayBuffer, body.byteOffset, body.byteLength), { | ||
| headers: { | ||
| "Content-Type": type, | ||
| // Content-addressed: the bytes behind a hash never change. | ||
| "Cache-Control": immutable ? "private, max-age=31536000, immutable" : "private, no-cache", | ||
| "X-Content-Type-Options": "nosniff", | ||
| }, | ||
| }); | ||
| } | ||
|
|
||
| async function readIfPresent(file: string): Promise<Buffer | null> { | ||
| try { | ||
| return await readFile(file); | ||
| } catch (error) { | ||
| if ((error as NodeJS.ErrnoException).code === "ENOENT") return null; | ||
| throw error; | ||
| } | ||
| } | ||
|
|
||
| /** GET /api/media/<sha256>[?thumb=1] — a tool-result image, or its small WebP preview. */ | ||
| export async function GET(req: Request, { params }: { params: Promise<{ hash: string }> }) { | ||
| const { hash } = await params; | ||
| const thumb = new URL(req.url).searchParams.has("thumb"); | ||
| const cached = thumb ? existingThumbnail(hash) : null; | ||
| const cachedBytes = cached ? await readIfPresent(cached) : null; | ||
| if (cachedBytes) return image(cachedBytes, "image/webp"); | ||
|
|
||
| // The sweep can delete our copy between lookup and read once omp's blob is | ||
| // complete; one more lookup then finds the blob. | ||
| let bytes: Buffer | null = null; | ||
| for (let attempt = 0; attempt < 2 && !bytes; attempt++) { | ||
| const file = mediaFilePath(hash); | ||
| if (!file) break; | ||
| bytes = await readIfPresent(file); | ||
| } | ||
| if (!bytes) return new Response("Not found", { status: 404 }); | ||
| const type = imageMimeType(bytes); | ||
| if (!type) return new Response("Not an image", { status: 415 }); | ||
| if (thumb) { | ||
| const preview = await createThumbnail(hash, bytes); | ||
| if (preview) return image(preview, "image/webp"); | ||
| // Full image standing in for a preview: not cached, so a later preview can replace it. | ||
| return image(bytes, type, false); | ||
| } | ||
| return image(bytes, type); | ||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,225 @@ | ||
| import { createHash, randomUUID } from "node:crypto"; | ||
| import { existsSync, mkdirSync, readdirSync, renameSync, rmSync, statSync, unlinkSync, writeFileSync } from "node:fs"; | ||
| import path from "node:path"; | ||
| import { getBlobsDir, getDiagnosticsDir } from "./omp/paths"; | ||
| import { isRecord } from "./type-guards"; | ||
|
|
||
| /** | ||
| * Tool-result images reach the browser as `/api/media/<sha256>` URLs instead of | ||
| * inline base64, so big screenshots never sit in the chat's React state. The | ||
| * hash is omp's own blob address (sha256 of the image bytes): history entries | ||
| * already carry it as `blob:sha256:<hash>`, and live output is hashed here and | ||
| * copied into omp-web's media directory until omp writes its own blob. | ||
| */ | ||
|
|
||
| const HASH_RE = /^[a-f0-9]{64}$/; | ||
| const BLOB_PREFIX = "blob:sha256:"; | ||
| /** Twice the 240px tool-result preview, for high-density screens. */ | ||
| const THUMB_PX = 480; | ||
| const SWEEP_EVERY_MS = 24 * 60 * 60 * 1000; | ||
| const MAX_AGE_MS = 30 * 24 * 60 * 60 * 1000; | ||
|
|
||
| declare global { | ||
| var __ompWebMediaSweptAt: number | undefined; | ||
| } | ||
|
|
||
| function mediaDir(): string { | ||
| return path.join(getDiagnosticsDir(), "media"); | ||
| } | ||
|
|
||
| function fileSize(file: string): number | undefined { | ||
| return statSync(file, { throwIfNoEntry: false })?.size; | ||
| } | ||
|
|
||
| function writeAtomic(file: string, data: Buffer): void { | ||
| sweepMediaCache(); | ||
| mkdirSync(path.dirname(file), { recursive: true }); | ||
| const tmp = `${file}.${randomUUID()}.tmp`; | ||
| try { | ||
| writeFileSync(tmp, data); | ||
| renameSync(tmp, file); | ||
| } catch (error) { | ||
| rmSync(tmp, { force: true }); | ||
| throw error; | ||
| } | ||
| } | ||
|
|
||
| /** Daily: drop copies omp now holds as complete blobs, and anything older than 30 days (previews regenerate on demand). */ | ||
| function sweepMediaCache(): void { | ||
| const now = Date.now(); | ||
| if (now - (globalThis.__ompWebMediaSweptAt ?? 0) < SWEEP_EVERY_MS) return; | ||
| globalThis.__ompWebMediaSweptAt = now; | ||
| let names: string[]; | ||
| try { | ||
| names = readdirSync(mediaDir()); | ||
| } catch { | ||
| return; | ||
| } | ||
| for (const name of names) { | ||
| const file = path.join(mediaDir(), name); | ||
| try { | ||
| const stat = statSync(file); | ||
| const hasBlob = HASH_RE.test(name) && fileSize(path.join(getBlobsDir(), name)) === stat.size; | ||
| if (hasBlob || now - stat.mtimeMs > MAX_AGE_MS) unlinkSync(file); | ||
| } catch { | ||
| // Removed concurrently or unreadable: the next sweep retries. | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Live tool output resends the full accumulated result on every update; this | ||
| // keeps one screenshot from being decoded and hashed once per frame. Each key | ||
| // is a whole base64 image, so the memo is capped by size. | ||
| const recentHashes = new Map<string, string>(); | ||
| const RECENT_MAX_CHARS = 16 * 1024 * 1024; | ||
| let recentChars = 0; | ||
|
|
||
| /** Store a base64 image under its content address and return the hash. */ | ||
| export function storeInlineImage(base64: string): string { | ||
| // The memo skips only decoding and hashing: the files are rechecked every | ||
| // time, because the sweep may have removed the copy since. | ||
| let hash = recentHashes.get(base64); | ||
| if (!hash) { | ||
| hash = createHash("sha256").update(Buffer.from(base64, "base64")).digest("hex"); | ||
| if (base64.length <= RECENT_MAX_CHARS) { | ||
| recentHashes.set(base64, hash); | ||
| recentChars += base64.length; | ||
| while (recentChars > RECENT_MAX_CHARS) { | ||
| const oldest = recentHashes.keys().next().value!; | ||
| recentChars -= oldest.length; | ||
| recentHashes.delete(oldest); | ||
| } | ||
| } | ||
| } | ||
| const size = Buffer.byteLength(base64, "base64"); | ||
| const copy = path.join(mediaDir(), hash); | ||
| // omp writes blobs in place: keep our own copy unless its blob is already complete. | ||
| if (fileSize(path.join(getBlobsDir(), hash)) !== size && fileSize(copy) !== size) writeAtomic(copy, Buffer.from(base64, "base64")); | ||
| return hash; | ||
| } | ||
|
|
||
| function imageBlockAsUrl(block: unknown): unknown { | ||
| if (!isRecord(block) || block.type !== "image") return block; | ||
| let data: unknown; | ||
| let mimeType: unknown; | ||
| if (typeof block.data === "string") { | ||
| data = block.data; | ||
| mimeType = block.mimeType; | ||
| } else if (isRecord(block.source) && block.source.type === "base64") { | ||
| data = block.source.data; | ||
| mimeType = block.source.media_type; | ||
| } | ||
| if (typeof data !== "string" || data === "") return block; | ||
| let hash: string; | ||
| if (data.startsWith(BLOB_PREFIX)) { | ||
| hash = data.slice(BLOB_PREFIX.length); | ||
| if (!HASH_RE.test(hash)) return block; | ||
| if (!existsSync(path.join(getBlobsDir(), hash)) && !existsSync(path.join(mediaDir(), hash))) { | ||
| return { type: "text", text: `[image unavailable: blob ${hash.slice(0, 12)}… not found]` }; | ||
| } | ||
| } else { | ||
| try { | ||
| hash = storeInlineImage(data); | ||
| } catch { | ||
| return block; // Disk trouble: keep showing the inline image. | ||
| } | ||
| } | ||
| return { type: "image", ...(typeof mimeType === "string" ? { mimeType } : {}), url: `/api/media/${hash}` }; | ||
| } | ||
|
|
||
| /** Tool-result content with every image replaced by a media URL. */ | ||
| export function toolResultContentAsUrls<T>(content: T): T { | ||
| if (!Array.isArray(content) || !content.some((block) => isRecord(block) && block.type === "image")) return content; | ||
| return content.map(imageBlockAsUrl) as T; | ||
| } | ||
|
|
||
| function withContentAsUrls(value: unknown): unknown { | ||
| if (!isRecord(value)) return value; | ||
| const content = toolResultContentAsUrls(value.content); | ||
| return content === value.content ? value : { ...value, content }; | ||
| } | ||
|
|
||
| function toolResultMessageAsUrls(message: unknown): unknown { | ||
| return isRecord(message) && message.role === "toolResult" ? withContentAsUrls(message) : message; | ||
| } | ||
|
|
||
| function mapMessages(value: unknown): unknown { | ||
| if (!Array.isArray(value)) return value; | ||
| const mapped = value.map(toolResultMessageAsUrls); | ||
| return mapped.some((message, i) => message !== value[i]) ? mapped : value; | ||
| } | ||
|
|
||
| /** An omp agent event with its tool-result images as media URLs (same object when it has none). */ | ||
| export function eventWithToolResultImageUrls<E extends { type: string; [key: string]: unknown }>(event: E): E { | ||
| const swap = (key: string, value: unknown): E => (value === event[key] ? event : { ...event, [key]: value }); | ||
| switch (event.type) { | ||
| case "tool_execution_update": return swap("partialResult", withContentAsUrls(event.partialResult)); | ||
| case "tool_execution_end": return swap("result", withContentAsUrls(event.result)); | ||
| case "message_start": | ||
| case "message_end": return swap("message", toolResultMessageAsUrls(event.message)); | ||
| case "turn_end": return swap("toolResults", mapMessages(event.toolResults)); | ||
| case "agent_end": return swap("messages", mapMessages(event.messages)); | ||
| case "subagent_event": { | ||
| // The child's raw event rides in `payload.event`. | ||
| const payload = event.payload; | ||
| if (!isRecord(payload) || !isRecord(payload.event) || typeof payload.event.type !== "string") return event; | ||
| const inner = eventWithToolResultImageUrls(payload.event as { type: string }); | ||
| return inner === payload.event ? event : swap("payload", { ...payload, event: inner }); | ||
| } | ||
| default: return event; | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
When a subagent runs an image-producing tool, AGENTS.md reference: AGENTS.md:L250-L254 Useful? React with 👍 / 👎.
Owner
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fixed in c5ed579: |
||
| } | ||
| } | ||
|
|
||
| /** | ||
| * File holding a stored image: omp's blob, else omp-web's copy. omp writes | ||
| * blobs in place, so a blob that differs in size from our copy is still being | ||
| * written and the copy is served instead. | ||
| */ | ||
| export function mediaFilePath(hash: string): string | null { | ||
| if (!HASH_RE.test(hash)) return null; | ||
| const blob = path.join(getBlobsDir(), hash); | ||
| const copy = path.join(mediaDir(), hash); | ||
| const blobSize = fileSize(blob); | ||
| const copySize = fileSize(copy); | ||
| if (blobSize !== undefined && (copySize === undefined || blobSize === copySize)) return blob; | ||
| return copySize === undefined ? null : copy; | ||
| } | ||
|
|
||
| /** Raster image type from magic bytes; null for anything else (the blob store also holds non-images). */ | ||
| export function imageMimeType(bytes: Buffer): string | null { | ||
| if (bytes.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))) return "image/png"; | ||
| if (bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff) return "image/jpeg"; | ||
| if (bytes.subarray(0, 4).toString("latin1") === "GIF8") return "image/gif"; | ||
| if (bytes.subarray(0, 4).toString("latin1") === "RIFF" && bytes.subarray(8, 12).toString("latin1") === "WEBP") return "image/webp"; | ||
| return null; | ||
| } | ||
|
|
||
| /** Existing thumbnail for `hash`, if one was generated before. */ | ||
| export function existingThumbnail(hash: string): string | null { | ||
| const thumb = path.join(mediaDir(), `${hash}.thumb.webp`); | ||
| return HASH_RE.test(hash) && existsSync(thumb) ? thumb : null; | ||
| } | ||
|
|
||
| /** Generate and cache a small WebP preview; null when sharp is unavailable or fails. */ | ||
| export async function createThumbnail(hash: string, image: Buffer): Promise<Buffer | null> { | ||
| let thumb: Buffer; | ||
| try { | ||
| // Dynamic on purpose: sharp's native binary is platform-specific and may be | ||
| // missing, and this module is also loaded by the session reader and RPC | ||
| // manager, which must not fail or load libvips just to rewrite URLs. | ||
| const { default: sharp } = await import("sharp"); | ||
| thumb = await sharp(image) | ||
| .rotate() // apply EXIF orientation; WebP output drops the tag | ||
| .resize(THUMB_PX, THUMB_PX, { fit: "inside", withoutEnlargement: true }) | ||
| .webp({ quality: 80 }) | ||
| .toBuffer(); | ||
| } catch { | ||
| return null; | ||
| } | ||
| try { | ||
| writeAtomic(path.join(mediaDir(), `${hash}.thumb.webp`), thumb); | ||
| } catch { | ||
| // Not cached this time; the preview is still served. | ||
| } | ||
| return thumb; | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When
createThumbnail()returnsnullbecause Sharp is unavailable or transiently fails, this branch serves the original image from the?thumb=1URL throughimage(), which applies a one-yearimmutablecache header. If thumbnail generation later becomes available, the same URL returns a WebP preview, but existing clients remain pinned to the full-size image and continue paying its decode-memory cost; return the fallback without immutable caching or redirect to the full-image URL.AGENTS.md reference: AGENTS.md:L261-L263
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Fixed in 7865950: when no preview can be made, the full image served for
?thumb=1now carriesCache-Control: private, no-cache, so a later WebP preview replaces it.