@vellumai/assistant 0.5.1 → 0.5.3

package/src/memory/conversation-crud.ts

@@ -1,4 +1,17 @@
-import { and, asc, count, eq, inArray, isNull, sql } from "drizzle-orm";
+import { mkdirSync, rmSync } from "node:fs";
+
+import {
+  and,
+  asc,
+  count,
+  desc,
+  eq,
+  gt,
+  inArray,
+  isNull,
+  lte,
+  sql,
+} from "drizzle-orm";
 import { v4 as uuid } from "uuid";
 import { z } from "zod";
 
@@ -7,10 +20,24 @@ import { parseChannelId, parseInterfaceId } from "../channels/types.js";
 import { CHANNEL_IDS, INTERFACE_IDS, isChannelId } from "../channels/types.js";
 import { getConfig } from "../config/loader.js";
 import type { TrustContext } from "../daemon/conversation-runtime-assembly.js";
+import { UserError } from "../util/errors.js";
 import { getLogger } from "../util/logger.js";
+import { getConversationsDir } from "../util/platform.js";
 import { createRowMapper } from "../util/row-mapper.js";
-import { deleteOrphanAttachments } from "./attachments-store.js";
-import { projectAssistantMessage } from "./conversation-attention-store.js";
+import {
+  deleteOrphanAttachments,
+  linkAttachmentToMessage,
+} from "./attachments-store.js";
+import {
+  projectAssistantMessage,
+  seedForkedConversationAttention,
+} from "./conversation-attention-store.js";
+import {
+  initConversationDir,
+  removeConversationDir,
+  syncMessageToDisk,
+  updateMetaFile,
+} from "./conversation-disk-view.js";
 import { ensureDisplayOrderMigration } from "./conversation-display-order-migration.js";
 import { getDb, rawAll, rawExec, rawGet, rawRun } from "./db.js";
 import { indexMessageNow } from "./indexer.js";
@@ -18,6 +45,7 @@ import { enqueueMemoryJob } from "./jobs-store.js";
 import {
   channelInboundEvents,
   conversations,
+  conversationStarters,
   llmRequestLogs,
   memoryEmbeddings,
   memoryItems,
@@ -47,6 +75,9 @@ const subagentNotificationSchema = z.object({
   conversationId: z.string().optional(),
 });
 
+export const PRIVATE_CONVERSATION_FORK_ERROR =
+  "Private conversations cannot be forked";
+
 export const messageMetadataSchema = z
   .object({
     userMessageChannel: channelIdSchema.optional(),
@@ -67,11 +98,42 @@ export const messageMetadataSchema = z
     provenanceGuardianExternalUserId: z.string().optional(),
     provenanceRequesterIdentifier: z.string().optional(),
     automated: z.boolean().optional(),
+    forkSourceMessageId: z.string().optional(),
+    /** Image source paths from desktop attachments, keyed by filename. */
+    imageSourcePaths: z.record(z.string(), z.string()).optional(),
   })
   .passthrough();
 
 export type MessageMetadata = z.infer<typeof messageMetadataSchema>;
 
+function cloneForkMessageMetadata(
+  metadata: string | null,
+  sourceMessageId: string,
+): string {
+  if (!metadata) {
+    return JSON.stringify({ forkSourceMessageId: sourceMessageId });
+  }
+
+  try {
+    const parsed = JSON.parse(metadata);
+    if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
+      const sourceRecord = parsed as Record<string, unknown>;
+      const forkSourceMessageId =
+        typeof sourceRecord.forkSourceMessageId === "string"
+          ? sourceRecord.forkSourceMessageId
+          : sourceMessageId;
+      return JSON.stringify({
+        ...sourceRecord,
+        forkSourceMessageId,
+      });
+    }
+  } catch {
+    // Fall through to source-only metadata.
+  }
+
+  return JSON.stringify({ forkSourceMessageId: sourceMessageId });
+}
+
 /**
  * Extract provenance metadata fields from a TrustContext.
  * When no guardian context is provided, defaults to 'unknown' because the
@@ -106,8 +168,13 @@ export interface ConversationRow {
   memoryScopeId: string;
   originChannel: string | null;
   originInterface: string | null;
+  forkParentConversationId: string | null;
+  forkParentMessageId: string | null;
   isAutoTitle: number;
   scheduleJobId: string | null;
+  memoryReducedThroughMessageId: string | null;
+  memoryDirtyTailSinceMessageId: string | null;
+  memoryLastReducedAt: number | null;
 }
 
 export const parseConversation = createRowMapper<
@@ -129,8 +196,13 @@ export const parseConversation = createRowMapper<
   memoryScopeId: "memoryScopeId",
   originChannel: "originChannel",
   originInterface: "originInterface",
+  forkParentConversationId: "forkParentConversationId",
+  forkParentMessageId: "forkParentMessageId",
   isAutoTitle: "isAutoTitle",
   scheduleJobId: "scheduleJobId",
+  memoryReducedThroughMessageId: "memoryReducedThroughMessageId",
+  memoryDirtyTailSinceMessageId: "memoryDirtyTailSinceMessageId",
+  memoryLastReducedAt: "memoryLastReducedAt",
 });
 
 export interface MessageRow {
@@ -232,6 +304,8 @@ export function createConversation(
     }
   }
 
+  initConversationDir({ ...conversation, originChannel: null });
+
   return conversation;
 }
 
@@ -258,6 +332,213 @@ export function getConversationMemoryScopeId(conversationId: string): string {
   return conv?.memoryScopeId ?? "default";
 }
 
+export function forkConversation(params: {
+  conversationId: string;
+  throughMessageId?: string;
+}): ConversationRow {
+  const { conversationId, throughMessageId } = params;
+  const db = getDb();
+  const sourceConversation = getConversation(conversationId);
+
+  if (!sourceConversation) {
+    throw new UserError(`Conversation ${conversationId} not found`);
+  }
+  if (sourceConversation.conversationType === "private") {
+    throw new UserError(PRIVATE_CONVERSATION_FORK_ERROR);
+  }
+
+  const sourceMessages = getMessages(conversationId);
+
+  if (sourceMessages.length === 0) {
+    throw new UserError(
+      `Conversation ${conversationId} has no persisted messages to fork`,
+    );
+  }
+
+  const copyBoundaryIndex =
+    throughMessageId == null
+      ? sourceMessages.length - 1
+      : sourceMessages.findIndex((message) => message.id === throughMessageId);
+
+  if (throughMessageId != null && copyBoundaryIndex === -1) {
+    throw new UserError(
+      `Message ${throughMessageId} does not belong to conversation ${conversationId}`,
+    );
+  }
+
+  const visibleWindowStartIndex = Math.max(
+    0,
+    Math.min(
+      sourceConversation.contextCompactedMessageCount,
+      sourceMessages.length,
+    ),
+  );
+  const preserveSourceCompactionState =
+    copyBoundaryIndex >= visibleWindowStartIndex;
+
+  const messagesToCopy =
+    copyBoundaryIndex >= 0
+      ? sourceMessages.slice(0, copyBoundaryIndex + 1)
+      : ([] as MessageRow[]);
+  const forkParentMessageId = messagesToCopy.at(-1)?.id ?? null;
+  const forkTitle = `${sourceConversation.title ?? "Untitled"} (Fork)`;
+
+  // Collect disk-sync work to run after the transaction commits.
+  const diskSyncQueue: Array<{
+    conversationId: string;
+    messageId: string;
+    createdAt: number;
+  }> = [];
+
+  // Wrap all DB mutations in a single transaction so a mid-flight failure
+  // rolls back cleanly instead of leaving a partial fork. Helper functions
+  // (linkAttachmentToMessage, relinkAttachments, seedForkedConversationAttention)
+  // use the same underlying bun:sqlite connection, so their writes participate
+  // in this transaction automatically.
+  const forkedConversation = db.transaction(() => {
+    const fc = createConversation({
+      title: forkTitle,
+      conversationType: "standard",
+    });
+
+    db.update(conversations)
+      .set({
+        forkParentConversationId: sourceConversation.id,
+        forkParentMessageId,
+        contextSummary: preserveSourceCompactionState
+          ? sourceConversation.contextSummary
+          : null,
+        contextCompactedMessageCount: preserveSourceCompactionState
+          ? sourceConversation.contextCompactedMessageCount
+          : 0,
+        contextCompactedAt: preserveSourceCompactionState
+          ? sourceConversation.contextCompactedAt
+          : null,
+      })
+      .where(eq(conversations.id, fc.id))
+      .run();
+
+    const forkedMessageIds = new Map<string, string>();
+    let latestForkedAssistant: {
+      messageId: string;
+      messageAt: number;
+    } | null = null;
+
+    for (const message of messagesToCopy) {
+      const forkedMessageId = uuid();
+      db.insert(messages)
+        .values({
+          id: forkedMessageId,
+          conversationId: fc.id,
+          role: message.role,
+          content: message.content,
+          createdAt: message.createdAt,
+          metadata: cloneForkMessageMetadata(message.metadata, message.id),
+        })
+        .run();
+      forkedMessageIds.set(message.id, forkedMessageId);
+
+      if (message.role === "assistant") {
+        latestForkedAssistant = {
+          messageId: forkedMessageId,
+          messageAt: message.createdAt,
+        };
+      }
+    }
+
+    const attachmentIdMap = new Map<string, string>();
+    for (const message of messagesToCopy) {
+      const forkedMessageId = forkedMessageIds.get(message.id);
+      if (!forkedMessageId) continue;
+
+      const attachmentLinks = db
+        .select({
+          attachmentId: messageAttachments.attachmentId,
+          position: messageAttachments.position,
+        })
+        .from(messageAttachments)
+        .where(eq(messageAttachments.messageId, message.id))
+        .orderBy(messageAttachments.position)
+        .all();
+      const uncachedAttachmentLinks = attachmentLinks.filter(
+        (link) => !attachmentIdMap.has(link.attachmentId),
+      );
+      const stagingMessageId =
+        uncachedAttachmentLinks.length > 0 ? uuid() : null;
+
+      if (stagingMessageId) {
+        db.insert(messages)
+          .values({
+            id: stagingMessageId,
+            conversationId: fc.id,
+            role: message.role,
+            content: "",
+            createdAt: message.createdAt,
+            metadata: null,
+          })
+          .run();
+      }
+
+      for (const link of attachmentLinks) {
+        const cachedAttachmentId = attachmentIdMap.get(link.attachmentId);
+        if (cachedAttachmentId) {
+          db.insert(messageAttachments)
+            .values({
+              id: uuid(),
+              messageId: forkedMessageId,
+              attachmentId: cachedAttachmentId,
+              position: link.position,
+              createdAt: Date.now(),
+            })
+            .run();
+          continue;
+        }
+
+        const scopedAttachmentId = linkAttachmentToMessage(
+          stagingMessageId ?? forkedMessageId,
+          link.attachmentId,
+          link.position,
+        );
+        attachmentIdMap.set(link.attachmentId, scopedAttachmentId);
+      }
+
+      if (stagingMessageId) {
+        relinkAttachments([stagingMessageId], forkedMessageId);
+        db.delete(messages).where(eq(messages.id, stagingMessageId)).run();
+      }
+
+      diskSyncQueue.push({
+        conversationId: fc.id,
+        messageId: forkedMessageId,
+        createdAt: fc.createdAt,
+      });
+    }
+
+    seedForkedConversationAttention({
+      conversationId: fc.id,
+      latestAssistantMessageId: latestForkedAssistant?.messageId ?? null,
+      latestAssistantMessageAt: latestForkedAssistant?.messageAt ?? null,
+    });
+
+    return fc;
+  });
+
+  // Disk-view sync runs after commit — file I/O is idempotent and
+  // conversation deletion cleans up orphaned directories.
+  for (const entry of diskSyncQueue) {
+    syncMessageToDisk(entry.conversationId, entry.messageId, entry.createdAt);
+  }
+
+  const persistedFork = getConversation(forkedConversation.id);
+  if (!persistedFork) {
+    throw new Error(
+      `Failed to load forked conversation ${forkedConversation.id} after creation`,
+    );
+  }
+
+  return persistedFork;
+}
+
 /**
  * Delete a conversation and all its messages, cleaning up orphaned memory
  * artifacts (items, embeddings). Returns segment and orphaned item IDs so
@@ -265,7 +546,18 @@ export function getConversationMemoryScopeId(conversationId: string): string {
  */
 export function deleteConversation(id: string): DeletedMemoryIds {
   const db = getDb();
-  const result: DeletedMemoryIds = { segmentIds: [], orphanedItemIds: [] };
+  const result: DeletedMemoryIds = {
+    segmentIds: [],
+    orphanedItemIds: [],
+    deletedSummaryIds: [],
+  };
+
+  // Capture createdAt before the transaction deletes the row — needed to
+  // resolve the conversation's disk-view directory path after deletion.
+  const convBeforeDelete = getConversation(id);
+  const createdAtForDiskCleanup = convBeforeDelete?.createdAt;
+  const memoryScopeId = convBeforeDelete?.memoryScopeId;
+  const isPrivateScope = memoryScopeId?.startsWith("private:") ?? false;
 
   db.transaction((tx) => {
     // Collect all message IDs for this conversation.
@@ -354,9 +646,73 @@ export function deleteConversation(id: string): DeletedMemoryIds {
         .run();
     }
 
+    if (isPrivateScope && memoryScopeId) {
+      // Sweep remaining memory items with this private scopeId.
+      const scopeItems = tx
+        .select({ id: memoryItems.id })
+        .from(memoryItems)
+        .where(eq(memoryItems.scopeId, memoryScopeId))
+        .all();
+      const alreadyDeleted = new Set(result.orphanedItemIds);
+      const scopeItemIds = scopeItems
+        .map((r) => r.id)
+        .filter((id) => !alreadyDeleted.has(id));
+
+      if (scopeItemIds.length > 0) {
+        tx.delete(memoryEmbeddings)
+          .where(
+            and(
+              eq(memoryEmbeddings.targetType, "item"),
+              inArray(memoryEmbeddings.targetId, scopeItemIds),
+            ),
+          )
+          .run();
+        tx.delete(memoryItemSources)
+          .where(inArray(memoryItemSources.memoryItemId, scopeItemIds))
+          .run();
+        tx.delete(memoryItems)
+          .where(inArray(memoryItems.id, scopeItemIds))
+          .run();
+        result.orphanedItemIds.push(...scopeItemIds);
+      }
+
+      // Sweep memory summaries with this private scopeId.
+      const scopeSummaries = tx
+        .select({ id: memorySummaries.id })
+        .from(memorySummaries)
+        .where(eq(memorySummaries.scopeId, memoryScopeId))
+        .all();
+      const scopeSummaryIds = scopeSummaries.map((r) => r.id);
+
+      if (scopeSummaryIds.length > 0) {
+        tx.delete(memoryEmbeddings)
+          .where(
+            and(
+              eq(memoryEmbeddings.targetType, "summary"),
+              inArray(memoryEmbeddings.targetId, scopeSummaryIds),
+            ),
+          )
+          .run();
+        tx.delete(memorySummaries)
+          .where(inArray(memorySummaries.id, scopeSummaryIds))
+          .run();
+        result.deletedSummaryIds.push(...scopeSummaryIds);
+      }
+
+      // Sweep conversation starters with this private scopeId.
+      tx.delete(conversationStarters)
+        .where(eq(conversationStarters.scopeId, memoryScopeId))
+        .run();
+    }
+
     tx.delete(conversations).where(eq(conversations.id, id)).run();
   });
 
+  // Remove the conversation's disk-view directory after the DB transaction
+  if (createdAtForDiskCleanup != null) {
+    removeConversationDir(id, createdAtForDiskCleanup);
+  }
+
   return result;
 }
 
@@ -541,7 +897,10 @@ export function wipeConversation(id: string): WipeConversationResult {
   return {
     ...deletedMemoryIds,
     unsupersededItemIds,
-    deletedSummaryIds,
+    deletedSummaryIds: [
+      ...deletedSummaryIds,
+      ...deletedMemoryIds.deletedSummaryIds,
+    ],
     cancelledJobCount,
   };
 }
@@ -563,16 +922,25 @@ export function purgePrivateConversations(): {
     .all();
 
   if (privateConvs.length === 0) {
-    return { count: 0, deletedMemory: { segmentIds: [], orphanedItemIds: [] } };
+    return {
+      count: 0,
+      deletedMemory: {
+        segmentIds: [],
+        orphanedItemIds: [],
+        deletedSummaryIds: [],
+      },
+    };
   }
 
   const allSegmentIds: string[] = [];
   const allOrphanedItemIds: string[] = [];
+  const allDeletedSummaryIds: string[] = [];
 
   for (const conv of privateConvs) {
     const deleted = deleteConversation(conv.id);
     allSegmentIds.push(...deleted.segmentIds);
     allOrphanedItemIds.push(...deleted.orphanedItemIds);
+    allDeletedSummaryIds.push(...deleted.deletedSummaryIds);
   }
 
   return {
@@ -580,6 +948,7 @@ export function purgePrivateConversations(): {
     deletedMemory: {
       segmentIds: allSegmentIds,
       orphanedItemIds: allOrphanedItemIds,
+      deletedSummaryIds: allDeletedSummaryIds,
     },
   };
 }
@@ -662,6 +1031,13 @@ export async function addMessage(
       throw err;
     }
   }
+
+  // Mark the conversation dirty for delayed memory reduction. This runs
+  // after the insert transaction succeeds so the reducer knows which
+  // conversations have unprocessed messages. The helper preserves the
+  // earliest unreduced boundary (no-op when already dirty).
+  markConversationMemoryDirty(conversationId, messageId);
+
   const message = {
     id: messageId,
     conversationId,
@@ -759,6 +1135,12 @@ export function updateConversationTitle(
   const set: Record<string, unknown> = { title, updatedAt: Date.now() };
   if (isAutoTitle !== undefined) set.isAutoTitle = isAutoTitle;
   db.update(conversations).set(set).where(eq(conversations.id, id)).run();
+
+  // Update disk view meta.json with the new title
+  const conv = getConversation(id);
+  if (conv) {
+    updateMetaFile(conv);
+  }
 }
 
 export function updateConversationUsage(
@@ -864,6 +1246,14 @@ export function clearAll(): { conversations: number; messages: number } {
     );
   }
 
+  // Clear the disk-view conversations directory and recreate it empty
+  try {
+    rmSync(getConversationsDir(), { recursive: true, force: true });
+    mkdirSync(getConversationsDir(), { recursive: true });
+  } catch (err) {
+    log.warn({ err }, "clearAll: failed to reset conversations directory");
+  }
+
   return { conversations: convCount, messages: msgCount };
 }
 
@@ -942,11 +1332,11 @@ export function deleteLastExchange(conversationId: string): number {
 export interface DeletedMemoryIds {
   segmentIds: string[];
   orphanedItemIds: string[];
+  deletedSummaryIds: string[];
 }
 
 export interface WipeConversationResult extends DeletedMemoryIds {
   unsupersededItemIds: string[];
-  deletedSummaryIds: string[];
   cancelledJobCount: number;
 }
 
@@ -1012,7 +1402,11 @@ export function relinkAttachments(
  */
 export function deleteMessageById(messageId: string): DeletedMemoryIds {
   const db = getDb();
-  const result: DeletedMemoryIds = { segmentIds: [], orphanedItemIds: [] };
+  const result: DeletedMemoryIds = {
+    segmentIds: [],
+    orphanedItemIds: [],
+    deletedSummaryIds: [],
+  };
 
   // Collect attachment IDs linked to this message before cascade-delete
   // so we can scope orphan cleanup to only those candidates.
@@ -1100,6 +1494,28 @@ export function deleteMessageById(messageId: string): DeletedMemoryIds {
   return result;
 }
 
+/**
+ * Mark a conversation as having unreduced messages starting from the given
+ * message. Sets `memoryDirtyTailSinceMessageId` only when it is currently
+ * null so the earliest unreduced boundary is preserved across multiple
+ * messages — later messages must not clobber the original dirty marker.
+ */
+export function markConversationMemoryDirty(
+  conversationId: string,
+  messageId: string,
+): void {
+  const db = getDb();
+  db.update(conversations)
+    .set({ memoryDirtyTailSinceMessageId: messageId })
+    .where(
+      and(
+        eq(conversations.id, conversationId),
+        isNull(conversations.memoryDirtyTailSinceMessageId),
+      ),
+    )
+    .run();
+}
+
 export function setConversationOriginChannelIfUnset(
   conversationId: string,
   channel: ChannelId,
@@ -1232,3 +1648,170 @@ export function getDisplayMetaForConversations(
   }
   return result;
 }
+
+// ── Turn boundary resolution ─────────────────────────────────────────
+
+/**
+ * Returns `true` if a message is a tool-result user message — i.e. its
+ * role is "user" and its content is a JSON array where every block has
+ * `type === "tool_result"`. These synthetic user messages are injected
+ * between assistant messages within a single agent turn and should NOT
+ * be treated as turn boundaries.
+ */
+function isToolResultMessage(role: string, content: string): boolean {
+  if (role !== "user") return false;
+  try {
+    const parsed = JSON.parse(content);
+    if (!Array.isArray(parsed) || parsed.length === 0) return false;
+    return parsed.every(
+      (block: unknown) =>
+        block != null &&
+        typeof block === "object" &&
+        (block as Record<string, unknown>).type === "tool_result",
+    );
+  } catch {
+    return false;
+  }
+}
+
+/**
+ * Resolve all assistant message IDs that belong to the same agent turn
+ * as the given `messageId`. A "turn" is bounded by:
+ *   - The start of the conversation, or
+ *   - A user message whose content is NOT a tool_result array.
+ *
+ * Within a multi-step agent loop, the pattern is:
+ *   user msg → assistant A1 → user (tool_result) → assistant A2 → ...
+ * All assistant messages from A1 through the queried message (and beyond,
+ * up to the next real user message) are part of the same turn.
+ *
+ * Returns `[messageId]` as a fallback if the message is not found,
+ * preserving backward compatibility for callers.
+ */
+export function getAssistantMessageIdsInTurn(messageId: string): string[] {
+  const db = getDb();
+
+  // Look up the target message to get its conversationId and createdAt.
+  const target = getMessageById(messageId);
+  if (!target) return [messageId];
+
+  // Walk backward from the target message to find the turn boundary.
+  // Limit to 50 rows — sufficient for even aggressive tool-use loops.
+  const backwardRows = db
+    .select({
+      id: messages.id,
+      role: messages.role,
+      content: messages.content,
+      createdAt: messages.createdAt,
+    })
+    .from(messages)
+    .where(
+      and(
+        eq(messages.conversationId, target.conversationId),
+        lte(messages.createdAt, target.createdAt),
+      ),
+    )
+    .orderBy(desc(messages.createdAt))
+    .limit(50)
+    .all();
+
+  const assistantIds: string[] = [];
+  let boundaryCreatedAt: number | null = null;
+
+  for (const row of backwardRows) {
+    if (row.role === "assistant") {
+      assistantIds.push(row.id);
+    } else if (row.role === "user") {
+      if (isToolResultMessage(row.role, row.content)) {
+        // Tool-result user message — still within the same turn, continue.
+        continue;
+      }
+      // Real user message — this is the turn boundary.
+      boundaryCreatedAt = row.createdAt;
+      break;
+    }
+  }
+
+  // Walk forward from the target to collect any later assistant messages
+  // still within the same turn (e.g. when querying an intermediate
+  // message like A1 in a multi-step turn A1 → tool_result → A2).
+  const forwardRows = db
+    .select({
+      id: messages.id,
+      role: messages.role,
+      content: messages.content,
+      createdAt: messages.createdAt,
+    })
+    .from(messages)
+    .where(
+      and(
+        eq(messages.conversationId, target.conversationId),
+        gt(messages.createdAt, target.createdAt),
+      ),
+    )
+    .orderBy(asc(messages.createdAt))
+    .limit(50)
+    .all();
+
+  for (const row of forwardRows) {
+    if (row.role === "assistant") {
+      if (!assistantIds.includes(row.id)) {
+        assistantIds.push(row.id);
+      }
+    } else if (row.role === "user") {
+      if (isToolResultMessage(row.role, row.content)) {
+        // Tool-result user message — still within the same turn.
+        continue;
+      }
+      // Real user message — end of the turn.
+      break;
+    }
+  }
+
+  // Also query forward from the backward-walk boundary to pick up any
+  // assistant messages between the boundary and the target that may have
+  // been missed (e.g. due to the 50-row limit in the backward walk).
+  if (boundaryCreatedAt != null) {
+    const gapRows = db
+      .select({
+        id: messages.id,
+        role: messages.role,
+        createdAt: messages.createdAt,
+      })
+      .from(messages)
+      .where(
+        and(
+          eq(messages.conversationId, target.conversationId),
+          gt(messages.createdAt, boundaryCreatedAt),
+          lte(messages.createdAt, target.createdAt),
+        ),
+      )
+      .orderBy(asc(messages.createdAt))
+      .all();
+
+    for (const row of gapRows) {
+      if (row.role === "assistant" && !assistantIds.includes(row.id)) {
+        assistantIds.push(row.id);
+      }
+    }
+  }
+
+  // Sort by createdAt to ensure stable ordering.
+  // Re-fetch createdAt for all collected IDs so the sort is accurate.
+  if (assistantIds.length <= 1) return assistantIds;
+
+  const idSet = new Set(assistantIds);
+  const sorted = db
+    .select({ id: messages.id, createdAt: messages.createdAt })
+    .from(messages)
+    .where(
+      and(
+        eq(messages.conversationId, target.conversationId),
+        inArray(messages.id, [...idSet]),
+      ),
+    )
+    .orderBy(asc(messages.createdAt))
+    .all();
+
+  return sorted.map((r) => r.id);
+}