zack3d/big-agi
1
0
Fork 0
mirror of https://github.com/enricoros/big-AGI.git synced 2026-05-10 21:50:14 -07:00
Code Issues Packages Projects Releases Wiki Activity

Fix GC on Beams with reference collectors.

Browse Source
This commit is contained in:
Enrico Ros
2026-02-11 12:59:36 -08:00
parent 5a2a47cb87
commit 4c930efbf0
2 changed files with 94 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/common/chat-overlay/ConversationsManager.ts
+36
View File
@@ -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<DConversationId, ConversationHandler> = 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<DBlobAssetId>();
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);
src/common/stores/chat/chat.gc.ts
+58 -47
View File
@@ -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<DMessageFragment[]>, assetIds: Set<DBlobAssetId>): 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<DConversation[]>) {
// find all the dblob references in all chats
const chatsAssetIDs: Set<DBlobAssetId> = 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
}
}