diff --git a/src/ai/conversation.ts b/src/ai/conversation.ts index 26ee4d6..296ecde 100644 --- a/src/ai/conversation.ts +++ b/src/ai/conversation.ts @@ -1,4 +1,4 @@ -import { getSaver } from "#chaincraft/ai/memory/checkpoint-memory.js"; +import { listThreadIds } from "#chaincraft/ai/memory/checkpoint-memory.js"; // Map of conversation IDs by graph type const conversationIds = new Map>(); @@ -30,16 +30,7 @@ async function bootstrapConversationIds(graphType: string): Promise { return; } - const saver = await getSaver('list-conversations', graphType); - const ids = new Set(); - - for await (const checkpoint of saver.list({}, {})) { - const threadId = checkpoint.config?.configurable?.thread_id; - if (threadId) { - ids.add(threadId); - } - } - - conversationIds.set(graphType, ids); + const threadIds = await listThreadIds(graphType); + conversationIds.set(graphType, new Set(threadIds)); bootstrappedTypes.add(graphType); } \ No newline at end of file diff --git a/src/ai/design/design-workflow.ts b/src/ai/design/design-workflow.ts index 76c825e..d39a946 100644 --- a/src/ai/design/design-workflow.ts +++ b/src/ai/design/design-workflow.ts @@ -262,6 +262,8 @@ export async function getDesignByVersion( conversationId: string, version: number ): Promise<(DesignState | undefined)> { + // TODO - Consider paginating through checkpoints if we expect many versions, + // to avoid excessive memory usage. // Check if conversation exists if (!(await isActiveConversation(conversationId))) { throw new Error(`Conversation ${conversationId} not found`); diff --git a/src/ai/memory/checkpoint-memory.ts b/src/ai/memory/checkpoint-memory.ts index 6442ff8..f793eb3 100644 --- a/src/ai/memory/checkpoint-memory.ts +++ b/src/ai/memory/checkpoint-memory.ts @@ -147,6 +147,40 @@ export async function getSaver( } } +/** + * List all thread IDs for a given graph type. + * Uses efficient direct SQL query for PostgreSQL, falls back to checkpoint list for SQLite. + */ +export async function listThreadIds(graphType: string): Promise { + await initialize(); + + const db = await getOrCreateDatabase(graphType); + const ids = new Set(); + + if (db.backend === "postgres" && db.sharedSaver) { + // For PostgreSQL: query distinct thread_ids directly (much more efficient than loading checkpoints) + const postgresaver = db.sharedSaver as PostgresSaver; + const pool = (postgresaver as any).pool as Pool; + const result = await pool.query( + 'SELECT DISTINCT thread_id FROM checkpoints ORDER BY thread_id' + ); + for (const row of result.rows) { + ids.add(row.thread_id); + } + } else { + // For SQLite: use limited checkpoint list + const saver = await getSaver('list-threads', graphType); + for await (const checkpoint of saver.list({}, { limit: 100 })) { + const threadId = checkpoint.config?.configurable?.thread_id; + if (threadId) { + ids.add(threadId); + } + } + } + + return Array.from(ids).sort(); +} + export async function deleteThread( threadId: string, graphType: string