From 4c930efbf011e7cd8d8332b85213638b52d938b5 Mon Sep 17 00:00:00 2001 From: Enrico Ros Date: Wed, 11 Feb 2026 12:59:36 -0800 Subject: [PATCH] Fix GC on Beams with reference collectors. --- .../chat-overlay/ConversationsManager.ts | 36 ++++++ src/common/stores/chat/chat.gc.ts | 105 ++++++++++-------- 2 files changed, 94 insertions(+), 47 deletions(-) diff --git a/src/common/chat-overlay/ConversationsManager.ts b/src/common/chat-overlay/ConversationsManager.ts index 36f5b06ca..3a26d033d 100644 --- a/src/common/chat-overlay/ConversationsManager.ts +++ b/src/common/chat-overlay/ConversationsManager.ts @@ -1,4 +1,6 @@ +import type { DBlobAssetId } from '~/common/stores/blob/dblobs-portability'; import type { DConversationId } from '~/common/stores/chat/chat.conversation'; +import { collectFragmentAssetIds, gcRegisterAssetCollector } from '~/common/stores/chat/chat.gc'; import { ConversationHandler } from './ConversationHandler'; @@ -14,6 +16,40 @@ export class ConversationsManager { private static _instance: ConversationsManager; private readonly handlers: Map = new Map(); + private constructor() { + // Register a GC collector to protect DBlob assets referenced in active Beam stores. + // Uses inversion of control to avoid circular dependency (chat/ -> chat-overlay/). + gcRegisterAssetCollector(() => this._collectBeamAssetIds()); + } + + /** + * Collect DBlob asset IDs from all active Beam stores (rays, fusions, follow-ups). + */ + private _collectBeamAssetIds(): DBlobAssetId[] { + const assetIds = new Set(); + for (const handler of this.handlers.values()) { + const { rays, fusions } = handler.getBeamStore().getState(); + + // Scatter rays + their follow-up messages + for (const ray of rays) { + collectFragmentAssetIds(ray.message.fragments, assetIds); + // if (ray.followUpMessages) + // for (const msg of ray.followUpMessages) + // collectFragmentAssetIds(msg.fragments, assetIds); + } + + // Gather fusions + their follow-up messages + for (const fusion of fusions) { + if (fusion.outputDMessage) + collectFragmentAssetIds(fusion.outputDMessage.fragments, assetIds); + // if (fusion.followUpMessages) + // for (const msg of fusion.followUpMessages) + // collectFragmentAssetIds(msg.fragments, assetIds); + } + } + return Array.from(assetIds); + } + static getHandler(conversationId: DConversationId): ConversationHandler { const instance = ConversationsManager._instance || (ConversationsManager._instance = new ConversationsManager()); let handler = instance.handlers.get(conversationId); diff --git a/src/common/stores/chat/chat.gc.ts b/src/common/stores/chat/chat.gc.ts index ca2f7fa8f..e4f68c100 100644 --- a/src/common/stores/chat/chat.gc.ts +++ b/src/common/stores/chat/chat.gc.ts @@ -1,65 +1,76 @@ import { DBlobAssetId, gcDBImageAssets } from '~/common/stores/blob/dblobs-portability'; +import type { Immutable } from '~/common/types/immutable.types'; + import type { DConversation } from './chat.conversation'; +import type { DMessageFragment } from './chat.fragments'; import { isContentOrAttachmentFragment, isImageRefPart, isZyncAssetReferencePart } from './chat.fragments'; import { useChatStore } from './store-chats'; + +// --- Asset collector registration --- + + +/** + * Allows external systems (Beam, scratch chat, etc.), to protect their DBlob assets from GC without creating circular dependencies. + */ +const _assetCollectors: AssetCollectorFn[] = []; +type AssetCollectorFn = () => DBlobAssetId[]; + +/** + * Register a callback that returns additional DBlob asset IDs to keep during GC. + * Uses inversion of control to avoid circular dependency (chat/ -> chat-overlay/). + * @returns unregister function + */ +export function gcRegisterAssetCollector(collector: AssetCollectorFn): () => void { + _assetCollectors.push(collector); + return () => { + const idx = _assetCollectors.indexOf(collector); + if (idx >= 0) _assetCollectors.splice(idx, 1); + }; +} + + +/** + * Collect DBlob asset IDs referenced in message fragments. + */ +export function collectFragmentAssetIds(fragments: Immutable, assetIds: Set): void { + for (const fragment of fragments) { + if (!isContentOrAttachmentFragment(fragment)) continue; + + // New References to Zync Assets (dblob refs for compatibility/migration) + if (isZyncAssetReferencePart(fragment.part) && fragment.part._legacyImageRefPart?.dataRef?.reftype === 'dblob') + assetIds.add(fragment.part._legacyImageRefPart.dataRef.dblobAssetId); + + // Legacy 'image_ref' parts (direct dblob refs) + if (isImageRefPart(fragment.part) && fragment.part.dataRef?.reftype === 'dblob') + assetIds.add(fragment.part.dataRef.dblobAssetId); + } +} + + /** * Garbage collect unreferenced dblobs in global chats * - This is ran as a side effect of the chat store rehydration * - This is also ran when a conversation or message is deleted, or when a conversation messages history is replaced */ -export async function gcChatImageAssets(conversations?: DConversation[]) { +export async function gcChatImageAssets(conversations?: Immutable) { // find all the dblob references in all chats const chatsAssetIDs: Set = new Set(); const _conversations = conversations || useChatStore.getState().conversations; - for (const chat of _conversations) { - for (const message of chat.messages) { - for (const fragment of message.fragments) { + for (const chat of _conversations) + for (const message of chat.messages) + collectFragmentAssetIds(message.fragments, chatsAssetIDs); - // only operate on content or attachment fragments - if (!isContentOrAttachmentFragment(fragment)) continue; - - // New References to Zync Assets (dblob refs for compatibility/migration) - if (isZyncAssetReferencePart(fragment.part) && fragment.part._legacyImageRefPart?.dataRef?.reftype === 'dblob') - chatsAssetIDs.add(fragment.part._legacyImageRefPart.dataRef.dblobAssetId); - - // Legacy 'image_ref' parts (direct dblob refs) - if (isImageRefPart(fragment.part) && fragment.part.dataRef?.reftype === 'dblob') - chatsAssetIDs.add(fragment.part.dataRef.dblobAssetId); - } - } - } - - // FIXME: [ASSET-GC-BEAM] GC deletes assets still referenced in Beam rays, causing images to disappear - // Bug occurs when: (1) Beam is open with imported rays containing images, (2) user regenerates/deletes - // those messages in the chat pane, (3) GC only scans main conversation store, not Beam vanilla stores, - // (4) assets are deleted while still displayed in Beam rays. - // Fix: Uncomment code below to scan all Beam stores for asset references before GC. - // Note: Also add import: import { ConversationsManager } from '~/common/chat-overlay/ConversationsManager'; - // Reproduction: Open Beam on right with images → regenerate (Ctrl+Shift+Z) on left -> images disappear. - // - // // Scan Beam rays for each conversation - // for (const conversation of _conversations) { - // const handler = ConversationsManager.getHandler(conversation.id); - // if (!handler.isValid()) continue; - // - // const rays = handler.beamStore.getState().rays; - // for (const ray of rays) { - // for (const fragment of ray.message.fragments) { - // if (!isContentOrAttachmentFragment(fragment)) continue; - // - // // New References to Zync Assets (dblob refs for compatibility/migration) - // if (isZyncAssetReferencePart(fragment.part) && fragment.part._legacyImageRefPart?.dataRef?.reftype === 'dblob') - // chatsAssetIDs.add(fragment.part._legacyImageRefPart.dataRef.dblobAssetId); - // - // // Legacy 'image_ref' parts (direct dblob refs) - // if (isImageRefPart(fragment.part) && fragment.part.dataRef?.reftype === 'dblob') - // chatsAssetIDs.add(fragment.part.dataRef.dblobAssetId); - // } - // } - // } + // [ASSET-GC-BEAM] Collect additional asset IDs from registered collectors (Beam, scratch chat, etc.) + // to prevent GC from deleting assets still displayed in ephemeral overlay stores (e.g. Beam rays/fusions). + // Bug: Beam images disappeared when regenerating/deleting chat messages while Beam was open, because + // GC only scanned conversation messages and not the vanilla Beam stores. Registration pattern avoids + // the circular dependency (chat/ -> chat-overlay/). + for (const collector of _assetCollectors) + for (const assetId of collector()) + chatsAssetIDs.add(assetId); // sanity check: if no blobs are referenced, do nothing; in case we have a state bug and we don't wipe the db if (!chatsAssetIDs.size) @@ -71,4 +82,4 @@ export async function gcChatImageAssets(conversations?: DConversation[]) { // FIXME: [ASSET] will only be able to GC local assets that haven't been uploaded to the cloud - otherwise they could be used, // in which case only the cloud can centralized-GC, or user will have to manually delete them -} \ No newline at end of file +}