@agent-native/core 0.137.7 → 0.138.0

package/corpus/templates/content/actions/delete-document.ts

@@ -1,16 +1,38 @@
 import { defineAction } from "@agent-native/core";
 import { writeAppState } from "@agent-native/core/application-state";
 import { assertAccess } from "@agent-native/core/sharing";
-import { and, eq, inArray, isNotNull, isNull, ne, or, sql } from "drizzle-orm";
+import { and, eq, inArray, isNotNull, isNull, ne, or } from "drizzle-orm";
 import { z } from "zod";
 
 import { getDb, schema } from "../server/db/index.js";
 import { chunks } from "./_batch-utils.js";
+import {
+  lockContentDatabaseMutation,
+  touchContentDatabase,
+} from "./_content-database-mutation-lock.js";
 import { assertNotWorkspaceCatalogDocuments } from "./_content-space-catalog-guards.js";
+import { lockDatabaseMemberships } from "./_database-membership-lock.js";
 import { renumberDatabaseRows } from "./_database-row-batch.js";
 
 const DELETE_BATCH_SIZE = 90;
 
+export class PermanentDeleteScopeChangedError extends Error {
+  constructor() {
+    super("Document deletion scope changed; retry deletion.");
+  }
+}
+
+type PermanentDeleteScope = {
+  documentIds: string[];
+  ownedDatabaseIds: string[];
+};
+
+function hasSameIds(left: string[], right: string[]) {
+  if (left.length !== right.length) return false;
+  const rightIds = new Set(right);
+  return left.every((id) => rightIds.has(id));
+}
+
 async function selectDocumentChildren(
   db: ReturnType<typeof getDb>,
   parentIds: string[],
@@ -97,26 +119,75 @@ async function selectDatabaseItemDocuments(
   return ownedRows.map((row) => ({ documentId: row.id }));
 }
 
-async function selectMembershipDatabaseIds(
+async function selectDatabaseItemDocumentTrashState(
   db: ReturnType<typeof getDb>,
-  documentIds: string[],
+  databaseIds: string[],
   ownerEmail: string,
 ) {
-  const rows: Array<{ databaseId: string }> = [];
-  for (const batch of chunks(documentIds, DELETE_BATCH_SIZE)) {
+  const itemDocuments = await selectDatabaseItemDocuments(
+    db,
+    databaseIds,
+    ownerEmail,
+  );
+  const rows: Array<{
+    id: string;
+    trashedAt: string | null;
+  }> = [];
+  for (const batch of chunks(
+    itemDocuments.map((item) => item.documentId),
+    DELETE_BATCH_SIZE,
+  )) {
     rows.push(
       ...(await db
-        .select({ databaseId: schema.contentDatabaseItems.databaseId })
-        .from(schema.contentDatabaseItems)
+        .select({
+          id: schema.documents.id,
+          trashedAt: schema.documents.trashedAt,
+        })
+        .from(schema.documents)
         .where(
           and(
-            inArray(schema.contentDatabaseItems.documentId, batch),
-            eq(schema.contentDatabaseItems.ownerEmail, ownerEmail),
+            inArray(schema.documents.id, batch),
+            eq(schema.documents.ownerEmail, ownerEmail),
           ),
         )),
     );
   }
-  return [...new Set(rows.map((row) => row.databaseId))];
+  return rows;
+}
+
+async function selectMembershipsForDocuments(
+  db: ReturnType<typeof getDb>,
+  documentIds: string[],
+) {
+  const rows: Array<{ id: string; databaseId: string }> = [];
+  for (const batch of chunks(documentIds, DELETE_BATCH_SIZE)) {
+    rows.push(
+      ...(await db
+        .select({
+          id: schema.contentDatabaseItems.id,
+          databaseId: schema.contentDatabaseItems.databaseId,
+        })
+        .from(schema.contentDatabaseItems)
+        .where(inArray(schema.contentDatabaseItems.documentId, batch))),
+    );
+  }
+  return rows;
+}
+
+async function selectMembershipIdsForDatabases(
+  db: ReturnType<typeof getDb>,
+  databaseIds: string[],
+) {
+  const rows: Array<{ id: string }> = [];
+  for (const batch of chunks(databaseIds, DELETE_BATCH_SIZE)) {
+    rows.push(
+      ...(await db
+        .select({ id: schema.contentDatabaseItems.id })
+        .from(schema.contentDatabaseItems)
+        .where(inArray(schema.contentDatabaseItems.databaseId, batch))),
+    );
+  }
+  return rows.map((row) => row.id);
 }
 
 async function collectDocumentSubtreeForDelete(
@@ -172,84 +243,117 @@ async function collectDocumentSubtreeForDelete(
     frontier = [...next];
   }
 
-  const collectedDocumentIds = [...documentIds];
-  const collectedOwnedDatabaseIds = [...ownedDatabaseIds];
   return {
-    documentIds: collectedDocumentIds,
-    ownedDatabaseIds: collectedOwnedDatabaseIds,
-    lockDatabaseIds: [
-      ...new Set([
-        ...collectedOwnedDatabaseIds,
-        ...(await selectMembershipDatabaseIds(
-          db,
-          collectedDocumentIds,
-          ownerEmail,
-        )),
-      ]),
-    ],
+    documentIds: [...documentIds],
+    ownedDatabaseIds: [...ownedDatabaseIds],
   };
 }
 
-async function lockDatabasesAndRecollect<
-  T extends { lockDatabaseIds: string[] },
->(
+async function lockPermanentDeleteScope(
   db: ReturnType<typeof getDb>,
-  ownerEmail: string,
-  collect: () => Promise<T>,
-): Promise<T> {
-  const lockedDatabaseIds = new Set<string>();
-  let collected = await collect();
-  while (true) {
-    const unlockedDatabaseIds = collected.lockDatabaseIds.filter(
-      (databaseId) => !lockedDatabaseIds.has(databaseId),
-    );
-    if (unlockedDatabaseIds.length === 0) return collected;
-    for (const batch of chunks(unlockedDatabaseIds, DELETE_BATCH_SIZE)) {
-      await db
-        .update(schema.contentDatabases)
-        .set({ updatedAt: sql`${schema.contentDatabases.updatedAt}` })
-        .where(
-          and(
-            inArray(schema.contentDatabases.id, batch),
-            eq(schema.contentDatabases.ownerEmail, ownerEmail),
-          ),
-        );
-      for (const databaseId of batch) lockedDatabaseIds.add(databaseId);
-    }
-    collected = await collect();
+  collectScope: () => Promise<PermanentDeleteScope>,
+) {
+  const initialScope = await collectScope();
+  const initialMemberships = await selectMembershipsForDocuments(
+    db,
+    initialScope.documentIds,
+  );
+  const lockedDatabaseIds = [
+    ...new Set([
+      ...initialScope.ownedDatabaseIds,
+      ...initialMemberships.map((membership) => membership.databaseId),
+    ]),
+  ].sort();
+  for (const databaseId of lockedDatabaseIds) {
+    await lockContentDatabaseMutation(db, databaseId);
+  }
+
+  const reloadedScope = await collectScope();
+  const reloadedMemberships = await selectMembershipsForDocuments(
+    db,
+    reloadedScope.documentIds,
+  );
+  const lockedDatabaseIdSet = new Set(lockedDatabaseIds);
+  const uncoveredDatabaseId = [
+    ...reloadedScope.ownedDatabaseIds,
+    ...reloadedMemberships.map((membership) => membership.databaseId),
+  ].find((databaseId) => !lockedDatabaseIdSet.has(databaseId));
+  if (uncoveredDatabaseId) {
+    throw new PermanentDeleteScopeChangedError();
+  }
+
+  const membershipIds = await selectMembershipIdsForDatabases(
+    db,
+    lockedDatabaseIds,
+  );
+  await lockDatabaseMemberships(db, membershipIds);
+
+  const lockedScope = await collectScope();
+  const lockedMemberships = await selectMembershipsForDocuments(
+    db,
+    lockedScope.documentIds,
+  );
+  const lockedMembershipIds = await selectMembershipIdsForDatabases(
+    db,
+    lockedDatabaseIds,
+  );
+  if (
+    !hasSameIds(reloadedScope.documentIds, lockedScope.documentIds) ||
+    !hasSameIds(reloadedScope.ownedDatabaseIds, lockedScope.ownedDatabaseIds) ||
+    !hasSameIds(membershipIds, lockedMembershipIds) ||
+    [
+      ...lockedScope.ownedDatabaseIds,
+      ...lockedMemberships.map((membership) => membership.databaseId),
+    ].some((databaseId) => !lockedDatabaseIdSet.has(databaseId))
+  ) {
+    throw new PermanentDeleteScopeChangedError();
   }
+  return lockedScope;
 }
 
-export async function trashDocumentSubtree(
+export async function lockDatabasesForTrash(
   db: ReturnType<typeof getDb>,
   id: string,
   ownerEmail: string,
-  trashedAt = new Date().toISOString(),
-): Promise<string[]> {
-  return db.transaction((tx) =>
-    trashDocumentSubtreeInTransaction(
-      tx as unknown as ReturnType<typeof getDb>,
-      id,
-      ownerEmail,
-      trashedAt,
-    ),
+) {
+  const subtree = await collectDocumentSubtreeForDelete(db, id, ownerEmail);
+  const memberships = await selectMembershipsForDocuments(
+    db,
+    subtree.documentIds,
   );
+  const databaseIds = [
+    ...new Set([
+      ...memberships.map((membership) => membership.databaseId),
+      ...subtree.ownedDatabaseIds,
+    ]),
+  ].sort();
+  for (const databaseId of databaseIds) {
+    await lockContentDatabaseMutation(db, databaseId);
+  }
+  return new Set(databaseIds);
 }
 
-async function trashDocumentSubtreeInTransaction(
+export async function trashDocumentSubtree(
   db: ReturnType<typeof getDb>,
   id: string,
   ownerEmail: string,
-  trashedAt: string,
+  trashedAt = new Date().toISOString(),
+  lockedDatabaseIds?: ReadonlySet<string>,
 ): Promise<string[]> {
-  // Stable-key upserts lock their canonical database row before creating or
-  // updating children. Acquire the same locks, then collect again: an upsert
-  // that won the lock before this transaction may have added a child after the
-  // first traversal, while one that arrives later must wait until the database
-  // has been marked deleted and will fail closed.
-  const { documentIds } = await lockDatabasesAndRecollect(db, ownerEmail, () =>
-    collectDocumentSubtreeForDelete(db, id, ownerEmail),
-  );
+  const { documentIds, ownedDatabaseIds } =
+    await collectDocumentSubtreeForDelete(db, id, ownerEmail);
+  if (lockedDatabaseIds) {
+    const memberships = await selectMembershipsForDocuments(db, documentIds);
+    const unlockedDatabaseId = [
+      ...new Set([
+        ...ownedDatabaseIds,
+        ...memberships.map((membership) => membership.databaseId),
+      ]),
+    ].find((databaseId) => !lockedDatabaseIds.has(databaseId));
+    if (unlockedDatabaseId) {
+      throw new Error("Document subtree changed; retry deletion.");
+    }
+  }
   await assertNotWorkspaceCatalogDocuments(db, documentIds, "deleted");
 
   const independentlyTrashedDatabaseDocumentIds = new Set<string>();
@@ -291,6 +395,39 @@ async function trashDocumentSubtreeInTransaction(
     );
   }
 
+  const activeMemberships = await selectMembershipsForDocuments(
+    db,
+    activeDocumentIds,
+  );
+  const transitioningOwnedDatabases: Array<{ id: string }> = [];
+  for (const batch of chunks(activeDocumentIds, DELETE_BATCH_SIZE)) {
+    transitioningOwnedDatabases.push(
+      ...(await db
+        .select({ id: schema.contentDatabases.id })
+        .from(schema.contentDatabases)
+        .where(
+          and(
+            inArray(schema.contentDatabases.documentId, batch),
+            eq(schema.contentDatabases.ownerEmail, ownerEmail),
+            isNull(schema.contentDatabases.deletedAt),
+          ),
+        )),
+    );
+  }
+  const transitioningOwnedDatabaseIds = new Set(
+    transitioningOwnedDatabases.map((database) => database.id),
+  );
+  const externalDatabaseIds = [
+    ...new Set(
+      activeMemberships
+        .map((membership) => membership.databaseId)
+        .filter((databaseId) => !transitioningOwnedDatabaseIds.has(databaseId)),
+    ),
+  ].sort();
+  for (const databaseId of externalDatabaseIds) {
+    await touchContentDatabase(db, databaseId, trashedAt);
+  }
+
   for (const batch of chunks(activeDocumentIds, DELETE_BATCH_SIZE)) {
     await db
       .update(schema.documents)
@@ -322,19 +459,55 @@ export async function restoreDocumentSubtree(
   rootId: string,
   ownerEmail: string,
 ): Promise<string[]> {
-  const documentIds = (
-    await db
-      .select({ id: schema.documents.id })
-      .from(schema.documents)
-      .where(
-        and(
-          eq(schema.documents.trashRootId, rootId),
-          eq(schema.documents.ownerEmail, ownerEmail),
-        ),
-      )
-  ).map((document) => document.id);
-  if (documentIds.length === 0) return [];
+  const collectRestoreScope = async () => {
+    const documentIds = (
+      await db
+        .select({ id: schema.documents.id })
+        .from(schema.documents)
+        .where(
+          and(
+            eq(schema.documents.trashRootId, rootId),
+            eq(schema.documents.ownerEmail, ownerEmail),
+          ),
+        )
+    ).map((document) => document.id);
+    const memberships = await selectMembershipsForDocuments(db, documentIds);
+    const ownedDatabaseIds = (
+      await selectOwnedDatabaseIds(db, documentIds, ownerEmail)
+    ).map((database) => database.id);
+    return { documentIds, memberships, ownedDatabaseIds };
+  };
+
+  const initialScope = await collectRestoreScope();
+  if (initialScope.documentIds.length === 0) return [];
+  const lockedDatabaseIds = [
+    ...new Set([
+      ...initialScope.ownedDatabaseIds,
+      ...initialScope.memberships.map((membership) => membership.databaseId),
+    ]),
+  ].sort();
+  for (const databaseId of lockedDatabaseIds) {
+    await lockContentDatabaseMutation(db, databaseId);
+  }
 
+  const restoreScope = await collectRestoreScope();
+  const lockedDatabaseIdSet = new Set(lockedDatabaseIds);
+  const unlockedDatabaseId = [
+    ...new Set([
+      ...restoreScope.ownedDatabaseIds,
+      ...restoreScope.memberships.map((membership) => membership.databaseId),
+    ]),
+  ].find((databaseId) => !lockedDatabaseIdSet.has(databaseId));
+  if (unlockedDatabaseId) {
+    throw new Error("Document restore scope changed; retry restoration.");
+  }
+  await lockDatabaseMemberships(
+    db,
+    await selectMembershipIdsForDatabases(db, lockedDatabaseIds),
+  );
+
+  const documentIds = restoreScope.documentIds;
+  if (documentIds.length === 0) return [];
   const now = new Date().toISOString();
   for (const batch of chunks(documentIds, DELETE_BATCH_SIZE)) {
     await db
@@ -399,20 +572,49 @@ export async function deleteDocumentRecursive(
   id: string,
   ownerEmail: string,
 ): Promise<string[]> {
-  return db.transaction(async (tx) => {
-    const scopedDb = tx as unknown as ReturnType<typeof getDb>;
-    const { documentIds, ownedDatabaseIds } = await lockDatabasesAndRecollect(
-      scopedDb,
-      ownerEmail,
-      () => collectDocumentSubtreeForDelete(scopedDb, id, ownerEmail),
-    );
-    return deleteCollectedDocuments(
-      scopedDb,
-      documentIds,
-      ownedDatabaseIds,
+  return db.transaction((tx) =>
+    deleteDocumentRootsRecursive(
+      tx as unknown as ReturnType<typeof getDb>,
+      [id],
       ownerEmail,
-    );
-  });
+    ),
+  );
+}
+
+export async function deleteDocumentRootsRecursive(
+  db: ReturnType<typeof getDb>,
+  rootIds: string[],
+  ownerEmail: string,
+): Promise<string[]> {
+  const collectScope = async () => {
+    const documentIds = new Set<string>();
+    const ownedDatabaseIds = new Set<string>();
+    for (const rootId of rootIds) {
+      const scope = await collectDocumentSubtreeForDelete(
+        db,
+        rootId,
+        ownerEmail,
+      );
+      for (const documentId of scope.documentIds) documentIds.add(documentId);
+      for (const databaseId of scope.ownedDatabaseIds) {
+        ownedDatabaseIds.add(databaseId);
+      }
+    }
+    return {
+      documentIds: [...documentIds],
+      ownedDatabaseIds: [...ownedDatabaseIds],
+    };
+  };
+  const { documentIds, ownedDatabaseIds } = await lockPermanentDeleteScope(
+    db,
+    collectScope,
+  );
+  return deleteCollectedDocuments(
+    db,
+    documentIds,
+    ownedDatabaseIds,
+    ownerEmail,
+  );
 }
 
 async function deleteCollectedDocuments(
@@ -515,6 +717,7 @@ async function deleteCollectedDocuments(
         ),
       );
   });
+
   await deleteWhereIn(documentIds, async (documentIdBatch) => {
     await db
       .delete(schema.contentDatabaseItemKeyClaims)
@@ -570,6 +773,17 @@ async function deleteCollectedDocuments(
     await db
       .delete(schema.contentDatabases)
       .where(inArray(schema.contentDatabases.id, databaseIdBatch));
+    // Receipts deliberately have no database foreign key. Removing them after
+    // the database row closes the race with a migration that already holds the
+    // row lock and commits its receipt before this deletion can continue.
+    await db
+      .delete(schema.contentDatabaseMigrationReceipts)
+      .where(
+        inArray(
+          schema.contentDatabaseMigrationReceipts.databaseId,
+          databaseIdBatch,
+        ),
+      );
   });
 
   await deleteWhereIn(documentIds, async (documentIdBatch) => {
@@ -629,74 +843,63 @@ export async function deleteTrashedDocumentSubtree(
   id: string,
   ownerEmail: string,
 ): Promise<string[]> {
-  return db.transaction((tx) =>
-    deleteTrashedDocumentSubtreeInTransaction(
-      tx as unknown as ReturnType<typeof getDb>,
-      id,
-      ownerEmail,
-    ),
-  );
-}
+  const collectScope = async () => {
+    const [root] = await db
+      .select({ id: schema.documents.id })
+      .from(schema.documents)
+      .where(
+        and(
+          eq(schema.documents.id, id),
+          eq(schema.documents.ownerEmail, ownerEmail),
+          eq(schema.documents.trashRootId, id),
+          isNotNull(schema.documents.trashedAt),
+        ),
+      )
+      .limit(1);
+    if (!root) {
+      throw new Error(
+        "Document must be in Trash and be a Trash root before permanent deletion",
+      );
+    }
 
-async function deleteTrashedDocumentSubtreeInTransaction(
-  db: ReturnType<typeof getDb>,
-  id: string,
-  ownerEmail: string,
-): Promise<string[]> {
-  const [root] = await db
-    .select({ id: schema.documents.id })
-    .from(schema.documents)
-    .where(
-      and(
-        eq(schema.documents.id, id),
-        eq(schema.documents.ownerEmail, ownerEmail),
-        eq(schema.documents.trashRootId, id),
-        isNotNull(schema.documents.trashedAt),
-      ),
-    )
-    .limit(1);
-  if (!root) {
-    throw new Error(
-      "Document must be in Trash and be a Trash root before permanent deletion",
+    const documentIds = (
+      await db
+        .select({ id: schema.documents.id })
+        .from(schema.documents)
+        .where(
+          and(
+            eq(schema.documents.ownerEmail, ownerEmail),
+            eq(schema.documents.trashRootId, id),
+            isNotNull(schema.documents.trashedAt),
+          ),
+        )
+    ).map((document) => document.id);
+    const ownedDatabaseIds = await selectOwnedDatabaseIds(
+      db,
+      documentIds,
+      ownerEmail,
+    ).then((rows) => rows.map((database) => database.id));
+    const documentIdSet = new Set(documentIds);
+    const activeOutsideScope = (
+      await selectDatabaseItemDocumentTrashState(
+        db,
+        ownedDatabaseIds,
+        ownerEmail,
+      )
+    ).find(
+      (document) => !documentIdSet.has(document.id) && !document.trashedAt,
     );
-  }
+    if (activeOutsideScope) {
+      throw new Error(
+        "Database contains an active row outside this Trash item",
+      );
+    }
+    return { documentIds, ownedDatabaseIds };
+  };
 
-  const { documentIds, ownedDatabaseIds } = await lockDatabasesAndRecollect(
+  const { documentIds, ownedDatabaseIds } = await lockPermanentDeleteScope(
     db,
-    ownerEmail,
-    async () => {
-      const collectedDocumentIds = (
-        await db
-          .select({ id: schema.documents.id })
-          .from(schema.documents)
-          .where(
-            and(
-              eq(schema.documents.ownerEmail, ownerEmail),
-              eq(schema.documents.trashRootId, id),
-              isNotNull(schema.documents.trashedAt),
-            ),
-          )
-      ).map((document) => document.id);
-      const collectedDatabaseIds = await selectOwnedDatabaseIds(
-        db,
-        collectedDocumentIds,
-        ownerEmail,
-      ).then((rows) => rows.map((database) => database.id));
-      return {
-        documentIds: collectedDocumentIds,
-        ownedDatabaseIds: collectedDatabaseIds,
-        lockDatabaseIds: [
-          ...new Set([
-            ...collectedDatabaseIds,
-            ...(await selectMembershipDatabaseIds(
-              db,
-              collectedDocumentIds,
-              ownerEmail,
-            )),
-          ]),
-        ],
-      };
-    },
+    collectScope,
   );
 
   await db
@@ -760,9 +963,19 @@ export default defineAction({
         if (!membership) {
           throw new Error("Document is not part of Favorites");
         }
-        await db
-          .delete(schema.contentDatabaseItems)
-          .where(eq(schema.contentDatabaseItems.id, membership.id));
+        await db.transaction(async (tx) => {
+          await lockContentDatabaseMutation(
+            tx as unknown as ReturnType<typeof getDb>,
+            contextDatabase.id,
+          );
+          await touchContentDatabase(
+            tx as unknown as ReturnType<typeof getDb>,
+            contextDatabase.id,
+          );
+          await tx
+            .delete(schema.contentDatabaseItems)
+            .where(eq(schema.contentDatabaseItems.id, membership.id));
+        });
         await writeAppState("refresh-signal", { ts: Date.now() });
         return { success: true, deleted: 0, removed: 1 };
       }
@@ -777,11 +990,21 @@ export default defineAction({
     if (systemDatabase?.systemRole) {
       throw new Error("System Content database documents cannot be deleted");
     }
-    const deleted = await trashDocumentSubtree(
-      db,
-      id,
-      existing.ownerEmail as string,
-    );
+    const deleted = await db.transaction(async (tx) => {
+      const transactionDb = tx as unknown as ReturnType<typeof getDb>;
+      const lockedDatabaseIds = await lockDatabasesForTrash(
+        transactionDb,
+        id,
+        existing.ownerEmail as string,
+      );
+      return trashDocumentSubtree(
+        transactionDb,
+        id,
+        existing.ownerEmail as string,
+        undefined,
+        lockedDatabaseIds,
+      );
+    });
 
     await writeAppState("refresh-signal", { ts: Date.now() });