@hiai-gg/docsmint 0.6.8 → 0.7.1

package/backend/src/lib/lifecycle-service.ts

@@ -1,5 +1,7 @@
 import {
 	accounts,
+	apiKeys,
+	attachmentStorageCleanupOutbox,
 	attachments,
 	auditLog,
 	categories,
@@ -7,6 +9,7 @@ import {
 	documents,
 	folders,
 	lifecycleOperations,
+	pendingAttachmentUploads,
 	sessions,
 	shareLinks,
 	tags,
@@ -23,9 +26,21 @@ import type {
 	UserDataExportRecord,
 	UserDataLifecycle,
 } from "@hiai-docs/sdk";
-import { and, eq, inArray, lt, or, sql } from "drizzle-orm";
+import { and, asc, eq, inArray, isNull, lt, or, sql } from "drizzle-orm";
+import {
+	drainAttachmentStorageCleanupOutbox,
+	stageAttachmentStorageCleanup,
+} from "./attachment-storage-cleanup";
+import {
+	type PendingAttachmentUploadRow,
+	stagePendingAttachmentCleanup,
+} from "./attachment-upload-cleanup";
+import { acquireDocumentPipelineLocks } from "./document-pipeline-serialization";
+import type { RuntimeAttachmentQuotaAdmission } from "./runtime-options";
+import { getDocsMintRuntimeOptions } from "./runtime-options";
 
 const LEASE_MS = 60_000;
+export const LIFECYCLE_CLEANUP_PAGE_SIZE = 100;
 
 type PersistentOperation = typeof lifecycleOperations.$inferSelect;
 
@@ -49,6 +64,16 @@ export class LifecycleLeaseLostError extends Error {
 	}
 }
 
+/** An issued direct-upload URL remains usable until its fixed expiry. */
+export class LifecyclePendingAttachmentUploadsError extends Error {
+	readonly code = "pending_attachment_uploads";
+
+	constructor() {
+		super("Pending attachment upload URLs have not expired");
+		this.name = "LifecyclePendingAttachmentUploadsError";
+	}
+}
+
 export type LifecycleRuntimeAdapters = Readonly<{
 	/** Checks the host-owned fence immediately before the first mutation. */
 	verifyPurgeFence: (
@@ -75,6 +100,8 @@ export type LifecycleRuntimeAdapters = Readonly<{
 		documentIds: readonly string[],
 		signal?: AbortSignal,
 	) => Promise<number>;
+	/** Generic OSS quota saga used while settling workspace attachment cleanup. */
+	attachmentStorageQuotaAdmission?: RuntimeAttachmentQuotaAdmission;
 }>;
 
 /**
@@ -116,6 +143,8 @@ function safeErrorCode(error: unknown): string {
 		return "aborted";
 	if (error instanceof LifecycleFenceRejectedError) return error.code;
 	if (error instanceof LifecycleLeaseLostError) return error.code;
+	if (error instanceof LifecyclePendingAttachmentUploadsError)
+		return error.code;
 	return "persistence_failed";
 }
 
@@ -509,7 +538,17 @@ export function createPersistentLifecycleService(
 			const deletedByDomain = operationCounts(operation);
 			try {
 				throwIfAborted(ctx.signal);
+				await runtime.verifyPurgeFence(ctx, gate.fenceToken);
 				await withActor(ctx.actorUserId, async (tx) => {
+					await tx.execute(
+						sql`SELECT public.acquire_account_purge_fence_lock(${ctx.actorUserId}::uuid)`,
+					);
+					const [actor] = await tx
+						.select({ id: users.id })
+						.from(users)
+						.where(eq(users.id, ctx.actorUserId))
+						.limit(1);
+					if (!actor) throw new LifecycleLeaseLostError();
 					const updated = await tx
 						.update(lifecycleOperations)
 						.set({
@@ -525,10 +564,21 @@ export function createPersistentLifecycleService(
 						.returning({ id: lifecycleOperations.id });
 					requireLeaseWrite(updated);
 				});
-				await runtime.verifyPurgeFence(ctx, gate.fenceToken);
+				const withCleanupActor = <T>(
+					operationAction: (tx: TenantTransaction) => Promise<T>,
+				) =>
+					withActor(ctx.actorUserId, async (tx) => {
+						await tx.execute(
+							sql`SELECT set_config('app.lifecycle_operation_id', ${operation.id}, true)`,
+						);
+						await tx.execute(
+							sql`SELECT set_config('app.lifecycle_lease_owner', ${leaseOwner}, true)`,
+						);
+						return operationAction(tx);
+					});
 				const dbDelete = <T extends { id: string }>(
 					action: (tx: TenantTransaction) => Promise<T[]>,
-				) => withActor(ctx.actorUserId, action).then((rows) => rows.length);
+				) => withCleanupActor(action).then((rows) => rows.length);
 				await runStep(
 					operation,
 					ctx.actorUserId,
@@ -537,6 +587,126 @@ export function createPersistentLifecycleService(
 					deletedByDomain,
 					() => runtime.cancelAccountJobs(ctx.actorUserId, ctx.signal),
 				);
+				await runStep(
+					operation,
+					ctx.actorUserId,
+					leaseOwner,
+					"settle_pending_attachment_uploads",
+					deletedByDomain,
+					async () => {
+						let staged = 0;
+						for (;;) {
+							const page = await withCleanupActor(async (tx) => {
+								const pending = await tx
+									.select({
+										id: pendingAttachmentUploads.id,
+										documentId: pendingAttachmentUploads.documentId,
+										ownerUserId: sql<string>`public.lifecycle_child_document_owner(
+											${pendingAttachmentUploads.documentId},
+											${ctx.actorUserId}::uuid
+										)`,
+										actorUserId: pendingAttachmentUploads.actorUserId,
+										workspaceId: pendingAttachmentUploads.workspaceId,
+										storageKey: pendingAttachmentUploads.storageKey,
+										tokenHash: pendingAttachmentUploads.tokenHash,
+										filename: pendingAttachmentUploads.filename,
+										mimeType: pendingAttachmentUploads.mimeType,
+										declaredSize: pendingAttachmentUploads.declaredSize,
+										quotaReservationId:
+											pendingAttachmentUploads.quotaReservationId,
+										quotaOperationKey:
+											pendingAttachmentUploads.quotaOperationKey,
+										quotaState: pendingAttachmentUploads.quotaState,
+										actualSize: pendingAttachmentUploads.actualSize,
+										urlIssuedAt: pendingAttachmentUploads.urlIssuedAt,
+										expiresAt: pendingAttachmentUploads.expiresAt,
+										leaseOwner: pendingAttachmentUploads.leaseOwner,
+									})
+									.from(pendingAttachmentUploads)
+									.where(
+										eq(pendingAttachmentUploads.actorUserId, ctx.actorUserId),
+									)
+									.orderBy(asc(pendingAttachmentUploads.id))
+									.limit(LIFECYCLE_CLEANUP_PAGE_SIZE);
+								for (const row of pending) {
+									await stagePendingAttachmentCleanup(
+										tx,
+										row as PendingAttachmentUploadRow,
+										ctx.actorUserId,
+									);
+								}
+								return pending.length;
+							});
+							staged += page;
+							if (page === 0) break;
+						}
+						const leftover = await withCleanupActor((tx) =>
+							tx
+								.select({ id: attachmentStorageCleanupOutbox.id })
+								.from(attachmentStorageCleanupOutbox)
+								.where(
+									or(
+										eq(
+											attachmentStorageCleanupOutbox.actorUserId,
+											ctx.actorUserId,
+										),
+										eq(
+											attachmentStorageCleanupOutbox.ownerUserId,
+											ctx.actorUserId,
+										),
+										eq(
+											attachmentStorageCleanupOutbox.requestedByUserId,
+											ctx.actorUserId,
+										),
+									),
+								)
+								.limit(1),
+						);
+						if (staged > 0 || leftover.length > 0) {
+							for (let page = 0; page < 10_000; page += 1) {
+								const result = await drainAttachmentStorageCleanupOutbox({
+									deleteObjects: runtime.deleteObjects,
+									actorUserId: ctx.actorUserId,
+									quotaAdmission:
+										runtime.attachmentStorageQuotaAdmission ??
+										getDocsMintRuntimeOptions()
+											?.attachmentStorageQuotaAdmission,
+									signal: ctx.signal,
+									pageSize: LIFECYCLE_CLEANUP_PAGE_SIZE,
+									maxPages: 1,
+								});
+								if (result.claimed === 0) break;
+								if (result.failed > 0 || result.deferred > 0)
+									throw new LifecyclePendingAttachmentUploadsError();
+							}
+						}
+						const retained = await withCleanupActor((tx) =>
+							tx
+								.select({ id: attachmentStorageCleanupOutbox.id })
+								.from(attachmentStorageCleanupOutbox)
+								.where(
+									or(
+										eq(
+											attachmentStorageCleanupOutbox.actorUserId,
+											ctx.actorUserId,
+										),
+										eq(
+											attachmentStorageCleanupOutbox.ownerUserId,
+											ctx.actorUserId,
+										),
+										eq(
+											attachmentStorageCleanupOutbox.requestedByUserId,
+											ctx.actorUserId,
+										),
+									),
+								)
+								.limit(1),
+						);
+						if (retained.length > 0)
+							throw new LifecyclePendingAttachmentUploadsError();
+						return staged;
+					},
+				);
 				await runStep(
 					operation,
 					ctx.actorUserId,
@@ -563,7 +733,12 @@ export function createPersistentLifecycleService(
 					tx
 						.select({ id: documents.id })
 						.from(documents)
-						.where(eq(documents.ownerId, ctx.actorUserId)),
+						.where(
+							and(
+								eq(documents.ownerId, ctx.actorUserId),
+								isNull(documents.workspaceId),
+							),
+						),
 				);
 				const documentIds = owned.map((row) => row.id);
 				await runStep(
@@ -573,14 +748,19 @@ export function createPersistentLifecycleService(
 					"remove_document_versions",
 					deletedByDomain,
 					() =>
-						documentIds.length
-							? dbDelete((tx) =>
-									tx
-										.delete(versions)
-										.where(inArray(versions.documentId, documentIds))
-										.returning({ id: versions.id }),
+						dbDelete((tx) =>
+							tx
+								.delete(versions)
+								.where(
+									documentIds.length
+										? or(
+												eq(versions.createdBy, ctx.actorUserId),
+												inArray(versions.documentId, documentIds),
+											)
+										: eq(versions.createdBy, ctx.actorUserId),
 								)
-							: Promise.resolve(0),
+								.returning({ id: versions.id }),
+						),
 				);
 				await runStep(
 					operation,
@@ -606,29 +786,183 @@ export function createPersistentLifecycleService(
 					deletedByDomain,
 					() => runtime.removeGraphState(documentIds, ctx.signal),
 				);
-				const objectRows = documentIds.length
-					? await withActor(ctx.actorUserId, (tx) =>
-							tx
-								.select({
-									id: attachments.id,
-									storageKey: attachments.storageKey,
-								})
-								.from(attachments)
-								.where(inArray(attachments.documentId, documentIds)),
-						)
-					: [];
 				await runStep(
 					operation,
 					ctx.actorUserId,
 					leaseOwner,
 					"delete_attachment_objects",
 					deletedByDomain,
+					async () => {
+						let removed = 0;
+						for (;;) {
+							const page = await withCleanupActor(async (tx) => {
+								const objectRows = await tx
+									.select({
+										id: attachments.id,
+										documentId: attachments.documentId,
+										storageKey: attachments.storageKey,
+										size: attachments.size,
+										uploadedBy: attachments.uploadedBy,
+										workspaceId: attachments.workspaceId,
+										ownerUserId: sql<string>`public.lifecycle_child_document_owner(
+											${attachments.documentId},
+											${ctx.actorUserId}::uuid
+										)`,
+									})
+									.from(attachments)
+									.where(
+										documentIds.length
+											? or(
+													eq(attachments.uploadedBy, ctx.actorUserId),
+													inArray(attachments.documentId, documentIds),
+												)
+											: eq(attachments.uploadedBy, ctx.actorUserId),
+									)
+									.orderBy(asc(attachments.id))
+									.limit(LIFECYCLE_CLEANUP_PAGE_SIZE);
+								if (objectRows.length === 0) return 0;
+								for (const row of objectRows) {
+									await stageAttachmentStorageCleanup(tx, {
+										sourceKind: "attachment",
+										sourceId: row.id,
+										storageKey: row.storageKey,
+										documentId: row.documentId,
+										actorUserId: row.uploadedBy,
+										ownerUserId: row.ownerUserId,
+										requestedByUserId: ctx.actorUserId,
+										workspaceId: row.workspaceId,
+										size: row.size,
+										quotaOperationKey: `attachment:${row.documentId}:${row.storageKey}`,
+										quotaReleaseKind: row.workspaceId ? "committed" : "none",
+									});
+								}
+								const deleted = await tx
+									.delete(attachments)
+									.where(
+										inArray(
+											attachments.id,
+											objectRows.map((row) => row.id),
+										),
+									)
+									.returning({ id: attachments.id });
+								return deleted.length;
+							});
+							removed += page;
+							if (page === 0) break;
+						}
+						return removed;
+					},
+				);
+				await runStep(
+					operation,
+					ctx.actorUserId,
+					leaseOwner,
+					"release_attachment_quota",
+					deletedByDomain,
+					async () => {
+						let deleted = 0;
+						for (let page = 0; page < 10_000; page += 1) {
+							const result = await drainAttachmentStorageCleanupOutbox({
+								deleteObjects: runtime.deleteObjects,
+								actorUserId: ctx.actorUserId,
+								quotaAdmission:
+									runtime.attachmentStorageQuotaAdmission ??
+									getDocsMintRuntimeOptions()?.attachmentStorageQuotaAdmission,
+								signal: ctx.signal,
+								pageSize: LIFECYCLE_CLEANUP_PAGE_SIZE,
+								maxPages: 1,
+							});
+							deleted += result.deleted;
+							if (result.claimed === 0) break;
+							if (result.failed > 0 || result.deferred > 0)
+								throw new LifecyclePendingAttachmentUploadsError();
+						}
+						const retained = await withCleanupActor((tx) =>
+							tx
+								.select({ id: attachmentStorageCleanupOutbox.id })
+								.from(attachmentStorageCleanupOutbox)
+								.where(
+									or(
+										eq(
+											attachmentStorageCleanupOutbox.actorUserId,
+											ctx.actorUserId,
+										),
+										eq(
+											attachmentStorageCleanupOutbox.ownerUserId,
+											ctx.actorUserId,
+										),
+										eq(
+											attachmentStorageCleanupOutbox.requestedByUserId,
+											ctx.actorUserId,
+										),
+									),
+								)
+								.limit(1),
+						);
+						if (retained.length > 0)
+							throw new LifecyclePendingAttachmentUploadsError();
+						return deleted;
+					},
+				);
+				await runStep(
+					operation,
+					ctx.actorUserId,
+					leaseOwner,
+					"remove_attachment_rows",
+					deletedByDomain,
+					() => Promise.resolve(0),
+				);
+				await runStep(
+					operation,
+					ctx.actorUserId,
+					leaseOwner,
+					"clear_redis_state",
+					deletedByDomain,
+					() => runtime.clearAccountRedisState(ctx.actorUserId, ctx.signal),
+				);
+				for (const step of orderedSteps)
+					if (step.purge)
+						await runStep(
+							operation,
+							ctx.actorUserId,
+							leaseOwner,
+							`host:${step.id}`,
+							deletedByDomain,
+							async () => (await step.purge?.(ctx))?.deletedCount ?? 0,
+						);
+				await runStep(
+					operation,
+					ctx.actorUserId,
+					leaseOwner,
+					"remove_subject_documents",
+					deletedByDomain,
 					() =>
-						runtime.deleteObjects(
-							objectRows.map((row) => row.storageKey),
-							ctx.signal,
-						),
+						dbDelete(async (tx) => {
+							const currentDocuments = await tx
+								.select({ id: documents.id })
+								.from(documents)
+								.where(
+									and(
+										eq(documents.ownerId, ctx.actorUserId),
+										isNull(documents.workspaceId),
+									),
+								);
+							const currentDocumentIds = currentDocuments.map(({ id }) => id);
+							await acquireDocumentPipelineLocks(tx, currentDocumentIds);
+							if (currentDocumentIds.length === 0) return [];
+							return tx
+								.delete(documents)
+								.where(
+									and(
+										eq(documents.ownerId, ctx.actorUserId),
+										isNull(documents.workspaceId),
+									),
+								)
+								.returning({ id: documents.id });
+						}),
 				);
+				// Parent metadata uses ON DELETE SET NULL. Delete documents first so
+				// those referential actions do not become fenced document UPDATEs.
 				await runStep(
 					operation,
 					ctx.actorUserId,
@@ -639,7 +973,12 @@ export function createPersistentLifecycleService(
 						dbDelete((tx) =>
 							tx
 								.delete(folders)
-								.where(eq(folders.ownerId, ctx.actorUserId))
+								.where(
+									and(
+										eq(folders.ownerId, ctx.actorUserId),
+										isNull(folders.workspaceId),
+									),
+								)
 								.returning({ id: folders.id }),
 						),
 				);
@@ -653,7 +992,12 @@ export function createPersistentLifecycleService(
 						dbDelete((tx) =>
 							tx
 								.delete(tags)
-								.where(eq(tags.ownerId, ctx.actorUserId))
+								.where(
+									and(
+										eq(tags.ownerId, ctx.actorUserId),
+										isNull(tags.workspaceId),
+									),
+								)
 								.returning({ id: tags.id }),
 						),
 				);
@@ -667,7 +1011,12 @@ export function createPersistentLifecycleService(
 						dbDelete((tx) =>
 							tx
 								.delete(categories)
-								.where(eq(categories.ownerId, ctx.actorUserId))
+								.where(
+									and(
+										eq(categories.ownerId, ctx.actorUserId),
+										isNull(categories.workspaceId),
+									),
+								)
 								.returning({ id: categories.id }),
 						),
 				);
@@ -675,55 +1024,16 @@ export function createPersistentLifecycleService(
 					operation,
 					ctx.actorUserId,
 					leaseOwner,
-					"remove_attachment_rows",
-					deletedByDomain,
-					() =>
-						objectRows.length
-							? dbDelete((tx) =>
-									tx
-										.delete(attachments)
-										.where(
-											inArray(
-												attachments.id,
-												objectRows.map((row) => row.id),
-											),
-										)
-										.returning({ id: attachments.id }),
-								)
-							: Promise.resolve(0),
-				);
-				await runStep(
-					operation,
-					ctx.actorUserId,
-					leaseOwner,
-					"remove_subject_documents",
+					"remove_api_keys",
 					deletedByDomain,
 					() =>
 						dbDelete((tx) =>
 							tx
-								.delete(documents)
-								.where(eq(documents.ownerId, ctx.actorUserId))
-								.returning({ id: documents.id }),
+								.delete(apiKeys)
+								.where(eq(apiKeys.ownerId, ctx.actorUserId))
+								.returning({ id: apiKeys.id }),
 						),
 				);
-				await runStep(
-					operation,
-					ctx.actorUserId,
-					leaseOwner,
-					"clear_redis_state",
-					deletedByDomain,
-					() => runtime.clearAccountRedisState(ctx.actorUserId, ctx.signal),
-				);
-				for (const step of orderedSteps)
-					if (step.purge)
-						await runStep(
-							operation,
-							ctx.actorUserId,
-							leaseOwner,
-							`host:${step.id}`,
-							deletedByDomain,
-							async () => (await step.purge?.(ctx))?.deletedCount ?? 0,
-						);
 				await runStep(
 					operation,
 					ctx.actorUserId,
@@ -752,6 +1062,20 @@ export function createPersistentLifecycleService(
 								.returning({ id: accounts.id }),
 						),
 				);
+				await runStep(
+					operation,
+					ctx.actorUserId,
+					leaseOwner,
+					"remove_subject_audit",
+					deletedByDomain,
+					() =>
+						dbDelete((tx) =>
+							tx
+								.delete(auditLog)
+								.where(eq(auditLog.actorId, ctx.actorUserId))
+								.returning({ id: auditLog.id }),
+						),
+				);
 				await runStep(
 					operation,
 					ctx.actorUserId,
@@ -780,7 +1104,7 @@ export function createPersistentLifecycleService(
 					"write_deletion_audit",
 					deletedByDomain,
 					async () => {
-						await withActor(ctx.actorUserId, (tx) =>
+						await withCleanupActor((tx) =>
 							tx.insert(auditLog).values({
 								actorId: ctx.actorUserId,
 								action: "account_data_purged",