diff --git a/packages/app/src/runtime/persistence/drafts.test.ts b/packages/app/src/runtime/persistence/drafts.test.ts index d611015da9d..9682d3b3e65 100644 --- a/packages/app/src/runtime/persistence/drafts.test.ts +++ b/packages/app/src/runtime/persistence/drafts.test.ts @@ -230,3 +230,72 @@ describe("draft store text externalization", () => { expect(JSON.parse(memory.documents.get("doc")!).prompt[0].content.blob.ids).toHaveLength(1) }) }) + +describe("draft store image retention", () => { + const image = (byte: number) => new Blob([new Uint8Array(6).fill(byte)], { type: "image/png" }) + const fresh = (grace = 0) => { + const memory = memoryDriver() + return { memory, store: createDraftStore(memory.driver, { grace }) } + } + // Release timers fire on the macrotask queue; a zero grace has fired after one tick. + const tick = () => new Promise((resolve) => setTimeout(resolve, 5)) + // An image with no object URL left gets a new one when its bytes are uploaded again. + const released = async (store: ReturnType, byte: number, url: string) => + (await store.putBlob(image(byte))).url !== url + + test("an uploaded image no document references is released after the grace", async () => { + const { store } = fresh() + const orphan = await store.putBlob(image(1)) + await tick() + expect(await released(store, 1, orphan.url)).toBe(true) + }) + + test("an image referenced within the grace is kept", async () => { + const { store } = fresh(50) + const pasted = await store.putBlob(image(2)) + await store.setDocument("pinned", { prompt: [{ type: "image", blob: pasted }] }) + await new Promise((resolve) => setTimeout(resolve, 80)) + expect(await released(store, 2, pasted.url)).toBe(false) + }) + + test("saving a document without an image or removing the document releases it", async () => { + const { store } = fresh() + const dropped = await store.putBlob(image(3)) + const removed = await store.putBlob(image(4)) + await store.setDocument("edited", { prompt: [{ type: "image", blob: dropped }] }) + await store.setDocument("closed", { prompt: [{ type: "image", blob: removed }] }) + await tick() + expect(await released(store, 3, dropped.url)).toBe(false) + expect(await released(store, 4, removed.url)).toBe(false) + await store.setDocument("edited", { prompt: [{ type: "text", content: "typed over it" }] }) + await store.removeItem("closed") + await tick() + expect(await released(store, 3, dropped.url)).toBe(true) + expect(await released(store, 4, removed.url)).toBe(true) + }) + + test("an image referenced by two documents survives until both drop it", async () => { + const { store } = fresh() + const shared = await store.putBlob(image(5)) + await store.setDocument("composer", { prompt: [{ type: "image", blob: shared }] }) + await store.setDocument("history", { entries: [{ prompt: [{ type: "image", blob: shared }] }] }) + await store.setDocument("composer", { prompt: [] }) + await tick() + expect(await released(store, 5, shared.url)).toBe(false) + await store.setDocument("history", { entries: [] }) + await tick() + expect(await released(store, 5, shared.url)).toBe(true) + }) + + test("loading a document pins the images it references", async () => { + const { memory, store } = fresh() + const id = await memory.driver.putBlob(image(6)) + memory.documents.set("loaded", JSON.stringify({ prompt: [{ type: "image", blob: { id } }] })) + const url = JSON.parse((await store.getItem("loaded"))!).prompt[0].blob.url + await tick() + expect(await released(store, 6, url)).toBe(false) + await store.removeItem("loaded") + await tick() + expect(await released(store, 6, url)).toBe(true) + }) +}) diff --git a/packages/app/src/runtime/persistence/drafts.ts b/packages/app/src/runtime/persistence/drafts.ts index d45949a17a2..e8c4713c069 100644 --- a/packages/app/src/runtime/persistence/drafts.ts +++ b/packages/app/src/runtime/persistence/drafts.ts @@ -28,23 +28,77 @@ export const draftTextThreshold = 16 * 1024 export const draftTextChunk = 64 * 1024 const textCacheLimit = 64 -const urls = new Map() -// The object URL already pins the Blob for the page's lifetime; keeping the Blob itself lets a -// collected image be uploaded again without fetching the URL. -const held = new Map() +// Decoded image bytes the renderer pins through object URLs. Every consumer of a `blob.url` is a +// persisted draft document (composer prompt, prompt history), so an image is pinned exactly while a +// stored document references it. Once the last reference disappears (removed from a draft, sent, or +// a discarded duplicate paste) the URL is revoked after this grace, which covers the persist delay +// between a paste and the save that references it, and the submit → history handoff. +export const retainedBlobGrace = 30_000 + +type Retained = { blob: Blob; url: string; release: ReturnType | undefined } +const retained = new Map() +// Document keys that reference each image id; an id with no keys is released after the grace. +const refs = new Map>() // Image ids that were restored under a different id (a store without WebCrypto assigns fresh // ones); live references still carry the original. const aliases = new Map() -function blobUrl(id: string, blob: Blob) { - const existing = urls.get(id) - if (existing) return existing +function blobUrl(id: string, blob: Blob, grace?: number) { + const existing = retained.get(id) + if (existing) return existing.url const url = URL.createObjectURL(blob) - urls.set(id, url) - held.set(id, blob) + // Without a grace the image has no store to reference it from and stays for the page's lifetime. + const release = grace === undefined || refs.get(id)?.size ? undefined : setTimeout(() => revoke(id), grace) + retained.set(id, { blob, url, release }) return url } +// Record which image ids `key` now references; ids it dropped are released once no other document +// references them, ids it gained stay pinned. +function retain(key: string, ids: ReadonlySet, grace: number) { + for (const [id, keys] of refs) { + if (ids.has(id) || !keys.delete(key) || keys.size) continue + refs.delete(id) + const entry = retained.get(id) + if (entry) entry.release = setTimeout(() => revoke(id), grace) + } + for (const id of ids) { + const keys = refs.get(id) ?? new Set() + keys.add(key) + refs.set(id, keys) + const entry = retained.get(id) + if (!entry) continue + clearTimeout(entry.release) + entry.release = undefined + } +} + +function revoke(id: string) { + const entry = retained.get(id) + if (!entry) return + URL.revokeObjectURL(entry.url) + retained.delete(id) + for (const [from, to] of aliases) if (to === id) aliases.delete(from) +} + +// Image ids a document references: `{ blob: { id } }` parts, not text chunk lists. +function imageIDs(value: unknown, into = new Set()): Set { + if (Array.isArray(value)) { + value.forEach((entry) => imageIDs(entry, into)) + return into + } + if (!value || typeof value !== "object") return into + const item = value as Record + const blob = item.blob + if (blob && typeof blob === "object" && !("kind" in blob)) { + const id = (blob as Record).id + if (typeof id === "string") into.add(id) + return into + } + Object.values(item).forEach((entry) => imageIDs(entry, into)) + return into +} + async function blobID(blob: Blob) { const bytes = crypto.subtle ? new Uint8Array(await crypto.subtle.digest("SHA-256", await blob.arrayBuffer())) @@ -60,24 +114,25 @@ export async function createBlobReference(blob: Blob): Promise { return { id, url: blobUrl(id, blob) } } -export function createDraftStore(driver: Driver): DraftStore { +export function createDraftStore(driver: Driver, options: { grace?: number } = {}): DraftStore { + const grace = options.grace ?? retainedBlobGrace const versions = new Map() const loading = new Map>() const loadBlobUrl = (id: string) => { - const existing = urls.get(id) - if (existing) return existing + const existing = retained.get(id) + if (existing) return existing.url const pending = loading.get(id) if (pending) return pending const next = driver .getBlob(id) - .then((blob) => (blob ? blobUrl(id, blob) : undefined)) + .then((blob) => (blob ? blobUrl(id, blob, grace) : undefined)) .finally(() => loading.delete(id)) loading.set(id, next) return next } const putBlob = async (blob: Blob) => { const id = await driver.putBlob(blob) - return { id, url: blobUrl(id, blob) } + return { id, url: blobUrl(id, blob, grace) } } // Keyed by chunk content so unchanged chunks are never hashed or sent again while the draft is // edited. Bounded because each entry pins up to draftTextChunk characters. A hit is safe even if @@ -145,8 +200,8 @@ export function createDraftStore(driver: Driver): DraftStore { if (typeof blob.id === "string") { // A live reference keeps the id it was created with; publish the id its bytes now live under. const id = aliases.get(blob.id) ?? blob.id - const kept = held.get(id) - const url = typeof blob.url === "string" ? blob.url : urls.get(id) + const kept = retained.get(id)?.blob + const url = typeof blob.url === "string" ? blob.url : retained.get(id)?.url if (kept) sources.set(id, { blob: async () => kept }) else if (url) sources.set(id, { blob: () => fetch(url).then((response) => response.blob()) }) return { ...item, blob: { id } } @@ -192,8 +247,7 @@ export function createDraftStore(driver: Driver): DraftStore { remember(chunks, next, source.chunk) return } - held.set(next, blob) - blobUrl(next, blob) + blobUrl(next, blob, grace) if (next === id) return // Later encodes of the still-live reference resolve straight to the new id. Re-point any // earlier alias chain so lookups stay one step. @@ -225,14 +279,19 @@ export function createDraftStore(driver: Driver): DraftStore { // stays visible until the bytes are back. Covers a blob collected while a cache, another tab, // or the composer's history still held its id. const missing = await driver.set(key, JSON.stringify(encoded), true) - if (missing.length === 0) return + if (missing.length === 0) { + retain(key, imageIDs(encoded), grace) + return + } const renamed = await restore(missing, sources) if (versions.get(key) !== version) return const unrestored = missing.filter((id) => !renamed.has(id)) if (unrestored.length) console.error(`[persistence] draft ${key} references blobs with no bytes to restore`, unrestored) // Anything still missing has no bytes anywhere; the owning codec drops such references on read. - await driver.set(key, JSON.stringify(rename(encoded, renamed)), false) + const final = rename(encoded, renamed) + await driver.set(key, JSON.stringify(final), false) + retain(key, imageIDs(final), grace) } return { getItem: async (key) => { @@ -241,6 +300,8 @@ export function createDraftStore(driver: Driver): DraftStore { const parsed = Schema.decodeUnknownOption(Schema.fromJsonString(Schema.Unknown))(value) // Let the owning persistence codec apply its invalid-document policy. if (Option.isNone(parsed)) return value + // A loaded document is live in the composer: pin its images before decode mints their URLs. + retain(key, imageIDs(parsed.value), grace) return JSON.stringify(await decode(parsed.value)) }, setItem: (key, value) => setDocument(key, JSON.parse(value)), @@ -248,6 +309,7 @@ export function createDraftStore(driver: Driver): DraftStore { removeItem: async (key) => { versions.set(key, (versions.get(key) ?? 0) + 1) await driver.remove(key) + retain(key, new Set(), grace) }, putBlob, } @@ -360,7 +422,8 @@ function referenced(json: string) { } export async function blobDataUrl(blob: BlobReference, mime: string) { - const data = await fetch(blob.url).then((response) => response.blob()) + const kept = retained.get(aliases.get(blob.id) ?? blob.id) + const data = kept ? kept.blob : await fetch(blob.url).then((response) => response.blob()) return new Promise((resolve, reject) => { const reader = new FileReader() reader.addEventListener("error", () => reject(reader.error))