diff --git a/blocksuite/integration-test/src/__tests__/main/editor-semantics.spec.ts b/blocksuite/integration-test/src/__tests__/main/editor-semantics.spec.ts index f4dee5a26..40dfc947b 100644 --- a/blocksuite/integration-test/src/__tests__/main/editor-semantics.spec.ts +++ b/blocksuite/integration-test/src/__tests__/main/editor-semantics.spec.ts @@ -1,5 +1,6 @@ import { LinkExtension } from '@blocksuite/affine-inline-link'; import { textKeymap } from '@blocksuite/affine-inline-preset'; +import type { AffineReference } from '@blocksuite/affine-inline-reference'; import type { ListBlockModel, ParagraphBlockModel, @@ -312,6 +313,16 @@ describe('hotkey/bracket/linked-page', () => { const richText = getRichTextByBlockId(paragraphId); expect(richText.querySelectorAll('affine-reference').length).toBe(2); expect(richText.inlineEditor.yTextString.length).toBe(2); + + collection.removeDoc(linkedDoc.id); + await wait(); + expect(collection.docs.has(linkedDoc.id)).toBe(false); + const danglingReferences = + richText.querySelectorAll('affine-reference'); + expect(danglingReferences.length).toBe(2); + expect([...danglingReferences].every(reference => !reference.refMeta)).toBe( + true + ); }); }); diff --git a/packages/backend/native/index.d.ts b/packages/backend/native/index.d.ts index 2a7c2207c..d7c169167 100644 --- a/packages/backend/native/index.d.ts +++ b/packages/backend/native/index.d.ts @@ -22,7 +22,6 @@ export declare class BackendRuntime { compactPendingDocUpdates(workspaceId: string, docId: string, batchLimit: number, historyMinIntervalMs: number, historyMaxAgeSeconds: number, owner: string, leaseTtlMs: number): Promise upsertDocSnapshot(workspaceId: string, docId: string, blob: Buffer, timestampMs: number, editorId?: string | undefined | null): Promise createDocHistory(input: RuntimeDocHistoryInput): Promise - deleteDocStorage(workspaceId: string, docId: string): Promise putRuntimeGateIfAbsent(key: string, ttlMs: number): Promise cleanupExpiredRuntimeGates(limit: number): Promise cleanupExpiredUserSessions(limit: number): Promise @@ -78,6 +77,9 @@ export declare class StorageRuntime { backfillMissingBlobMetadata(workspaceId: string | undefined | null, limit: number): Promise rebuildDocBlobRefs(workspaceId: string, docId: string): Promise rebuildWorkspaceDocBlobRefs(workspaceId: string, limit: number): Promise + reconcileWorkspaceDocuments(workspaceId: string): Promise + executeDocumentCleanupCandidates(workspaceId: string | undefined | null, gracePeriodDays: number, limit: number): Promise + ackDocumentCleanupEffect(workspaceId: string, docId: string, cleanupVersion: string, effect: string): Promise constructor() start(): Promise configure(configJson: string): void @@ -976,6 +978,37 @@ export interface RuntimeDocHistoryInput { historyMaxAgeMs: number } +export interface RuntimeDocumentCleanupAckResult { + completed: boolean +} + +export interface RuntimeDocumentCleanupEffect { + workspaceId: string + docId: string + cleanupVersion: string + commentObjectsDone: boolean + searchDone: boolean + copilotDone: boolean +} + +export interface RuntimeDocumentCleanupExecuteResult { + scannedCandidates: number + serializationRetries: number + executed: number + recovered: number + reset: number + failed: number + deletedRows: number + effects: Array +} + +export interface RuntimeDocumentCleanupReconcileResult { + scannedDocs: number + marked: number + reset: number + recovered: number +} + export interface RuntimeInviteAbuseActionRequired { action: string subjectKey: string diff --git a/packages/backend/native/src/runtime/backend_runtime/doc_storage.rs b/packages/backend/native/src/runtime/backend_runtime/doc_storage.rs index c8fc8ad75..896f3203b 100644 --- a/packages/backend/native/src/runtime/backend_runtime/doc_storage.rs +++ b/packages/backend/native/src/runtime/backend_runtime/doc_storage.rs @@ -126,37 +126,4 @@ impl BackendRuntime { Ok(true) } - - #[napi] - pub async fn delete_doc_storage(&self, workspace_id: String, doc_id: String) -> napi::Result<()> { - let pool = self.pool().await?; - let mut tx = pool - .begin() - .await - .map_err(|err| RuntimeError::database("DocStorage delete begin transaction failed", err))?; - - sqlx::query("DELETE FROM snapshots WHERE workspace_id = $1 AND guid = $2") - .bind(&workspace_id) - .bind(&doc_id) - .execute(&mut *tx) - .await - .map_err(|err| RuntimeError::database("DocStorage delete snapshot failed", err))?; - sqlx::query("DELETE FROM updates WHERE workspace_id = $1 AND guid = $2") - .bind(&workspace_id) - .bind(&doc_id) - .execute(&mut *tx) - .await - .map_err(|err| RuntimeError::database("DocStorage delete updates failed", err))?; - sqlx::query("DELETE FROM snapshot_histories WHERE workspace_id = $1 AND guid = $2") - .bind(&workspace_id) - .bind(&doc_id) - .execute(&mut *tx) - .await - .map_err(|err| RuntimeError::database("DocStorage delete histories failed", err))?; - - tx.commit() - .await - .map_err(|err| RuntimeError::database("DocStorage delete commit failed", err))?; - Ok(()) - } } diff --git a/packages/backend/native/src/runtime/backend_runtime/tests.rs b/packages/backend/native/src/runtime/backend_runtime/tests.rs index 4a06d1873..885c0f14f 100644 --- a/packages/backend/native/src/runtime/backend_runtime/tests.rs +++ b/packages/backend/native/src/runtime/backend_runtime/tests.rs @@ -18,8 +18,9 @@ fn migrations_include_runtime_tables_without_worker_heartbeats() { assert!(RUNTIME_MIGRATIONS.contains("runtime_states")); assert!(RUNTIME_MIGRATIONS.contains("runtime_gates")); assert!(RUNTIME_MIGRATIONS.contains("runtime_leases")); - assert!(RUNTIME_MIGRATIONS.contains("blob_reconciliation_runs")); - assert!(RUNTIME_MIGRATIONS.contains("blob_reconciliation_checkpoints")); + assert!(RUNTIME_MIGRATIONS.contains("storage_reconciliation_runs")); + assert!(RUNTIME_MIGRATIONS.contains("storage_reconciliation_checkpoints")); + assert!(RUNTIME_MIGRATIONS.contains("document_cleanup_candidates")); assert!(RUNTIME_MIGRATIONS.contains("doc_blob_refs")); assert!(RUNTIME_MIGRATIONS.contains("blob_cleanup_candidates")); assert!(!RUNTIME_MIGRATIONS.contains("runtime_worker_heartbeats")); @@ -136,12 +137,10 @@ async fn insert_invite_quota_fixture( .bind(email) .execute(&pool) .await?; - sqlx::query( - "INSERT INTO workspaces (id, public, created_at) VALUES ($1, false, clock_timestamp() - interval '60 days')", - ) - .bind(&workspace_id) - .execute(&pool) - .await?; + sqlx::query("INSERT INTO workspaces (id, created_at) VALUES ($1, clock_timestamp() - interval '60 days')") + .bind(&workspace_id) + .execute(&pool) + .await?; sqlx::query( r#" INSERT INTO effective_workspace_quota_states ( diff --git a/packages/backend/native/src/runtime/error.rs b/packages/backend/native/src/runtime/error.rs index 55a339383..bc121a5b2 100644 --- a/packages/backend/native/src/runtime/error.rs +++ b/packages/backend/native/src/runtime/error.rs @@ -97,6 +97,16 @@ impl RuntimeError { _ => false, } } + + pub(crate) fn is_serialization_failure(&self) -> bool { + matches!( + self, + Self::Database { + source: sqlx::Error::Database(source), + .. + } if source.code().as_deref() == Some("40001") + ) + } } pub(crate) fn to_napi_error(error: RuntimeError) -> Error { diff --git a/packages/backend/native/src/runtime/migrations.rs b/packages/backend/native/src/runtime/migrations.rs index b49d687df..0a0ed453a 100644 --- a/packages/backend/native/src/runtime/migrations.rs +++ b/packages/backend/native/src/runtime/migrations.rs @@ -5,16 +5,10 @@ use super::{RuntimeError, RuntimeResult}; pub(crate) const RUNTIME_MIGRATIONS: &str = include_str!("sql/runtime_migrations.sql"); pub(crate) async fn migrate_runtime_tables(pool: &PgPool) -> RuntimeResult<()> { - for statement in RUNTIME_MIGRATIONS - .split(';') - .map(str::trim) - .filter(|statement| !statement.is_empty()) - { - sqlx::query(statement) - .execute(pool) - .await - .map_err(|err| RuntimeError::database("Runtime migration failed", err))?; - } + sqlx::raw_sql(RUNTIME_MIGRATIONS) + .execute(pool) + .await + .map_err(|err| RuntimeError::database("Runtime migration failed", err))?; Ok(()) } diff --git a/packages/backend/native/src/runtime/sql/runtime_migrations.sql b/packages/backend/native/src/runtime/sql/runtime_migrations.sql index 98b523db3..4ebf4fc7f 100644 --- a/packages/backend/native/src/runtime/sql/runtime_migrations.sql +++ b/packages/backend/native/src/runtime/sql/runtime_migrations.sql @@ -39,7 +39,7 @@ CREATE TABLE IF NOT EXISTS runtime_leases ( CREATE INDEX IF NOT EXISTS runtime_leases_expires_at_idx ON runtime_leases (expires_at); -CREATE TABLE IF NOT EXISTS blob_reconciliation_runs ( +CREATE TABLE IF NOT EXISTS storage_reconciliation_runs ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), kind TEXT NOT NULL, mode TEXT NOT NULL, @@ -54,10 +54,10 @@ CREATE TABLE IF NOT EXISTS blob_reconciliation_runs ( metadata JSONB NOT NULL DEFAULT '{}' ); -CREATE INDEX IF NOT EXISTS blob_reconciliation_runs_workspace_idx - ON blob_reconciliation_runs (workspace_id, started_at DESC); +CREATE INDEX IF NOT EXISTS storage_reconciliation_runs_workspace_idx + ON storage_reconciliation_runs (workspace_id, started_at DESC); -CREATE TABLE IF NOT EXISTS blob_reconciliation_checkpoints ( +CREATE TABLE IF NOT EXISTS storage_reconciliation_checkpoints ( kind TEXT NOT NULL, scope TEXT NOT NULL, status TEXT NOT NULL, @@ -70,8 +70,28 @@ CREATE TABLE IF NOT EXISTS blob_reconciliation_checkpoints ( PRIMARY KEY (kind, scope) ); -CREATE INDEX IF NOT EXISTS blob_reconciliation_checkpoints_status_idx - ON blob_reconciliation_checkpoints (kind, status, updated_at DESC); +CREATE INDEX IF NOT EXISTS storage_reconciliation_checkpoints_status_idx + ON storage_reconciliation_checkpoints (kind, status, updated_at DESC); + +CREATE TABLE IF NOT EXISTS document_cleanup_candidates ( + workspace_id TEXT NOT NULL, + doc_id TEXT NOT NULL, + status TEXT NOT NULL CHECK (status IN ('marked', 'effects_pending', 'failed')), + missing_since TIMESTAMPTZ(3) NOT NULL, + last_observed_missing_at TIMESTAMPTZ(3) NOT NULL, + last_doc_activity_at TIMESTAMPTZ(3), + cleanup_payload JSONB NOT NULL DEFAULT '{}', + attempt_count INTEGER NOT NULL DEFAULT 0 CHECK (attempt_count >= 0), + error TEXT, + updated_at TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (workspace_id, doc_id) +); + +CREATE INDEX IF NOT EXISTS document_cleanup_candidates_status_missing_idx + ON document_cleanup_candidates (status, missing_since); + +CREATE INDEX IF NOT EXISTS document_cleanup_candidates_workspace_status_idx + ON document_cleanup_candidates (workspace_id, status, updated_at DESC); CREATE TABLE IF NOT EXISTS doc_blob_refs ( workspace_id TEXT NOT NULL, diff --git a/packages/backend/native/src/runtime/storage_runtime/blob_cleanup.rs b/packages/backend/native/src/runtime/storage_runtime/blob_cleanup.rs index 808d0f364..5cf622b37 100644 --- a/packages/backend/native/src/runtime/storage_runtime/blob_cleanup.rs +++ b/packages/backend/native/src/runtime/storage_runtime/blob_cleanup.rs @@ -35,7 +35,7 @@ fn push_workspace_once(workspace_ids: &mut Vec, workspace_id: &str) { async fn checkpoint_completed(pool: &PgPool, kind: &str, scope: &str) -> RuntimeResult { sqlx::query_scalar::<_, bool>( - "SELECT EXISTS(SELECT 1 FROM blob_reconciliation_checkpoints WHERE kind = $1 AND scope = $2 AND status = \ + "SELECT EXISTS(SELECT 1 FROM storage_reconciliation_checkpoints WHERE kind = $1 AND scope = $2 AND status = \ 'completed')", ) .bind(kind) @@ -46,7 +46,43 @@ async fn checkpoint_completed(pool: &PgPool, kind: &str, scope: &str) -> Runtime } async fn projection_is_stale(pool: &PgPool, workspace_id: &str) -> RuntimeResult { - let checkpoint_fresh = checkpoint_completed(pool, "doc_blob_refs", workspace_id).await?; + let checkpoint_completed_at = sqlx::query_scalar::<_, Option>>( + r#" + SELECT MIN(completed_at) + FROM storage_reconciliation_checkpoints + WHERE scope = $1 + AND kind IN ('document_cleanup', 'doc_blob_refs') + AND status = 'completed' + HAVING COUNT(*) = 2 + "#, + ) + .bind(workspace_id) + .fetch_optional(pool) + .await + .map_err(|err| RuntimeError::database("Blob cleanup retention checkpoint load failed", err))? + .flatten(); + let Some(checkpoint_completed_at) = checkpoint_completed_at else { + return Ok(true); + }; + let activity_after_checkpoint = sqlx::query_scalar::<_, bool>( + r#" + SELECT EXISTS( + SELECT 1 FROM snapshots + WHERE workspace_id = $1 AND updated_at > $2 + UNION ALL + SELECT 1 FROM updates + WHERE workspace_id = $1 AND created_at > $2 + UNION ALL + SELECT 1 FROM snapshot_histories + WHERE workspace_id = $1 AND timestamp > $2 + ) + "#, + ) + .bind(workspace_id) + .bind(checkpoint_completed_at) + .fetch_one(pool) + .await + .map_err(|err| RuntimeError::database("Blob cleanup retention activity check failed", err))?; let has_stale_rows = sqlx::query_scalar::<_, bool>( "SELECT EXISTS(SELECT 1 FROM doc_blob_refs WHERE workspace_id = $1 AND status <> 'fresh')", ) @@ -54,7 +90,7 @@ async fn projection_is_stale(pool: &PgPool, workspace_id: &str) -> RuntimeResult .fetch_one(pool) .await .map_err(|err| RuntimeError::database("Blob cleanup projection freshness check failed", err))?; - Ok(!checkpoint_fresh || has_stale_rows) + Ok(activity_after_checkpoint || has_stale_rows) } async fn stale_projection_workspaces(pool: &PgPool, workspace_id: &str) -> RuntimeResult> { @@ -170,7 +206,7 @@ async fn load_completed_blobs( async fn load_plan_cursor(pool: &PgPool, workspace_id: &str) -> RuntimeResult> { let row = sqlx::query_as::<_, (String, serde_json::Value)>( - "SELECT status, cursor FROM blob_reconciliation_checkpoints WHERE kind = 'blob_cleanup_plan' AND scope = $1", + "SELECT status, cursor FROM storage_reconciliation_checkpoints WHERE kind = 'blob_cleanup_plan' AND scope = $1", ) .bind(workspace_id) .fetch_optional(pool) @@ -199,13 +235,13 @@ async fn upsert_plan_checkpoint( let status = if completed { "completed" } else { "running" }; sqlx::query( r#" - INSERT INTO blob_reconciliation_checkpoints + INSERT INTO storage_reconciliation_checkpoints (kind, scope, status, cursor, last_key, completed_at) VALUES ('blob_cleanup_plan', $1, $2, $3, $4, CASE WHEN $5 THEN CURRENT_TIMESTAMP ELSE NULL END) ON CONFLICT (kind, scope) DO UPDATE SET status = EXCLUDED.status, cursor = EXCLUDED.cursor, - last_key = COALESCE(EXCLUDED.last_key, blob_reconciliation_checkpoints.last_key), + last_key = COALESCE(EXCLUDED.last_key, storage_reconciliation_checkpoints.last_key), completed_at = CASE WHEN $5 THEN CURRENT_TIMESTAMP ELSE NULL END, updated_at = CURRENT_TIMESTAMP "#, @@ -224,7 +260,7 @@ async fn upsert_plan_checkpoint( async fn create_run(pool: &PgPool, workspace_id: &str) -> RuntimeResult { sqlx::query_scalar::<_, String>( r#" - INSERT INTO blob_reconciliation_runs (kind, mode, status, workspace_id) + INSERT INTO storage_reconciliation_runs (kind, mode, status, workspace_id) VALUES ('blob_cleanup_plan', 'mark_only', 'running', $1) RETURNING id::text "#, @@ -252,7 +288,7 @@ async fn finish_run( .unwrap_or(0); sqlx::query( r#" - UPDATE blob_reconciliation_runs + UPDATE storage_reconciliation_runs SET status = 'finished', finished_at = CURRENT_TIMESTAMP, scanned = $2, @@ -318,7 +354,7 @@ async fn finish_execute_run( ) -> RuntimeResult<()> { sqlx::query( r#" - UPDATE blob_reconciliation_runs + UPDATE storage_reconciliation_runs SET status = 'finished', finished_at = CURRENT_TIMESTAMP, scanned = $2, diff --git a/packages/backend/native/src/runtime/storage_runtime/blob_reconciliation.rs b/packages/backend/native/src/runtime/storage_runtime/blob_reconciliation.rs index dcf38488f..8c7c360ae 100644 --- a/packages/backend/native/src/runtime/storage_runtime/blob_reconciliation.rs +++ b/packages/backend/native/src/runtime/storage_runtime/blob_reconciliation.rs @@ -85,7 +85,8 @@ impl BackfillCheckpoint { async fn load_checkpoint(pool: &PgPool, scope: &str) -> RuntimeResult> { sqlx::query_as::<_, BackfillCheckpoint>( - "SELECT last_key, cursor FROM blob_reconciliation_checkpoints WHERE kind = 'blob_metadata_backfill' AND scope = $1", + "SELECT last_key, cursor FROM storage_reconciliation_checkpoints WHERE kind = 'blob_metadata_backfill' AND scope \ + = $1", ) .bind(scope) .fetch_optional(pool) @@ -103,13 +104,13 @@ async fn upsert_checkpoint( let status = if completed { "completed" } else { "running" }; sqlx::query( r#" - INSERT INTO blob_reconciliation_checkpoints + INSERT INTO storage_reconciliation_checkpoints (kind, scope, status, cursor, last_key, completed_at, metadata) VALUES ('blob_metadata_backfill', $1, $2, $3, $4, CASE WHEN $5 THEN CURRENT_TIMESTAMP ELSE NULL END, $6) ON CONFLICT (kind, scope) DO UPDATE SET status = EXCLUDED.status, cursor = EXCLUDED.cursor, - last_key = COALESCE(EXCLUDED.last_key, blob_reconciliation_checkpoints.last_key), + last_key = COALESCE(EXCLUDED.last_key, storage_reconciliation_checkpoints.last_key), completed_at = CASE WHEN $5 THEN CURRENT_TIMESTAMP ELSE NULL END, updated_at = CURRENT_TIMESTAMP, metadata = EXCLUDED.metadata @@ -247,7 +248,7 @@ impl StorageRuntime { sqlx::query( r#" - INSERT INTO blob_reconciliation_runs + INSERT INTO storage_reconciliation_runs (kind, mode, status, workspace_id, finished_at, scanned, changed, failed, metadata) VALUES ('blob_metadata_backfill', 'execute', 'finished', $1, CURRENT_TIMESTAMP, $2, $3, $4, $5) "#, diff --git a/packages/backend/native/src/runtime/storage_runtime/doc_blob_refs.rs b/packages/backend/native/src/runtime/storage_runtime/doc_blob_refs.rs index 1997a7d21..80fa15826 100644 --- a/packages/backend/native/src/runtime/storage_runtime/doc_blob_refs.rs +++ b/packages/backend/native/src/runtime/storage_runtime/doc_blob_refs.rs @@ -1,111 +1,32 @@ use affine_doc_loader as doc_loader; -use chrono::{DateTime, Utc}; -use sqlx::{FromRow, PgPool}; -use y_octo::Doc; +use sqlx::PgPool; -use super::{RuntimeDocBlobRefsResult, RuntimeError, RuntimeResult, StorageRuntime, napi_error}; +use super::{ + CurrentDoc, RuntimeDocBlobRefsResult, RuntimeError, RuntimeResult, StorageRuntime, load_current_doc, + load_workspace_live_doc_ids, napi_error, +}; const PARSER_VERSION: i32 = 1; -#[derive(FromRow)] -struct SnapshotRow { - workspace_id: String, - doc_id: String, - blob: Vec, - updated_at: DateTime, -} - -#[derive(FromRow)] -struct UpdateRow { - blob: Vec, - created_at: DateTime, -} - struct ExtractedRef { blob_key: String, block_id: String, flavour: String, } -async fn load_snapshot(pool: &PgPool, workspace_id: &str, doc_id: &str) -> RuntimeResult> { - sqlx::query_as::<_, SnapshotRow>( - r#" - SELECT workspace_id, guid AS doc_id, blob, updated_at - FROM snapshots - WHERE workspace_id = $1 AND guid = $2 - "#, +async fn load_workspace_doc_ids(pool: &PgPool, workspace_id: &str) -> RuntimeResult> { + let mut ids = load_workspace_live_doc_ids(pool, workspace_id).await?; + let retained = sqlx::query_scalar::<_, String>( + "SELECT doc_id FROM document_cleanup_candidates WHERE workspace_id = $1 AND status IN ('marked', 'failed') ORDER \ + BY doc_id", ) .bind(workspace_id) - .bind(doc_id) - .fetch_optional(pool) - .await - .map_err(|err| RuntimeError::database("Doc blob refs load snapshot failed", err)) -} - -async fn load_updates(pool: &PgPool, workspace_id: &str, doc_id: &str) -> RuntimeResult> { - sqlx::query_as::<_, UpdateRow>( - r#" - SELECT blob, created_at - FROM updates - WHERE workspace_id = $1 AND guid = $2 - ORDER BY created_at ASC - "#, - ) - .bind(workspace_id) - .bind(doc_id) .fetch_all(pool) .await - .map_err(|err| RuntimeError::database("Doc blob refs load updates failed", err)) -} - -fn apply_doc_updates(updates: impl IntoIterator>) -> RuntimeResult> { - let mut doc = Doc::default(); - for update in updates { - doc - .apply_update_from_binary_v1(&update) - .map_err(|err| RuntimeError::invalid_state(format!("Doc blob refs merge failed: {err}")))?; - } - doc - .encode_update_v1() - .map_err(|err| RuntimeError::invalid_state(format!("Doc blob refs encode failed: {err}"))) -} - -async fn load_current_doc(pool: &PgPool, workspace_id: &str, doc_id: &str) -> RuntimeResult> { - let snapshot = load_snapshot(pool, workspace_id, doc_id).await?; - let updates = load_updates(pool, workspace_id, doc_id).await?; - if snapshot.is_none() && updates.is_empty() { - return Ok(None); - } - - let mut merge_inputs = Vec::with_capacity(updates.len() + usize::from(snapshot.is_some())); - let mut updated_at = snapshot - .as_ref() - .map(|snapshot| snapshot.updated_at) - .unwrap_or_else(Utc::now); - if let Some(snapshot) = snapshot { - merge_inputs.push(snapshot.blob); - } - for update in updates { - updated_at = update.created_at; - merge_inputs.push(update.blob); - } - - Ok(Some(SnapshotRow { - workspace_id: workspace_id.to_string(), - doc_id: doc_id.to_string(), - blob: apply_doc_updates(merge_inputs)?, - updated_at, - })) -} - -async fn load_workspace_doc_ids(pool: &PgPool, workspace_id: &str) -> RuntimeResult> { - let Some(root) = load_current_doc(pool, workspace_id, workspace_id).await? else { - return Ok(Vec::new()); - }; - let ids = doc_loader::get_doc_ids_from_binary(root.blob, true) - .map_err(|err| RuntimeError::invalid_state(format!("Doc blob refs root doc parse failed: {err}")))?; - let mut ids = ids; + .map_err(|err| RuntimeError::database("Doc blob refs candidate load failed", err))?; + ids.extend(retained); ids.sort(); + ids.dedup(); Ok(ids) } @@ -124,7 +45,7 @@ async fn upsert_projection_checkpoint( }; sqlx::query( r#" - INSERT INTO blob_reconciliation_checkpoints + INSERT INTO storage_reconciliation_checkpoints (kind, scope, status, cursor, completed_at, metadata) VALUES ('doc_blob_refs', $1, $2, $3, CASE WHEN $4 THEN CURRENT_TIMESTAMP ELSE NULL END, $5) ON CONFLICT (kind, scope) DO UPDATE @@ -151,7 +72,7 @@ async fn upsert_projection_checkpoint( async fn upsert_projection_failure_checkpoint(pool: &PgPool, workspace_id: &str, error: &str) -> RuntimeResult<()> { sqlx::query( r#" - INSERT INTO blob_reconciliation_checkpoints + INSERT INTO storage_reconciliation_checkpoints (kind, scope, status, cursor, completed_at, metadata) VALUES ('doc_blob_refs', $1, 'failed', '{}', NULL, $2) ON CONFLICT (kind, scope) DO UPDATE @@ -174,14 +95,17 @@ async fn upsert_projection_failure_checkpoint(pool: &PgPool, workspace_id: &str, } async fn load_projection_cursor(pool: &PgPool, workspace_id: &str) -> RuntimeResult> { - let cursor = sqlx::query_scalar::<_, serde_json::Value>( - "SELECT cursor FROM blob_reconciliation_checkpoints WHERE kind = 'doc_blob_refs' AND scope = $1", + let checkpoint = sqlx::query_as::<_, (String, serde_json::Value)>( + "SELECT status, cursor FROM storage_reconciliation_checkpoints WHERE kind = 'doc_blob_refs' AND scope = $1", ) .bind(workspace_id) .fetch_optional(pool) .await .map_err(|err| RuntimeError::database("Doc blob refs checkpoint load failed", err))?; - Ok(cursor.and_then(|cursor| { + Ok(checkpoint.and_then(|(status, cursor)| { + if status != "running" { + return None; + } cursor .get("lastDocId") .and_then(|value| value.as_str()) @@ -205,7 +129,7 @@ async fn purge_removed_doc_refs(pool: &PgPool, workspace_id: &str, current_doc_i Ok(result.rows_affected() as i64) } -fn extract_refs(snapshot: &SnapshotRow) -> RuntimeResult> { +fn extract_refs(snapshot: &CurrentDoc) -> RuntimeResult> { let parsed = doc_loader::parse_doc_from_binary(snapshot.blob.clone(), snapshot.doc_id.clone()) .map_err(|err| RuntimeError::invalid_state(format!("Doc blob refs parse failed: {err}")))?; let mut refs = Vec::new(); @@ -226,6 +150,9 @@ fn extract_refs(snapshot: &SnapshotRow) -> RuntimeResult> { #[cfg(test)] mod tests { + use chrono::Utc; + use y_octo::Doc; + use super::*; #[test] @@ -233,7 +160,7 @@ mod tests { let doc_id = "doc-blob-ref-test".to_string(); let blob = doc_loader::build_full_doc("Doc", "![Alt](blob://image-blob-key)", &doc_id).expect("doc fixture should build"); - let snapshot = SnapshotRow { + let snapshot = CurrentDoc { workspace_id: "workspace".to_string(), doc_id, blob, @@ -272,9 +199,21 @@ mod tests { let ids = doc_loader::get_doc_ids_from_binary(root, true).expect("root doc ids should parse"); assert_eq!(ids, vec!["active-doc", "trashed-doc"]); } + + #[test] + fn doc_blob_refs_rejects_corrupt_docs() { + let snapshot = CurrentDoc { + workspace_id: "workspace".to_string(), + doc_id: "corrupt".to_string(), + blob: vec![0xff], + updated_at: Utc::now(), + }; + + assert!(extract_refs(&snapshot).is_err()); + } } -async fn replace_doc_refs(pool: &PgPool, snapshot: &SnapshotRow, refs: Vec) -> RuntimeResult<(i64, i64)> { +async fn replace_doc_refs(pool: &PgPool, snapshot: &CurrentDoc, refs: Vec) -> RuntimeResult<(i64, i64)> { let mut tx = pool .begin() .await diff --git a/packages/backend/native/src/runtime/storage_runtime/document_cleanup.rs b/packages/backend/native/src/runtime/storage_runtime/document_cleanup.rs new file mode 100644 index 000000000..2988bde62 --- /dev/null +++ b/packages/backend/native/src/runtime/storage_runtime/document_cleanup.rs @@ -0,0 +1,1507 @@ +use std::collections::{HashMap, HashSet}; + +use chrono::{DateTime, Utc}; +use serde_json::{Value, json}; +use sqlx::{FromRow, PgPool, Postgres, Row, Transaction}; +use uuid::Uuid; + +use super::{ + CurrentDoc, CurrentDocUpdate, RuntimeDocumentCleanupAckResult, RuntimeDocumentCleanupEffect, + RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, RuntimeError, RuntimeResult, + StorageRuntime, load_workspace_live_doc_ids, merge_current_doc, napi_error, +}; + +#[derive(FromRow)] +struct StoredDocActivity { + doc_id: String, + last_activity_at: DateTime, +} + +#[derive(FromRow)] +struct Candidate { + workspace_id: String, + doc_id: String, + last_doc_activity_at: Option>, +} + +#[derive(FromRow)] +struct PendingEffect { + workspace_id: String, + doc_id: String, + cleanup_payload: Value, +} + +async fn load_stored_doc_activity(pool: &PgPool, workspace_id: &str) -> RuntimeResult> { + sqlx::query_as::<_, StoredDocActivity>( + r#" + SELECT doc_id, MAX(activity_at) AS last_activity_at + FROM ( + SELECT guid AS doc_id, updated_at AS activity_at FROM snapshots WHERE workspace_id = $1 + UNION ALL + SELECT guid, created_at FROM updates WHERE workspace_id = $1 + UNION ALL + SELECT guid, timestamp FROM snapshot_histories WHERE workspace_id = $1 + ) stored + WHERE doc_id <> $1 + GROUP BY doc_id + ORDER BY doc_id + "#, + ) + .bind(workspace_id) + .fetch_all(pool) + .await + .map_err(|err| RuntimeError::database("Document cleanup stored doc scan failed", err)) +} + +async fn record_reconcile_failure( + pool: &PgPool, + workspace_id: &str, + failure_kind: &str, + error: &str, +) -> RuntimeResult<()> { + let mut tx = pool + .begin() + .await + .map_err(|err| RuntimeError::database("Document cleanup failure transaction failed", err))?; + let root_failed = i32::from(failure_kind == "root"); + let doc_failed = i32::from(failure_kind == "doc"); + sqlx::query( + r#" + INSERT INTO storage_reconciliation_checkpoints (kind, scope, status, cursor, completed_at, metadata) + VALUES ('document_cleanup', $1, 'failed', '{}', NULL, $2) + ON CONFLICT (kind, scope) DO UPDATE + SET status = 'failed', cursor = '{}', completed_at = NULL, + updated_at = CURRENT_TIMESTAMP, metadata = EXCLUDED.metadata + "#, + ) + .bind(workspace_id) + .bind(json!({ + "checkpointCompleted": false, + "failureKind": failure_kind, + "rootFailed": root_failed, + "docFailed": doc_failed, + "error": error, + })) + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup failure checkpoint write failed", err))?; + sqlx::query( + r#" + INSERT INTO storage_reconciliation_runs + (kind, mode, status, workspace_id, finished_at, failed, metadata) + VALUES ('document_cleanup', 'mark_only', 'failed', $1, CURRENT_TIMESTAMP, 1, $2) + "#, + ) + .bind(workspace_id) + .bind(json!({ + "checkpointCompleted": false, + "failureKind": failure_kind, + "rootFailed": root_failed, + "docFailed": doc_failed, + "error": error, + })) + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup failure run write failed", err))?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup failure commit failed", err))?; + Ok(()) +} + +async fn reconcile_workspace( + runtime: &StorageRuntime, + workspace_id: &str, +) -> RuntimeResult { + let pool = runtime.pool().await?; + let live_ids = match load_workspace_live_doc_ids(&pool, workspace_id).await { + Ok(ids) => ids.into_iter().collect::>(), + Err(err) => { + record_reconcile_failure(&pool, workspace_id, "root", &err.to_string()).await?; + return Err(err); + } + }; + let stored = match load_stored_doc_activity(&pool, workspace_id).await { + Ok(stored) => stored, + Err(err) => { + record_reconcile_failure(&pool, workspace_id, "scan", &err.to_string()).await?; + return Err(err); + } + }; + + let mut tx = pool + .begin() + .await + .map_err(|err| RuntimeError::database("Document cleanup reconcile transaction failed", err))?; + let now = Utc::now(); + let mut result = RuntimeDocumentCleanupReconcileResult { + scanned_docs: stored.len() as i64, + marked: 0, + reset: 0, + recovered: 0, + }; + + for doc in &stored { + if live_ids.contains(&doc.doc_id) { + continue; + } + let existing = sqlx::query( + r#" + SELECT status, last_doc_activity_at + FROM document_cleanup_candidates + WHERE workspace_id = $1 AND doc_id = $2 + "#, + ) + .bind(workspace_id) + .bind(&doc.doc_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup candidate load failed", err))?; + sqlx::query( + r#" + INSERT INTO document_cleanup_candidates + (workspace_id, doc_id, status, missing_since, last_observed_missing_at, last_doc_activity_at) + VALUES ($1, $2, 'marked', $3, $3, $4) + ON CONFLICT (workspace_id, doc_id) DO UPDATE + SET status = CASE + WHEN document_cleanup_candidates.status = 'effects_pending' THEN 'effects_pending' + ELSE 'marked' + END, + missing_since = CASE + WHEN document_cleanup_candidates.status = 'effects_pending' THEN document_cleanup_candidates.missing_since + WHEN document_cleanup_candidates.last_doc_activity_at IS DISTINCT FROM EXCLUDED.last_doc_activity_at + THEN EXCLUDED.missing_since + ELSE document_cleanup_candidates.missing_since + END, + last_observed_missing_at = EXCLUDED.last_observed_missing_at, + last_doc_activity_at = EXCLUDED.last_doc_activity_at, + error = CASE WHEN document_cleanup_candidates.status = 'effects_pending' + THEN document_cleanup_candidates.error ELSE NULL END, + updated_at = CURRENT_TIMESTAMP + "#, + ) + .bind(workspace_id) + .bind(&doc.doc_id) + .bind(now) + .bind(doc.last_activity_at) + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup candidate upsert failed", err))?; + match existing { + None => result.marked += 1, + Some(row) + if row.get::("status") != "effects_pending" + && row.get::>, _>("last_doc_activity_at") != Some(doc.last_activity_at) => + { + result.reset += 1; + } + Some(_) => {} + } + } + + let live_ids = live_ids.into_iter().collect::>(); + result.recovered = sqlx::query( + r#" + DELETE FROM document_cleanup_candidates + WHERE workspace_id = $1 + AND status IN ('marked', 'failed') + AND doc_id = ANY($2) + "#, + ) + .bind(workspace_id) + .bind(&live_ids) + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup recovered candidate delete failed", err))? + .rows_affected() as i64; + + sqlx::query( + r#" + INSERT INTO storage_reconciliation_checkpoints (kind, scope, status, cursor, completed_at, metadata) + VALUES ('document_cleanup', $1, 'completed', '{}', CURRENT_TIMESTAMP, $2) + ON CONFLICT (kind, scope) DO UPDATE + SET status = 'completed', cursor = '{}', completed_at = CURRENT_TIMESTAMP, + updated_at = CURRENT_TIMESTAMP, metadata = EXCLUDED.metadata + "#, + ) + .bind(workspace_id) + .bind(json!({ + "scannedDocs": result.scanned_docs, + "marked": result.marked, + "reset": result.reset, + "recovered": result.recovered, + "rootFailed": 0, + "docFailed": 0, + "checkpointCompleted": true, + })) + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup checkpoint write failed", err))?; + sqlx::query( + r#" + INSERT INTO storage_reconciliation_runs + (kind, mode, status, workspace_id, finished_at, scanned, changed, metadata) + VALUES ('document_cleanup', 'mark_only', 'finished', $1, CURRENT_TIMESTAMP, $2, $3, $4) + "#, + ) + .bind(workspace_id) + .bind(result.scanned_docs as i32) + .bind((result.marked + result.reset + result.recovered) as i32) + .bind(json!({ + "reset": result.reset, + "recovered": result.recovered, + "rootFailed": 0, + "docFailed": 0, + "checkpointCompleted": true, + })) + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup run write failed", err))?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup reconcile commit failed", err))?; + + Ok(result) +} + +async fn load_current_doc_for_update( + tx: &mut Transaction<'_, Postgres>, + workspace_id: &str, + doc_id: &str, +) -> RuntimeResult> { + let snapshot = sqlx::query_as::<_, CurrentDoc>( + "SELECT workspace_id, guid AS doc_id, blob, updated_at FROM snapshots WHERE workspace_id = $1 AND guid = $2", + ) + .bind(workspace_id) + .bind(doc_id) + .fetch_optional(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup current snapshot load failed", err))?; + let updates = sqlx::query_as::<_, CurrentDocUpdate>( + "SELECT blob, created_at FROM updates WHERE workspace_id = $1 AND guid = $2 ORDER BY created_at ASC", + ) + .bind(workspace_id) + .bind(doc_id) + .fetch_all(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup current updates load failed", err))?; + merge_current_doc(workspace_id, doc_id, snapshot, updates) +} + +async fn current_activity( + tx: &mut Transaction<'_, Postgres>, + workspace_id: &str, + doc_id: &str, +) -> RuntimeResult>> { + sqlx::query_scalar::<_, Option>>( + r#" + SELECT MAX(activity_at) + FROM ( + SELECT updated_at AS activity_at FROM snapshots WHERE workspace_id = $1 AND guid = $2 + UNION ALL SELECT created_at FROM updates WHERE workspace_id = $1 AND guid = $2 + UNION ALL SELECT timestamp FROM snapshot_histories WHERE workspace_id = $1 AND guid = $2 + ) activity + "#, + ) + .bind(workspace_id) + .bind(doc_id) + .fetch_one(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup activity load failed", err)) +} + +fn root_contains(root: CurrentDoc, doc_id: &str) -> RuntimeResult { + let ids = affine_doc_loader::get_doc_ids_from_binary(root.blob, true) + .map_err(|err| RuntimeError::invalid_state(format!("Document cleanup root parse failed: {err}")))?; + Ok(ids.iter().any(|id| id == doc_id)) +} + +async fn delete_doc_rows(tx: &mut Transaction<'_, Postgres>, candidate: &Candidate) -> RuntimeResult { + let attachment_keys = sqlx::query_scalar::<_, String>( + "SELECT key FROM comment_attachments WHERE workspace_id = $1 AND doc_id = $2 ORDER BY key", + ) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .fetch_all(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup attachment key load failed", err))?; + let cleanup_version = Uuid::new_v4().to_string(); + let storage_bytes = sqlx::query( + r#" + SELECT + COALESCE((SELECT SUM(octet_length(blob)) FROM snapshots WHERE workspace_id = $1 AND guid = $2), 0)::bigint AS snapshot_bytes, + COALESCE((SELECT SUM(octet_length(blob)) FROM updates WHERE workspace_id = $1 AND guid = $2), 0)::bigint AS update_bytes, + COALESCE((SELECT SUM(octet_length(blob)) FROM snapshot_histories WHERE workspace_id = $1 AND guid = $2), 0)::bigint AS history_bytes + "#, + ) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .fetch_one(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup storage bytes load failed", err))?; + let mut row_counts = HashMap::::new(); + row_counts.insert( + "ai_workspace_embeddings".to_string(), + sqlx::query_scalar::<_, i64>( + "SELECT COUNT(*) FROM ai_workspace_embeddings WHERE workspace_id = $1 AND doc_id = $2", + ) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .fetch_one(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup embedding cascade count failed", err))?, + ); + row_counts.insert( + "replies".to_string(), + sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM replies WHERE workspace_id = $1 AND doc_id = $2") + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .fetch_one(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup reply cascade count failed", err))?, + ); + let mut deleted_rows = 0; + for (table, doc_column) in [ + ("workspace_pages", "page_id"), + ("doc_access_policies", "doc_id"), + ("doc_grants", "doc_id"), + ("doc_blob_refs", "doc_id"), + ("ai_workspace_ignored_docs", "doc_id"), + ("comments", "doc_id"), + ("comment_attachments", "doc_id"), + ("workspace_doc_view_daily", "doc_id"), + ] { + let query = format!("DELETE FROM {table} WHERE workspace_id = $1 AND {doc_column} = $2"); + let affected = sqlx::query(&query) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .execute(&mut **tx) + .await + .map_err(|err| RuntimeError::database(format!("Document cleanup {table} delete failed"), err))? + .rows_affected() as i64; + deleted_rows += affected; + row_counts.insert(table.to_string(), affected); + } + for (table, doc_column) in [ + ("workspace_member_last_access", "last_doc_id"), + ("ai_sessions_metadata", "doc_id"), + ("ai_action_runs", "doc_id"), + ] { + let query = format!("UPDATE {table} SET {doc_column} = NULL WHERE workspace_id = $1 AND {doc_column} = $2"); + let affected = sqlx::query(&query) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .execute(&mut **tx) + .await + .map_err(|err| RuntimeError::database(format!("Document cleanup {table} unlink failed"), err))? + .rows_affected() as i64; + deleted_rows += affected; + row_counts.insert(format!("{table}.set_null"), affected); + } + for table in ["updates", "snapshot_histories", "snapshots"] { + let affected = sqlx::query(&format!("DELETE FROM {table} WHERE workspace_id = $1 AND guid = $2")) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .execute(&mut **tx) + .await + .map_err(|err| RuntimeError::database(format!("Document cleanup {table} delete failed"), err))? + .rows_affected() as i64; + deleted_rows += affected; + row_counts.insert(table.to_string(), affected); + } + + sqlx::query( + r#" + UPDATE document_cleanup_candidates + SET status = 'effects_pending', cleanup_payload = $3, error = NULL, + attempt_count = 0, updated_at = CURRENT_TIMESTAMP + WHERE workspace_id = $1 AND doc_id = $2 + "#, + ) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .bind(json!({ + "cleanupVersion": cleanup_version, + "commentAttachmentKeys": attachment_keys, + "commentObjectsDone": false, + "searchDone": false, + "copilotDone": false, + })) + .execute(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup candidate effect transition failed", err))?; + sqlx::query( + r#" + INSERT INTO storage_reconciliation_runs + (kind, mode, status, workspace_id, finished_at, scanned, changed, metadata) + VALUES ('document_cleanup_execute', 'execute', 'finished', $1, CURRENT_TIMESTAMP, 1, $2, $3) + "#, + ) + .bind(&candidate.workspace_id) + .bind(deleted_rows as i32) + .bind(json!({ + "docId": candidate.doc_id, + "cleanupVersion": cleanup_version, + "rowCounts": row_counts, + "snapshotBytes": storage_bytes.get::("snapshot_bytes"), + "updateBytes": storage_bytes.get::("update_bytes"), + "historyBytes": storage_bytes.get::("history_bytes"), + })) + .execute(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup execute audit write failed", err))?; + Ok(deleted_rows) +} + +async fn mark_candidate_failed( + tx: &mut Transaction<'_, Postgres>, + candidate: &Candidate, + error: String, +) -> RuntimeResult<()> { + sqlx::query( + r#" + UPDATE document_cleanup_candidates + SET status = 'failed', attempt_count = attempt_count + 1, + error = $3, updated_at = CURRENT_TIMESTAMP + WHERE workspace_id = $1 AND doc_id = $2 + "#, + ) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .bind(&error) + .execute(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup candidate failure write failed", err))?; + sqlx::query( + r#" + INSERT INTO storage_reconciliation_runs + (kind, mode, status, workspace_id, finished_at, scanned, failed, metadata) + VALUES ('document_cleanup_execute', 'execute', 'failed', $1, CURRENT_TIMESTAMP, 1, 1, $2) + "#, + ) + .bind(&candidate.workspace_id) + .bind(json!({ + "docId": candidate.doc_id, + "error": error, + })) + .execute(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup execute failure audit write failed", err))?; + Ok(()) +} + +async fn execute_one( + pool: &PgPool, + workspace_id: Option<&str>, + grace_period_days: i64, +) -> RuntimeResult> { + let mut tx = pool + .begin() + .await + .map_err(|err| RuntimeError::database("Document cleanup execute transaction failed", err))?; + sqlx::query("SET TRANSACTION ISOLATION LEVEL SERIALIZABLE") + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup isolation setup failed", err))?; + let candidate = sqlx::query_as::<_, Candidate>( + r#" + SELECT workspace_id, doc_id, last_doc_activity_at + FROM document_cleanup_candidates + WHERE (status = 'marked' + OR (status = 'failed' AND updated_at <= CURRENT_TIMESTAMP - INTERVAL '5 minutes')) + AND ($1::text IS NULL OR workspace_id = $1) + AND missing_since <= CURRENT_TIMESTAMP - make_interval(days => $2::int) + ORDER BY missing_since, workspace_id, doc_id + FOR UPDATE SKIP LOCKED + LIMIT 1 + "#, + ) + .bind(workspace_id) + .bind(grace_period_days as i32) + .fetch_optional(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup candidate claim failed", err))?; + let Some(candidate) = candidate else { + tx.rollback() + .await + .map_err(|err| RuntimeError::database("Document cleanup empty claim rollback failed", err))?; + return Ok(None); + }; + + let root = match load_current_doc_for_update(&mut tx, &candidate.workspace_id, &candidate.workspace_id).await { + Ok(Some(root)) => root, + Ok(None) => { + mark_candidate_failed( + &mut tx, + &candidate, + "Workspace root doc is missing during document cleanup execute".to_string(), + ) + .await?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup failed candidate commit failed", err))?; + return Ok(Some((candidate, -3))); + } + Err(err) => { + mark_candidate_failed(&mut tx, &candidate, err.to_string()).await?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup failed candidate commit failed", err))?; + return Ok(Some((candidate, -3))); + } + }; + let contains = match root_contains(root, &candidate.doc_id) { + Ok(contains) => contains, + Err(err) => { + mark_candidate_failed(&mut tx, &candidate, err.to_string()).await?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup failed candidate commit failed", err))?; + return Ok(Some((candidate, -3))); + } + }; + if contains { + sqlx::query("DELETE FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2") + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup recovered candidate delete failed", err))?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup recovered candidate commit failed", err))?; + return Ok(Some((candidate, -1))); + } + let activity = current_activity(&mut tx, &candidate.workspace_id, &candidate.doc_id).await?; + if activity != candidate.last_doc_activity_at { + sqlx::query( + r#" + UPDATE document_cleanup_candidates + SET status = 'marked', missing_since = CURRENT_TIMESTAMP, + last_observed_missing_at = CURRENT_TIMESTAMP, last_doc_activity_at = $3, + error = NULL, updated_at = CURRENT_TIMESTAMP + WHERE workspace_id = $1 AND doc_id = $2 + "#, + ) + .bind(&candidate.workspace_id) + .bind(&candidate.doc_id) + .bind(activity) + .execute(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup activity reset failed", err))?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup activity reset commit failed", err))?; + return Ok(Some((candidate, -2))); + } + let deleted_rows = delete_doc_rows(&mut tx, &candidate).await?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup execute commit failed", err))?; + Ok(Some((candidate, deleted_rows))) +} + +fn payload_effect(effect: PendingEffect) -> RuntimeResult { + let cleanup_version = effect + .cleanup_payload + .get("cleanupVersion") + .and_then(Value::as_str) + .ok_or_else(|| RuntimeError::invalid_state("Document cleanup effect payload has no cleanupVersion"))?; + Ok(RuntimeDocumentCleanupEffect { + workspace_id: effect.workspace_id, + doc_id: effect.doc_id, + cleanup_version: cleanup_version.to_string(), + comment_objects_done: effect + .cleanup_payload + .get("commentObjectsDone") + .and_then(Value::as_bool) + .unwrap_or(false), + search_done: effect + .cleanup_payload + .get("searchDone") + .and_then(Value::as_bool) + .unwrap_or(false), + copilot_done: effect + .cleanup_payload + .get("copilotDone") + .and_then(Value::as_bool) + .unwrap_or(false), + }) +} + +async fn complete_effect( + tx: &mut Transaction<'_, Postgres>, + workspace_id: &str, + doc_id: &str, + cleanup_version: &str, + path: &str, +) -> RuntimeResult { + let completed = sqlx::query_scalar::<_, bool>( + r#" + UPDATE document_cleanup_candidates + SET cleanup_payload = jsonb_set(cleanup_payload, ARRAY[$4], 'true'), + error = NULL, updated_at = CURRENT_TIMESTAMP + WHERE workspace_id = $1 AND doc_id = $2 AND status = 'effects_pending' + AND cleanup_payload->>'cleanupVersion' = $3 + RETURNING COALESCE((cleanup_payload->>'commentObjectsDone')::boolean, false) + AND COALESCE((cleanup_payload->>'searchDone')::boolean, false) + AND COALESCE((cleanup_payload->>'copilotDone')::boolean, false) + "#, + ) + .bind(workspace_id) + .bind(doc_id) + .bind(cleanup_version) + .bind(path) + .fetch_optional(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup effect completion failed", err))?; + let Some(completed) = completed else { + return Ok(true); + }; + if completed { + sqlx::query( + "DELETE FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2 AND \ + cleanup_payload->>'cleanupVersion' = $3", + ) + .bind(workspace_id) + .bind(doc_id) + .bind(cleanup_version) + .execute(&mut **tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup completed candidate delete failed", err))?; + } + Ok(completed) +} + +async fn process_comment_objects(runtime: &StorageRuntime, effect: &PendingEffect) -> RuntimeResult<()> { + let cleanup_version = effect + .cleanup_payload + .get("cleanupVersion") + .and_then(Value::as_str) + .ok_or_else(|| RuntimeError::invalid_state("Document cleanup effect is missing cleanupVersion"))?; + let comment_objects_done = effect + .cleanup_payload + .get("commentObjectsDone") + .and_then(Value::as_bool) + .unwrap_or(false); + if !comment_objects_done { + let keys = effect + .cleanup_payload + .get("commentAttachmentKeys") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_str) + .map(|key| format!("comment-attachments/{}/{}/{key}", effect.workspace_id, effect.doc_id)) + .collect::>(); + if !keys.is_empty() { + let outcomes = runtime.object_storage_delete_many(keys).await?; + if let Some(failed) = outcomes.iter().find(|outcome| outcome.error.is_some()) { + return Err(RuntimeError::invalid_state(format!( + "Comment attachment object delete failed for {}: {}", + failed.key, + failed.error.as_deref().unwrap_or("unknown") + ))); + } + } + } + let pool = runtime.pool().await?; + let mut tx = pool + .begin() + .await + .map_err(|err| RuntimeError::database("Document cleanup object effect transaction failed", err))?; + complete_effect( + &mut tx, + &effect.workspace_id, + &effect.doc_id, + cleanup_version, + "commentObjectsDone", + ) + .await?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup object effect commit failed", err))?; + Ok(()) +} + +async fn load_pending_effects( + pool: &PgPool, + workspace_id: Option<&str>, + limit: i64, +) -> RuntimeResult> { + sqlx::query_as::<_, PendingEffect>( + r#" + SELECT workspace_id, doc_id, cleanup_payload + FROM document_cleanup_candidates + WHERE status = 'effects_pending' AND ($1::text IS NULL OR workspace_id = $1) + ORDER BY updated_at, workspace_id, doc_id + LIMIT $2 + "#, + ) + .bind(workspace_id) + .bind(limit) + .fetch_all(pool) + .await + .map_err(|err| RuntimeError::database("Document cleanup pending effects load failed", err)) +} + +#[napi_derive::napi] +impl StorageRuntime { + #[napi] + pub async fn reconcile_workspace_documents( + &self, + workspace_id: String, + ) -> napi::Result { + Ok(reconcile_workspace(self, &workspace_id).await?) + } + + #[napi] + pub async fn execute_document_cleanup_candidates( + &self, + workspace_id: Option, + grace_period_days: i64, + limit: i64, + ) -> napi::Result { + if grace_period_days < 30 { + return Err(napi_error("document cleanup grace period must be at least 30 days")); + } + if limit <= 0 { + return Err(napi_error("document cleanup execute limit must be positive")); + } + let pool = self.pool().await?; + let mut result = RuntimeDocumentCleanupExecuteResult { + scanned_candidates: 0, + serialization_retries: 0, + executed: 0, + recovered: 0, + reset: 0, + failed: 0, + deleted_rows: 0, + effects: Vec::new(), + }; + for _ in 0..limit { + let mut retries = 0; + let outcome = loop { + match execute_one(&pool, workspace_id.as_deref(), grace_period_days).await { + Err(err) if err.is_serialization_failure() && retries < 3 => { + retries += 1; + result.serialization_retries += 1; + } + result => break result, + } + }?; + let Some((_, outcome)) = outcome else { break }; + result.scanned_candidates += 1; + match outcome { + -1 => result.recovered += 1, + -2 => result.reset += 1, + -3 => result.failed += 1, + rows => { + result.executed += 1; + result.deleted_rows += rows; + } + } + } + let effects = load_pending_effects(&pool, workspace_id.as_deref(), limit).await?; + for effect in effects { + if let Err(err) = process_comment_objects(self, &effect).await { + result.failed += 1; + sqlx::query( + r#" + UPDATE document_cleanup_candidates + SET attempt_count = attempt_count + 1, error = $3, updated_at = CURRENT_TIMESTAMP + WHERE workspace_id = $1 AND doc_id = $2 + "#, + ) + .bind(&effect.workspace_id) + .bind(&effect.doc_id) + .bind(err.to_string()) + .execute(&pool) + .await + .map_err(|db_err| RuntimeError::database("Document cleanup effect failure write failed", db_err))?; + } + } + result.effects = load_pending_effects(&pool, workspace_id.as_deref(), limit) + .await? + .into_iter() + .map(payload_effect) + .collect::>()?; + Ok(result) + } + + #[napi] + pub async fn ack_document_cleanup_effect( + &self, + workspace_id: String, + doc_id: String, + cleanup_version: String, + effect: String, + ) -> napi::Result { + let path = match effect.as_str() { + "search" => "searchDone", + "copilot" => "copilotDone", + _ => return Err(napi_error("document cleanup effect must be search or copilot")), + }; + let pool = self.pool().await?; + let mut tx = pool + .begin() + .await + .map_err(|err| RuntimeError::database("Document cleanup ack transaction failed", err))?; + let current_version = sqlx::query_scalar::<_, String>( + r#" + SELECT cleanup_payload->>'cleanupVersion' + FROM document_cleanup_candidates + WHERE workspace_id = $1 AND doc_id = $2 AND status = 'effects_pending' + "#, + ) + .bind(&workspace_id) + .bind(&doc_id) + .fetch_optional(&mut *tx) + .await + .map_err(|err| RuntimeError::database("Document cleanup ack candidate load failed", err))?; + let Some(current_version) = current_version else { + tx.rollback() + .await + .map_err(|err| RuntimeError::database("Document cleanup duplicate ack rollback failed", err))?; + return Ok(RuntimeDocumentCleanupAckResult { completed: true }); + }; + if current_version != cleanup_version { + return Err(napi_error("document cleanup effect candidate version mismatch")); + } + let completed = complete_effect(&mut tx, &workspace_id, &doc_id, &cleanup_version, path).await?; + tx.commit() + .await + .map_err(|err| RuntimeError::database("Document cleanup ack commit failed", err))?; + Ok(RuntimeDocumentCleanupAckResult { completed }) + } +} + +#[cfg(test)] +mod tests { + use std::sync::RwLock; + + use anyhow::{Context, Result as AnyResult}; + use napi::bindgen_prelude::Buffer; + use sqlx::postgres::PgPoolOptions; + use tokio::sync::Mutex; + + use super::*; + use crate::runtime::{migrations::migrate_runtime_tables, storage_runtime::StorageRuntimeConfig}; + + async fn runtime_from_database_url() -> AnyResult> { + let Ok(database_url) = std::env::var("DATABASE_URL") else { + return Ok(None); + }; + let pool = PgPoolOptions::new() + .max_connections(5) + .connect(&database_url) + .await + .context("connect postgres for document cleanup tests")?; + migrate_runtime_tables(&pool) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + let runtime = StorageRuntime { + config: RwLock::new(StorageRuntimeConfig { + database_url, + backends: HashMap::new(), + }), + pool: Mutex::new(Some(pool.clone())), + }; + Ok(Some((runtime, pool))) + } + + async fn insert_user_workspace(pool: &PgPool, suffix: &str) -> AnyResult<(String, String)> { + let user_id = format!("rust-test:document-cleanup:user:{suffix}"); + let workspace_id = format!("rust-test:document-cleanup:workspace:{suffix}"); + sqlx::query("DELETE FROM workspaces WHERE id = $1") + .bind(&workspace_id) + .execute(pool) + .await?; + sqlx::query("DELETE FROM users WHERE id = $1") + .bind(&user_id) + .execute(pool) + .await?; + sqlx::query( + r#" + INSERT INTO users (id, name, email, registered, email_verified, disabled, created_at) + VALUES ($1, 'Rust Document Cleanup Actor', $2, true, CURRENT_TIMESTAMP, false, CURRENT_TIMESTAMP) + "#, + ) + .bind(&user_id) + .bind(format!("{user_id}@example.com")) + .execute(pool) + .await?; + sqlx::query("INSERT INTO workspaces (id, created_at) VALUES ($1, CURRENT_TIMESTAMP)") + .bind(&workspace_id) + .execute(pool) + .await?; + Ok((user_id, workspace_id)) + } + + async fn cleanup_workspace_fixture(pool: &PgPool, user_id: &str, workspace_id: &str) -> AnyResult<()> { + sqlx::query("DELETE FROM storage_reconciliation_runs WHERE workspace_id = $1") + .bind(workspace_id) + .execute(pool) + .await?; + sqlx::query("DELETE FROM storage_reconciliation_checkpoints WHERE scope = $1") + .bind(workspace_id) + .execute(pool) + .await?; + for table in [ + "blob_cleanup_candidates", + "document_cleanup_candidates", + "doc_blob_refs", + ] { + sqlx::query(&format!("DELETE FROM {table} WHERE workspace_id = $1")) + .bind(workspace_id) + .execute(pool) + .await?; + } + sqlx::query("DELETE FROM workspaces WHERE id = $1") + .bind(workspace_id) + .execute(pool) + .await?; + sqlx::query("DELETE FROM users WHERE id = $1") + .bind(user_id) + .execute(pool) + .await?; + Ok(()) + } + + #[tokio::test] + async fn doc_blob_refs_projection_semantics_and_document_cleanup_mark_only_postgres() -> AnyResult<()> { + let Some((runtime, pool)) = runtime_from_database_url().await? else { + eprintln!("skipping postgres integration test: DATABASE_URL is not set"); + return Ok(()); + }; + let workspace_id = format!("rust-test:document-cleanup:{}", Uuid::new_v4()); + let doc_id = "missing-doc"; + let root = affine_doc_loader::add_doc_to_root_doc(Vec::new(), "live-doc", None)?; + let live_doc = affine_doc_loader::build_full_doc("Live", "", "live-doc")?; + let doc = affine_doc_loader::build_full_doc("Missing", "![Asset](blob://candidate-blob)", doc_id)?; + sqlx::query( + r#" + INSERT INTO snapshots (workspace_id, guid, blob, updated_at) + VALUES ($1, $1, $2, CURRENT_TIMESTAMP), + ($1, 'live-doc', $3, CURRENT_TIMESTAMP), + ($1, $4, $5, CURRENT_TIMESTAMP - INTERVAL '90 days') + "#, + ) + .bind(&workspace_id) + .bind(root) + .bind(live_doc) + .bind(doc_id) + .bind(doc) + .execute(&pool) + .await?; + + let observed_after = Utc::now(); + let first = reconcile_workspace(&runtime, &workspace_id).await?; + assert_eq!( + (first.scanned_docs, first.marked, first.reset, first.recovered), + (2, 1, 0, 0) + ); + let first_missing_since = sqlx::query_scalar::<_, DateTime>( + "SELECT missing_since FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(doc_id) + .fetch_one(&pool) + .await?; + assert!(first_missing_since >= observed_after); + + let projection = runtime + .rebuild_workspace_doc_blob_refs(workspace_id.clone(), 100) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + assert_eq!(projection.failed_docs, 0); + assert_eq!( + sqlx::query_scalar::<_, i64>( + "SELECT COUNT(*) FROM doc_blob_refs WHERE workspace_id = $1 AND doc_id = $2 AND blob_key = 'candidate-blob'", + ) + .bind(&workspace_id) + .bind(doc_id) + .fetch_one(&pool) + .await?, + 1 + ); + + let second = reconcile_workspace(&runtime, &workspace_id).await?; + assert_eq!((second.marked, second.reset, second.recovered), (0, 0, 0)); + let unchanged_missing_since = sqlx::query_scalar::<_, DateTime>( + "SELECT missing_since FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(doc_id) + .fetch_one(&pool) + .await?; + assert_eq!(unchanged_missing_since, first_missing_since); + + sqlx::query("UPDATE snapshots SET updated_at = CURRENT_TIMESTAMP WHERE workspace_id = $1 AND guid = $2") + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + let reset = reconcile_workspace(&runtime, &workspace_id).await?; + assert_eq!(reset.reset, 1); + let reset_missing_since = sqlx::query_scalar::<_, DateTime>( + "SELECT missing_since FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(doc_id) + .fetch_one(&pool) + .await?; + assert!(reset_missing_since >= first_missing_since); + + sqlx::query("UPDATE snapshots SET blob = $2 WHERE workspace_id = $1 AND guid = $1") + .bind(&workspace_id) + .bind(affine_doc_loader::add_doc_to_root_doc(Vec::new(), doc_id, None)?) + .execute(&pool) + .await?; + let recovered = reconcile_workspace(&runtime, &workspace_id).await?; + assert_eq!(recovered.recovered, 1); + + sqlx::query("UPDATE snapshots SET blob = $2 WHERE workspace_id = $1 AND guid = $1") + .bind(&workspace_id) + .bind(vec![0xff_u8]) + .execute(&pool) + .await?; + assert!(reconcile_workspace(&runtime, &workspace_id).await.is_err()); + let failure = sqlx::query( + "SELECT status, metadata FROM storage_reconciliation_checkpoints WHERE kind = 'document_cleanup' AND scope = $1", + ) + .bind(&workspace_id) + .fetch_one(&pool) + .await?; + assert_eq!(failure.get::("status"), "failed"); + assert_eq!(failure.get::("metadata")["rootFailed"], 1); + + sqlx::query("DELETE FROM storage_reconciliation_runs WHERE workspace_id = $1") + .bind(&workspace_id) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM storage_reconciliation_checkpoints WHERE scope = $1") + .bind(&workspace_id) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM doc_blob_refs WHERE workspace_id = $1") + .bind(&workspace_id) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM document_cleanup_candidates WHERE workspace_id = $1") + .bind(&workspace_id) + .execute(&pool) + .await?; + sqlx::query("DELETE FROM snapshots WHERE workspace_id = $1") + .bind(&workspace_id) + .execute(&pool) + .await?; + Ok(()) + } + + #[tokio::test] + async fn document_cleanup_execute_postgres_semantics() -> AnyResult<()> { + let Some((runtime, pool)) = runtime_from_database_url().await? else { + eprintln!("skipping postgres integration test: DATABASE_URL is not set"); + return Ok(()); + }; + let suffix = Uuid::new_v4().to_string(); + let (user_id, workspace_id) = insert_user_workspace(&pool, &suffix).await?; + let object_root = tempfile::tempdir()?; + runtime.config.write().unwrap().backends.insert( + "blob".to_string(), + super::super::StorageBackendConfig::Fs(super::super::FsStorageConfig { + provider: "fs".to_string(), + root: object_root.path().to_string_lossy().to_string(), + bucket: "document-cleanup-test".to_string(), + }), + ); + let doc_id = "missing-doc"; + let live_doc_id = "live-doc"; + let root = affine_doc_loader::add_doc_to_root_doc(Vec::new(), live_doc_id, None)?; + let live_doc = affine_doc_loader::build_full_doc("Live", "", live_doc_id)?; + let missing_doc = affine_doc_loader::build_full_doc("Doc", "![Alt](blob://image-blob-key)", doc_id)?; + sqlx::query( + r#" + INSERT INTO snapshots (workspace_id, guid, blob, updated_at) + VALUES ($1, $1, $2, CURRENT_TIMESTAMP), + ($1, $3, $4, CURRENT_TIMESTAMP), + ($1, $5, $6, CURRENT_TIMESTAMP - INTERVAL '90 days') + "#, + ) + .bind(&workspace_id) + .bind(root) + .bind(live_doc_id) + .bind(live_doc) + .bind(doc_id) + .bind(missing_doc) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO updates (workspace_id, guid, blob, created_at) VALUES ($1, $2, $3, CURRENT_TIMESTAMP - INTERVAL \ + '89 days')", + ) + .bind(&workspace_id) + .bind(doc_id) + .bind(affine_doc_loader::add_doc_to_root_doc( + Vec::new(), + "update-block", + None, + )?) + .execute(&pool) + .await?; + sqlx::query( + r#" + INSERT INTO snapshot_histories (workspace_id, guid, timestamp, blob, expired_at) + VALUES ($1, $2, CURRENT_TIMESTAMP - INTERVAL '88 days', $3, CURRENT_TIMESTAMP + INTERVAL '1 day') + "#, + ) + .bind(&workspace_id) + .bind(doc_id) + .bind(affine_doc_loader::add_doc_to_root_doc( + Vec::new(), + "history-block", + None, + )?) + .execute(&pool) + .await?; + + let mark = reconcile_workspace(&runtime, &workspace_id).await?; + assert_eq!(mark.marked, 1); + sqlx::query( + "UPDATE document_cleanup_candidates SET missing_since = CURRENT_TIMESTAMP - INTERVAL '29 days' WHERE \ + workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + let refs = runtime + .rebuild_workspace_doc_blob_refs(workspace_id.clone(), 100) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + assert_eq!(refs.failed_docs, 0); + let ref_count = sqlx::query_scalar::<_, i64>( + "SELECT COUNT(*) FROM doc_blob_refs WHERE workspace_id = $1 AND doc_id = $2 AND blob_key = 'image-blob-key'", + ) + .bind(&workspace_id) + .bind(doc_id) + .fetch_one(&pool) + .await?; + assert_eq!(ref_count, 1); + let not_due = execute_one(&pool, Some(&workspace_id), 30).await?; + assert!(not_due.is_none()); + + let session_id = format!("session:{suffix}"); + let prompt_name = format!("p_{}", &suffix[..30]); + sqlx::query( + "INSERT INTO ai_prompts_metadata (name, model, created_at, updated_at) VALUES ($1, 'test-model', \ + CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) ON CONFLICT (name) DO NOTHING", + ) + .bind(&prompt_name) + .execute(&pool) + .await?; + sqlx::query( + r#" + INSERT INTO ai_sessions_metadata + (id, user_id, workspace_id, doc_id, prompt_name, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + "#, + ) + .bind(&session_id) + .bind(&user_id) + .bind(&workspace_id) + .bind(doc_id) + .bind(&prompt_name) + .execute(&pool) + .await?; + sqlx::query( + r#" + INSERT INTO ai_action_runs + (id, user_id, workspace_id, doc_id, session_id, action_id, action_version, status, updated_at) + VALUES ($1, $2, $3, $4, $5, 'action', '1', 'running', CURRENT_TIMESTAMP) + "#, + ) + .bind(format!("action:{suffix}")) + .bind(&user_id) + .bind(&workspace_id) + .bind(doc_id) + .bind(&session_id) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO workspace_member_last_access (workspace_id, user_id, last_accessed_at, last_doc_id) VALUES ($1, \ + $2, CURRENT_TIMESTAMP, $3)", + ) + .bind(&workspace_id) + .bind(&user_id) + .bind(doc_id) + .execute(&pool) + .await?; + sqlx::query("INSERT INTO workspace_pages (workspace_id, page_id) VALUES ($1, $2)") + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + sqlx::query("INSERT INTO doc_access_policies (workspace_id, doc_id) VALUES ($1, $2)") + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO doc_grants (workspace_id, doc_id, principal_type, principal_id, role) VALUES ($1, $2, 'user', $3, \ + 'owner')", + ) + .bind(&workspace_id) + .bind(doc_id) + .bind(&user_id) + .execute(&pool) + .await?; + sqlx::query("INSERT INTO ai_workspace_ignored_docs (workspace_id, doc_id) VALUES ($1, $2)") + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + sqlx::query("INSERT INTO workspace_doc_view_daily (workspace_id, doc_id, date) VALUES ($1, $2, CURRENT_DATE)") + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + let comment_id = format!("comment:{suffix}"); + sqlx::query("INSERT INTO comments (id, workspace_id, doc_id, user_id, content) VALUES ($1, $2, $3, $4, '{}')") + .bind(&comment_id) + .bind(&workspace_id) + .bind(doc_id) + .bind(&user_id) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO replies (id, user_id, comment_id, workspace_id, doc_id, content) VALUES ($1, $2, $3, $4, $5, '{}')", + ) + .bind(format!("reply:{suffix}")) + .bind(&user_id) + .bind(&comment_id) + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + let attachment_key = "attachment-key"; + sqlx::query( + "INSERT INTO comment_attachments (workspace_id, doc_id, key, size, mime, name, created_by) VALUES ($1, $2, $3, \ + 4, 'text/plain', 'attachment.txt', $4)", + ) + .bind(&workspace_id) + .bind(doc_id) + .bind(attachment_key) + .bind(&user_id) + .execute(&pool) + .await?; + let attachment_object_key = format!("comment-attachments/{workspace_id}/{doc_id}/{attachment_key}"); + runtime + .put_object( + "blob".to_string(), + attachment_object_key.clone(), + Buffer::from(b"test".to_vec()), + None, + ) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + sqlx::query( + r#" + INSERT INTO ai_workspace_embeddings + (workspace_id, doc_id, chunk, content, embedding, created_at, updated_at) + VALUES ($1, $2, 0, 'content', ('[' || rtrim(repeat('0,', 1024), ',') || ']')::vector, + CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + "#, + ) + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + + sqlx::query( + "UPDATE document_cleanup_candidates SET missing_since = CURRENT_TIMESTAMP - INTERVAL '31 days' WHERE \ + workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(doc_id) + .execute(&pool) + .await?; + let executed = runtime + .execute_document_cleanup_candidates(Some(workspace_id.clone()), 30, 10) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + assert_eq!(executed.executed, 1); + assert_eq!(executed.failed, 0); + assert_eq!(executed.effects.len(), 1); + assert!(executed.effects[0].comment_objects_done); + assert!(!executed.effects[0].search_done); + assert!(!executed.effects[0].copilot_done); + assert!( + runtime + .head_object("blob".to_string(), attachment_object_key) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))? + .is_none() + ); + + for (table, column) in [ + ("snapshots", "guid"), + ("updates", "guid"), + ("snapshot_histories", "guid"), + ("workspace_pages", "page_id"), + ("doc_access_policies", "doc_id"), + ("doc_grants", "doc_id"), + ("doc_blob_refs", "doc_id"), + ("ai_workspace_ignored_docs", "doc_id"), + ("comments", "doc_id"), + ("comment_attachments", "doc_id"), + ("replies", "doc_id"), + ("workspace_doc_view_daily", "doc_id"), + ("ai_workspace_embeddings", "doc_id"), + ] { + let count = sqlx::query_scalar::<_, i64>(&format!( + "SELECT COUNT(*) FROM {table} WHERE workspace_id = $1 AND {column} = $2" + )) + .bind(&workspace_id) + .bind(doc_id) + .fetch_one(&pool) + .await?; + assert_eq!(count, 0, "{table} should be cleaned"); + } + for (table, column) in [ + ("workspace_member_last_access", "last_doc_id"), + ("ai_sessions_metadata", "doc_id"), + ("ai_action_runs", "doc_id"), + ] { + let count = sqlx::query_scalar::<_, i64>(&format!( + "SELECT COUNT(*) FROM {table} WHERE workspace_id = $1 AND {column} = $2" + )) + .bind(&workspace_id) + .bind(doc_id) + .fetch_one(&pool) + .await?; + assert_eq!(count, 0, "{table}.{column} should be nulled"); + } + + let effect = &executed.effects[0]; + let search = runtime + .ack_document_cleanup_effect( + workspace_id.clone(), + doc_id.to_string(), + effect.cleanup_version.clone(), + "search".to_string(), + ) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + assert!(!search.completed); + let copilot = runtime + .ack_document_cleanup_effect( + workspace_id.clone(), + doc_id.to_string(), + effect.cleanup_version.clone(), + "copilot".to_string(), + ) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + assert!(copilot.completed); + let candidate_count = sqlx::query_scalar::<_, i64>( + "SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(doc_id) + .fetch_one(&pool) + .await?; + assert_eq!(candidate_count, 0); + + let retry_doc_id = "object-retry-doc"; + let retry_doc = affine_doc_loader::build_full_doc("Retry", "", retry_doc_id)?; + sqlx::query( + "INSERT INTO snapshots (workspace_id, guid, blob, updated_at) VALUES ($1, $2, $3, CURRENT_TIMESTAMP - INTERVAL \ + '90 days')", + ) + .bind(&workspace_id) + .bind(retry_doc_id) + .bind(retry_doc) + .execute(&pool) + .await?; + sqlx::query( + "INSERT INTO comment_attachments (workspace_id, doc_id, key, size, mime, name, created_by) VALUES ($1, $2, \ + '..', 4, 'text/plain', 'attachment.txt', $3)", + ) + .bind(&workspace_id) + .bind(retry_doc_id) + .bind(&user_id) + .execute(&pool) + .await?; + assert_eq!(reconcile_workspace(&runtime, &workspace_id).await?.marked, 1); + sqlx::query( + "UPDATE document_cleanup_candidates SET missing_since = CURRENT_TIMESTAMP - INTERVAL '31 days' WHERE \ + workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(retry_doc_id) + .execute(&pool) + .await?; + + let failed_object_delete = runtime + .execute_document_cleanup_candidates(Some(workspace_id.clone()), 30, 10) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + assert_eq!(failed_object_delete.executed, 1); + assert_eq!(failed_object_delete.failed, 1); + assert!(!failed_object_delete.effects[0].comment_objects_done); + let retained = sqlx::query( + "SELECT status, attempt_count, error, cleanup_payload->>'cleanupVersion' AS cleanup_version FROM \ + document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(retry_doc_id) + .fetch_one(&pool) + .await?; + assert_eq!(retained.get::("status"), "effects_pending"); + assert_eq!(retained.get::("attempt_count"), 1); + assert!(retained.get::, _>("error").is_some()); + let retry_cleanup_version = retained.get::("cleanup_version"); + + for effect in ["search", "copilot"] { + let ack = runtime + .ack_document_cleanup_effect( + workspace_id.clone(), + retry_doc_id.to_string(), + retry_cleanup_version.clone(), + effect.to_string(), + ) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + assert!(!ack.completed); + } + + sqlx::query( + "UPDATE document_cleanup_candidates SET cleanup_payload = jsonb_set(cleanup_payload, '{commentAttachmentKeys}', \ + '[\"already-missing\"]') WHERE workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(retry_doc_id) + .execute(&pool) + .await?; + let retried = runtime + .execute_document_cleanup_candidates(Some(workspace_id.clone()), 30, 10) + .await + .map_err(|err| anyhow::anyhow!(err.to_string()))?; + assert_eq!(retried.executed, 0); + assert_eq!(retried.failed, 0); + assert!(retried.effects.is_empty()); + let retry_candidate_count = sqlx::query_scalar::<_, i64>( + "SELECT COUNT(*) FROM document_cleanup_candidates WHERE workspace_id = $1 AND doc_id = $2", + ) + .bind(&workspace_id) + .bind(retry_doc_id) + .fetch_one(&pool) + .await?; + assert_eq!(retry_candidate_count, 0); + + cleanup_workspace_fixture(&pool, &user_id, &workspace_id).await?; + Ok(()) + } +} diff --git a/packages/backend/native/src/runtime/storage_runtime/mod.rs b/packages/backend/native/src/runtime/storage_runtime/mod.rs index 403231518..b1d2ea75d 100644 --- a/packages/backend/native/src/runtime/storage_runtime/mod.rs +++ b/packages/backend/native/src/runtime/storage_runtime/mod.rs @@ -7,18 +7,21 @@ use std::{ }; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; +use chrono::{DateTime, Utc}; use napi::bindgen_prelude::Buffer; use serde::Deserialize; use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; -use sqlx::{PgPool, Row, postgres::PgPoolOptions}; +use sqlx::{FromRow, PgPool, Row, postgres::PgPoolOptions}; use tokio::{sync::Mutex, task::JoinSet}; +use y_octo::Doc; mod assetpack; mod blob_cleanup; mod blob_reclaimer; mod blob_reconciliation; mod doc_blob_refs; +mod document_cleanup; pub(crate) mod object_storage; use self::object_storage::{ @@ -33,9 +36,10 @@ pub(super) use super::{ napi_error, to_napi_error, types::{ RuntimeBlobCleanupExecuteResult, RuntimeBlobCleanupPlanResult, RuntimeBlobCleanupResult, RuntimeBlobCompleteResult, - RuntimeBlobMetadataBackfillResult, RuntimeDocBlobRefsResult, RuntimeMultipartUploadInit, - RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry, RuntimeObjectMetadata, - RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest, + RuntimeBlobMetadataBackfillResult, RuntimeDocBlobRefsResult, RuntimeDocumentCleanupAckResult, + RuntimeDocumentCleanupEffect, RuntimeDocumentCleanupExecuteResult, RuntimeDocumentCleanupReconcileResult, + RuntimeMultipartUploadInit, RuntimeMultipartUploadPart, RuntimeObjectGetResult, RuntimeObjectListEntry, + RuntimeObjectMetadata, RuntimeObjectStoragePutOptions, RuntimePresignedObjectRequest, }, }; @@ -45,6 +49,100 @@ const OBJECT_DELETE_MANY_CONCURRENCY: usize = 3; type Result = RuntimeResult; +#[derive(FromRow)] +struct CurrentDoc { + workspace_id: String, + doc_id: String, + blob: Vec, + updated_at: DateTime, +} + +#[derive(FromRow)] +struct CurrentDocUpdate { + blob: Vec, + created_at: DateTime, +} + +async fn load_current_doc(pool: &PgPool, workspace_id: &str, doc_id: &str) -> RuntimeResult> { + let snapshot = sqlx::query_as::<_, CurrentDoc>( + r#" + SELECT workspace_id, guid AS doc_id, blob, updated_at + FROM snapshots + WHERE workspace_id = $1 AND guid = $2 + "#, + ) + .bind(workspace_id) + .bind(doc_id) + .fetch_optional(pool) + .await + .map_err(|err| RuntimeError::database("Current doc snapshot load failed", err))?; + let updates = sqlx::query_as::<_, CurrentDocUpdate>( + r#" + SELECT blob, created_at + FROM updates + WHERE workspace_id = $1 AND guid = $2 + ORDER BY created_at ASC + "#, + ) + .bind(workspace_id) + .bind(doc_id) + .fetch_all(pool) + .await + .map_err(|err| RuntimeError::database("Current doc updates load failed", err))?; + merge_current_doc(workspace_id, doc_id, snapshot, updates) +} + +fn merge_current_doc( + workspace_id: &str, + doc_id: &str, + snapshot: Option, + updates: Vec, +) -> RuntimeResult> { + if snapshot.is_none() && updates.is_empty() { + return Ok(None); + } + let mut doc = Doc::default(); + let mut updated_at = snapshot + .as_ref() + .map(|snapshot| snapshot.updated_at) + .or_else(|| updates.first().map(|update| update.created_at)) + .unwrap_or_else(Utc::now); + if let Some(snapshot) = &snapshot { + doc + .apply_update_from_binary_v1(&snapshot.blob) + .map_err(|err| RuntimeError::invalid_state(format!("Current doc snapshot merge failed: {err}")))?; + } + for update in updates { + updated_at = updated_at.max(update.created_at); + doc + .apply_update_from_binary_v1(&update.blob) + .map_err(|err| RuntimeError::invalid_state(format!("Current doc update merge failed: {err}")))?; + } + let blob = doc + .encode_update_v1() + .map_err(|err| RuntimeError::invalid_state(format!("Current doc encode failed: {err}")))?; + + Ok(Some(CurrentDoc { + workspace_id: workspace_id.to_string(), + doc_id: doc_id.to_string(), + blob, + updated_at, + })) +} + +async fn load_workspace_live_doc_ids(pool: &PgPool, workspace_id: &str) -> RuntimeResult> { + workspace_live_doc_ids(load_current_doc(pool, workspace_id, workspace_id).await?) +} + +fn workspace_live_doc_ids(root: Option) -> RuntimeResult> { + let root = root.ok_or_else(|| RuntimeError::invalid_state("Workspace root doc is missing"))?; + let mut ids = affine_doc_loader::get_doc_ids_from_binary(root.blob, true) + .map_err(|err| RuntimeError::invalid_state(format!("Workspace root doc parse failed: {err}")))?; + ids.sort(); + ids.dedup(); + Ok(ids) +} + #[napi_derive::napi(object)] pub struct StorageRuntimeHealth { pub started: bool, @@ -1385,6 +1483,81 @@ fn system_time_ms(time: SystemTime) -> Result { mod tests { use super::*; + #[test] + fn workspace_live_set_merges_pending_updates_and_includes_trash() { + use y_octo::{Any, Value}; + + let snapshot = affine_doc_loader::add_doc_to_root_doc(Vec::new(), "live", None).unwrap(); + let pending = affine_doc_loader::add_doc_to_root_doc(snapshot.clone(), "trash", None).unwrap(); + let merged = merge_current_doc( + "workspace", + "workspace", + Some(CurrentDoc { + workspace_id: "workspace".to_string(), + doc_id: "workspace".to_string(), + blob: snapshot, + updated_at: Utc::now(), + }), + vec![CurrentDocUpdate { + blob: pending, + created_at: Utc::now(), + }], + ) + .unwrap() + .unwrap(); + let mut root = Doc::default(); + root.apply_update_from_binary_v1(&merged.blob).unwrap(); + let meta = root.get_map("meta").unwrap(); + let mut pages = meta.get("pages").and_then(|value| value.to_array()).unwrap(); + let mut trash = pages + .iter() + .find_map(|value| { + let page = value.to_map()?; + (page.get("id")?.to_any()? == Any::String("trash".to_string())).then_some(page) + }) + .unwrap(); + trash.insert("trash".to_string(), Value::Any(Any::True)).unwrap(); + + let ids = workspace_live_doc_ids(Some(CurrentDoc { + workspace_id: "workspace".to_string(), + doc_id: "workspace".to_string(), + blob: root.encode_update_v1().unwrap(), + updated_at: Utc::now(), + })) + .unwrap(); + assert_eq!(ids, ["live", "trash"]); + + let trash_index = pages + .iter() + .position(|value| { + value.to_map().and_then(|page| page.get("id")) == Some(Value::Any(Any::String("trash".to_string()))) + }) + .unwrap(); + pages.remove(trash_index as u64, 1).unwrap(); + let ids = workspace_live_doc_ids(Some(CurrentDoc { + workspace_id: "workspace".to_string(), + doc_id: "workspace".to_string(), + blob: root.encode_update_v1().unwrap(), + updated_at: Utc::now(), + })) + .unwrap(); + assert_eq!(ids, ["live"]); + } + + #[test] + fn workspace_live_set_fails_closed_for_missing_or_corrupt_root() { + assert!(workspace_live_doc_ids(None).is_err()); + assert!( + workspace_live_doc_ids(Some(CurrentDoc { + workspace_id: "workspace".to_string(), + doc_id: "workspace".to_string(), + blob: vec![0xff], + updated_at: Utc::now(), + })) + .is_err() + ); + } + #[test] fn fs_key_normalization_rejects_traversal() { for (key, valid) in [ diff --git a/packages/backend/native/src/runtime/types.rs b/packages/backend/native/src/runtime/types.rs index 295d6fe1d..ae926bf50 100644 --- a/packages/backend/native/src/runtime/types.rs +++ b/packages/backend/native/src/runtime/types.rs @@ -243,6 +243,41 @@ pub struct RuntimeDocBlobRefsResult { pub next_cursor: Option, } +#[napi_derive::napi(object)] +pub struct RuntimeDocumentCleanupReconcileResult { + pub scanned_docs: i64, + pub marked: i64, + pub reset: i64, + pub recovered: i64, +} + +#[napi_derive::napi(object)] +pub struct RuntimeDocumentCleanupEffect { + pub workspace_id: String, + pub doc_id: String, + pub cleanup_version: String, + pub comment_objects_done: bool, + pub search_done: bool, + pub copilot_done: bool, +} + +#[napi_derive::napi(object)] +pub struct RuntimeDocumentCleanupExecuteResult { + pub scanned_candidates: i64, + pub serialization_retries: i64, + pub executed: i64, + pub recovered: i64, + pub reset: i64, + pub failed: i64, + pub deleted_rows: i64, + pub effects: Vec, +} + +#[napi_derive::napi(object)] +pub struct RuntimeDocumentCleanupAckResult { + pub completed: bool, +} + #[napi_derive::napi(object)] pub struct RuntimeBlobCleanupPlanResult { pub run_id: Option, diff --git a/packages/backend/server/src/__tests__/copilot/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot/copilot.spec.ts index 055e5ec3f..3d2af6816 100644 --- a/packages/backend/server/src/__tests__/copilot/copilot.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/copilot.spec.ts @@ -28,7 +28,7 @@ import { WorkspaceModel, WorkspaceRole, } from '../../models'; -import type { LlmToolCallbackRequest } from '../../native'; +import { addDocToRootDoc, type LlmToolCallbackRequest } from '../../native'; import { CopilotModule } from '../../plugins/copilot'; import { CopilotContextService } from '../../plugins/copilot/context'; import { CopilotContextResolver } from '../../plugins/copilot/context/resolver'; @@ -260,6 +260,80 @@ test.after.always(async t => { await t.context.module?.close(); }); +test('document cleanup reconciles missing and restored copilot state before ack', async t => { + const { db, jobs, models, module, workspace } = t.context; + const queue = module.get(JobQueue); + const deleteEmbedding = Sinon.spy( + models.copilotContext, + 'purgeWorkspaceEmbedding' + ); + const scheduleEmbedding = Sinon.stub( + jobs, + 'addDocEmbeddingQueueFromEvent' + ).resolves(); + + for (const [docId, restored, cleanupVersion] of [ + ['missing-doc', false, 'missing-version'], + ['restored-doc', true, 'restored-version'], + ] as const) { + const ws = await workspace.create(userId); + const root = addDocToRootDoc(Buffer.from([0, 0]), docId, docId); + const missingRoot = addDocToRootDoc( + Buffer.from([0, 0]), + 'live-doc', + 'Live' + ); + await db.snapshot.create({ + data: { + workspaceId: ws.id, + id: ws.id, + blob: restored ? root : missingRoot, + state: Buffer.from([0, 0]), + updatedAt: new Date(), + createdAt: new Date(), + }, + }); + if (restored) { + await db.snapshot.create({ + data: { + workspaceId: ws.id, + id: docId, + blob: addDocToRootDoc(Buffer.from([0, 0]), 'content', 'Content'), + state: Buffer.from([0, 0]), + updatedAt: new Date(), + createdAt: new Date(), + }, + }); + } + + await jobs.reconcileDocumentCleanup({ + workspaceId: ws.id, + docId, + cleanupVersion, + }); + + if (restored) { + t.true(scheduleEmbedding.calledOnceWith({ workspaceId: ws.id, docId })); + t.false(deleteEmbedding.called); + } else { + t.true(deleteEmbedding.calledOnceWith(ws.id, docId)); + t.false(scheduleEmbedding.called); + } + const { payload } = await module.queue.waitFor( + 'backendRuntime.ackDocumentCleanupEffect' + ); + t.deepEqual(payload, { + workspaceId: ws.id, + docId, + cleanupVersion, + effect: 'copilot', + }); + deleteEmbedding.resetHistory(); + scheduleEmbedding.resetHistory(); + (queue.add as Sinon.SinonStub).resetHistory(); + } +}); + test('MCP credentials stay bound to their endpoint, workspace, and profile', async t => { const { db, mcpCredentials, mcpProvider, models, workspace } = t.context; const ws = await workspace.create(userId); diff --git a/packages/backend/server/src/__tests__/sync/gateway.spec.ts b/packages/backend/server/src/__tests__/sync/gateway.spec.ts index 51def5203..1ad515937 100644 --- a/packages/backend/server/src/__tests__/sync/gateway.spec.ts +++ b/packages/backend/server/src/__tests__/sync/gateway.spec.ts @@ -671,7 +671,7 @@ test('active users metric should dedupe multiple sockets for one user', async t test('workspace sync delete-doc should enforce doc permissions', async t => { const db = app.get(PrismaClient); const models = app.get(Models); - const { user: owner } = await login(app); + const { user: owner, cookieHeader: ownerCookieHeader } = await login(app); const { user: collaborator, cookieHeader } = await login(app); const workspace = await models.workspace.create(owner.id); const docId = 'private-doc'; @@ -692,9 +692,10 @@ test('workspace sync delete-doc should enforce doc permissions', async t => { }); const socket = createClient(url, cookieHeader); + const ownerSocket = createClient(url, ownerCookieHeader); try { - await waitForConnect(socket); + await Promise.all([waitForConnect(socket), waitForConnect(ownerSocket)]); const join = unwrapResponse( t, @@ -719,8 +720,37 @@ test('workspace sync delete-doc should enforce doc permissions', async t => { }) ); t.true(error.message.includes('Doc.Delete')); + + const ownerJoin = unwrapResponse( + t, + await emitWithAck<{ clientId: string; success: boolean }>( + ownerSocket, + 'space:join', + { + spaceType: 'workspace', + spaceId: workspace.id, + clientVersion: '0.26.0', + } + ) + ); + t.true(ownerJoin.success); + unwrapResponse( + t, + await emitWithAck(ownerSocket, 'space:delete-doc', { + spaceType: 'workspace', + spaceId: workspace.id, + docId, + }) + ); + t.is( + await db.snapshot.count({ + where: { workspaceId: workspace.id, id: docId }, + }), + 1 + ); } finally { socket.disconnect(); + ownerSocket.disconnect(); } }); diff --git a/packages/backend/server/src/core/doc/adapters/workspace.ts b/packages/backend/server/src/core/doc/adapters/workspace.ts index c5c576b8a..9cf8ae085 100644 --- a/packages/backend/server/src/core/doc/adapters/workspace.ts +++ b/packages/backend/server/src/core/doc/adapters/workspace.ts @@ -137,8 +137,8 @@ export class PgWorkspaceDocStorageAdapter extends DocStorageAdapter { })); } - async deleteDoc(workspaceId: string, docId: string) { - await this.models.doc.delete(workspaceId, docId); + async deleteDoc(_workspaceId: string, _docId: string) { + return; } async deleteSpace(workspaceId: string) { diff --git a/packages/backend/server/src/core/storage-runtime/provider.ts b/packages/backend/server/src/core/storage-runtime/provider.ts index 4b05958a2..4ff84221f 100644 --- a/packages/backend/server/src/core/storage-runtime/provider.ts +++ b/packages/backend/server/src/core/storage-runtime/provider.ts @@ -258,6 +258,33 @@ export class StorageRuntimeProvider ); } + async reconcileWorkspaceDocuments(workspaceId: string) { + return await this.measured('reconcileWorkspaceDocuments', rt => + rt.reconcileWorkspaceDocuments(workspaceId) + ); + } + + async executeDocumentCleanupCandidates( + workspaceId: string | null | undefined, + gracePeriodDays: number, + limit: number + ) { + return await this.measured('executeDocumentCleanupCandidates', rt => + rt.executeDocumentCleanupCandidates(workspaceId, gracePeriodDays, limit) + ); + } + + async ackDocumentCleanupEffect( + workspaceId: string, + docId: string, + cleanupVersion: string, + effect: 'search' | 'copilot' + ) { + return await this.measured('ackDocumentCleanupEffect', rt => + rt.ackDocumentCleanupEffect(workspaceId, docId, cleanupVersion, effect) + ); + } + async planUnreferencedWorkspaceBlobs( workspaceId: string, gracePeriodDays: number, diff --git a/packages/backend/server/src/core/storage/__tests__/blob-job.spec.ts b/packages/backend/server/src/core/storage/__tests__/blob-job.spec.ts index 7dd1b778a..b1597f075 100644 --- a/packages/backend/server/src/core/storage/__tests__/blob-job.spec.ts +++ b/packages/backend/server/src/core/storage/__tests__/blob-job.spec.ts @@ -6,10 +6,13 @@ import { StorageBlobJob } from '../blob-job'; interface Context { runtime: { health: Sinon.SinonStub; + reconcileWorkspaceDocuments: Sinon.SinonStub; backfillMissingBlobMetadata: Sinon.SinonStub; rebuildWorkspaceDocBlobRefs: Sinon.SinonStub; planUnreferencedWorkspaceBlobs: Sinon.SinonStub; executeBlobCleanupCandidates: Sinon.SinonStub; + executeDocumentCleanupCandidates: Sinon.SinonStub; + ackDocumentCleanupEffect: Sinon.SinonStub; }; event: { emitAsync: Sinon.SinonStub; @@ -35,10 +38,18 @@ test.beforeEach(t => { providerConfigured: true, provider: 'fs', }), + reconcileWorkspaceDocuments: Sinon.stub().resolves({ + scannedDocs: 1, + marked: 0, + reset: 0, + recovered: 0, + }), backfillMissingBlobMetadata: Sinon.stub(), rebuildWorkspaceDocBlobRefs: Sinon.stub(), planUnreferencedWorkspaceBlobs: Sinon.stub(), executeBlobCleanupCandidates: Sinon.stub(), + executeDocumentCleanupCandidates: Sinon.stub(), + ackDocumentCleanupEffect: Sinon.stub(), }; t.context.event = { emitAsync: Sinon.stub().resolves(undefined), @@ -47,7 +58,18 @@ test.beforeEach(t => { add: Sinon.stub().resolves(undefined), }; t.context.db = { - $queryRaw: Sinon.stub(), + $queryRaw: Sinon.stub().resolves([ + { + marked: 0n, + failed: 0n, + effectsPending: 0n, + failedWorkspaceCheckpoints: 0n, + rootFailureCheckpoints: 0n, + staleProjectionBlocks: 0n, + oldestFailedSeconds: null, + oldestEffectsPendingSeconds: null, + }, + ]), workspace: { findMany: Sinon.stub(), }, @@ -129,46 +151,6 @@ for (const scenario of objectStorageRequiredCases) { }); } -test('doc blob refs sweep continues after one workspace fails', async t => { - t.context.db.workspace.findMany.resolves([ - { id: 'workspace-1', sid: 1 }, - { id: 'workspace-2', sid: 2 }, - ]); - t.context.runtime.rebuildWorkspaceDocBlobRefs - .onFirstCall() - .rejects(new Error('bad root doc')) - .onSecondCall() - .resolves({ - scannedDocs: 1, - parsedDocs: 1, - refsWritten: 0, - refsDeleted: 0, - failedDocs: 0, - nextCursor: null, - }); - - await t.context.job.rebuildWorkspaceDocBlobRefsBySid({ - workspaceLimit: 2, - docLimit: 100, - }); - - t.is(t.context.runtime.rebuildWorkspaceDocBlobRefs.callCount, 2); - t.deepEqual(t.context.runtime.rebuildWorkspaceDocBlobRefs.firstCall.args, [ - 'workspace-1', - 100, - ]); - t.deepEqual(t.context.runtime.rebuildWorkspaceDocBlobRefs.secondCall.args, [ - 'workspace-2', - 100, - ]); - t.true( - t.context.queue.add.calledWith( - 'backendRuntime.rebuildWorkspaceDocBlobRefsBySid', - { lastSid: 2, workspaceLimit: 2, docLimit: 100 } - ) - ); -}); - test('blob cleanup planning drains each workspace cursor before continuing', async t => { t.context.db.workspace.findMany.resolves([ { id: 'workspace-1', sid: 1 }, @@ -229,6 +211,13 @@ test('blob cleanup planning drains each workspace cursor before continuing', asy test('daily blob cleanup execution uses a fixed job id', async t => { await t.context.job.dailyBlobCleanupExecution(); + t.true( + t.context.queue.add.calledWith( + 'backendRuntime.executeDocumentCleanupCandidates', + {}, + { jobId: 'daily-backend-runtime-document-cleanup-execution' } + ) + ); t.true( t.context.queue.add.calledWith( 'backendRuntime.executeBlobCleanupCandidatesByMarkedRuns', @@ -238,6 +227,132 @@ test('daily blob cleanup execution uses a fixed job id', async t => { ); }); +test('daily storage reconciliation uses a fixed job id', async t => { + await t.context.job.dailyStorageReconciliation(); + + t.true( + t.context.queue.add.calledWith( + 'backendRuntime.reconcileWorkspaceStorageBySid', + {}, + { jobId: 'daily-backend-runtime-storage-reconciliation' } + ) + ); +}); + +test('storage reconciliation orders document retention before blob cleanup', async t => { + t.context.db.workspace.findMany.resolves([{ id: 'workspace-1', sid: 1 }]); + t.context.runtime.rebuildWorkspaceDocBlobRefs.resolves({ + scannedDocs: 1, + parsedDocs: 1, + refsWritten: 1, + refsDeleted: 0, + failedDocs: 0, + nextCursor: null, + }); + t.context.runtime.planUnreferencedWorkspaceBlobs.resolves({ + runId: 'run-1', + scannedBlobs: 1, + candidatesMarked: 0, + nextCursor: null, + }); + + await t.context.job.reconcileWorkspaceStorageBySid({ workspaceLimit: 10 }); + + Sinon.assert.callOrder( + t.context.runtime.reconcileWorkspaceDocuments, + t.context.runtime.rebuildWorkspaceDocBlobRefs, + t.context.runtime.planUnreferencedWorkspaceBlobs + ); + t.pass(); +}); + +test('storage reconciliation still refreshes document retention without object storage', async t => { + t.context.runtime.health.resolves({ + databaseConnected: true, + providerConfigured: true, + provider: undefined, + }); + t.context.db.workspace.findMany.resolves([{ id: 'workspace-1', sid: 1 }]); + t.context.runtime.rebuildWorkspaceDocBlobRefs.resolves({ + scannedDocs: 1, + parsedDocs: 1, + refsWritten: 0, + refsDeleted: 0, + failedDocs: 0, + nextCursor: null, + }); + + await t.context.job.reconcileWorkspaceStorageBySid({}); + + t.true(t.context.runtime.reconcileWorkspaceDocuments.calledOnce); + t.true(t.context.runtime.rebuildWorkspaceDocBlobRefs.calledOnce); + t.false(t.context.runtime.planUnreferencedWorkspaceBlobs.called); +}); + +test('document cleanup dispatches independent stable search and copilot effects', async t => { + t.context.runtime.executeDocumentCleanupCandidates.resolves({ + scannedCandidates: 1, + serializationRetries: 0, + executed: 1, + recovered: 0, + reset: 0, + failed: 0, + deletedRows: 3, + effects: [ + { + workspaceId: 'workspace-1', + docId: 'doc-1', + cleanupVersion: 'version-1', + commentObjectsDone: true, + searchDone: false, + copilotDone: false, + }, + ], + }); + + await t.context.job.executeDocumentCleanupCandidates({}); + + t.true( + t.context.queue.add.calledWith( + 'indexer.reconcileDocumentCleanup', + Sinon.match({ docId: 'doc-1' }), + { + jobId: 'document-cleanup:search:workspace-1:doc-1:version-1', + } + ) + ); + t.true( + t.context.queue.add.calledWith( + 'copilot.embedding.reconcileDocumentCleanup', + Sinon.match({ docId: 'doc-1' }), + { + jobId: 'document-cleanup:copilot:workspace-1:doc-1:version-1', + } + ) + ); + t.true( + t.context.event.emitAsync.calledWith('workspace.blobs.updated', { + workspaceId: 'workspace-1', + }) + ); +}); + +test('document cleanup effect ack delegates to storage runtime', async t => { + await t.context.job.ackDocumentCleanupEffect({ + workspaceId: 'workspace-1', + docId: 'doc-1', + cleanupVersion: 'version-1', + effect: 'search', + }); + + t.deepEqual(t.context.runtime.ackDocumentCleanupEffect.firstCall.args, [ + 'workspace-1', + 'doc-1', + 'version-1', + 'search', + ]); +}); + test('blob cleanup execution sweep drains marked runs and continues by page', async t => { t.context.db.$queryRaw .onFirstCall() diff --git a/packages/backend/server/src/core/storage/blob-job.ts b/packages/backend/server/src/core/storage/blob-job.ts index bdd54de88..f20887720 100644 --- a/packages/backend/server/src/core/storage/blob-job.ts +++ b/packages/backend/server/src/core/storage/blob-job.ts @@ -2,7 +2,7 @@ import { Injectable, Logger } from '@nestjs/common'; import { Cron, CronExpression } from '@nestjs/schedule'; import { PrismaClient } from '@prisma/client'; -import { EventBus, JobQueue, OnJob } from '../../base'; +import { EventBus, JobQueue, metrics, OnJob } from '../../base'; import { StorageRuntimeProvider } from '../storage-runtime'; // Queue keys are persisted API; keep the legacy backendRuntime.* names while @@ -18,15 +18,22 @@ declare global { workspaceLimit?: number; objectLimit?: number; }; - 'backendRuntime.rebuildWorkspaceDocBlobRefs': { - workspaceId: string; - limit?: number; - }; - 'backendRuntime.rebuildWorkspaceDocBlobRefsBySid': { + 'backendRuntime.reconcileWorkspaceStorageBySid': { lastSid?: number; workspaceLimit?: number; docLimit?: number; }; + 'backendRuntime.executeDocumentCleanupCandidates': { + workspaceId?: string; + gracePeriodDays?: number; + limit?: number; + }; + 'backendRuntime.ackDocumentCleanupEffect': { + workspaceId: string; + docId: string; + cleanupVersion: string; + effect: 'search' | 'copilot'; + }; 'backendRuntime.planUnreferencedWorkspaceBlobs': { workspaceId: string; gracePeriodDays?: number; @@ -89,25 +96,6 @@ export class StorageBlobJob { }); } - async enqueueRebuildWorkspaceDocBlobRefs(workspaceId: string, limit = 1000) { - await this.queue.add('backendRuntime.rebuildWorkspaceDocBlobRefs', { - workspaceId, - limit, - }); - } - - async enqueueRebuildWorkspaceDocBlobRefsBySid( - lastSid = 0, - workspaceLimit = 100, - docLimit = 1000 - ) { - await this.queue.add('backendRuntime.rebuildWorkspaceDocBlobRefsBySid', { - lastSid, - workspaceLimit, - docLimit, - }); - } - @OnJob('backendRuntime.backfillMissingBlobMetadataBySid') async backfillMissingBlobMetadataBySid({ lastSid = 0, @@ -211,25 +199,21 @@ export class StorageBlobJob { } @Cron(CronExpression.EVERY_DAY_AT_2AM) - async dailyDocBlobRefsRebuild() { + async dailyStorageReconciliation() { await this.queue.add( - 'backendRuntime.rebuildWorkspaceDocBlobRefsBySid', + 'backendRuntime.reconcileWorkspaceStorageBySid', {}, - { jobId: 'daily-backend-runtime-doc-blob-refs-rebuild' } - ); - } - - @Cron(CronExpression.EVERY_DAY_AT_3AM) - async dailyBlobCleanupPlanning() { - await this.queue.add( - 'backendRuntime.planUnreferencedWorkspaceBlobsBySid', - {}, - { jobId: 'daily-backend-runtime-blob-cleanup-planning' } + { jobId: 'daily-backend-runtime-storage-reconciliation' } ); } @Cron(CronExpression.EVERY_DAY_AT_4AM) async dailyBlobCleanupExecution() { + await this.queue.add( + 'backendRuntime.executeDocumentCleanupCandidates', + {}, + { jobId: 'daily-backend-runtime-document-cleanup-execution' } + ); await this.queue.add( 'backendRuntime.executeBlobCleanupCandidatesByMarkedRuns', {}, @@ -249,44 +233,39 @@ export class StorageBlobJob { await this.drainBlobMetadataBackfill(workspaceId, limit); } - @OnJob('backendRuntime.rebuildWorkspaceDocBlobRefs') - async rebuildWorkspaceDocBlobRefs({ - workspaceId, - limit = 1000, - }: Jobs['backendRuntime.rebuildWorkspaceDocBlobRefs']) { - await this.drainWorkspaceDocBlobRefs(workspaceId, limit); - } - - @OnJob('backendRuntime.rebuildWorkspaceDocBlobRefsBySid') - async rebuildWorkspaceDocBlobRefsBySid({ + @OnJob('backendRuntime.reconcileWorkspaceStorageBySid') + async reconcileWorkspaceStorageBySid({ lastSid = 0, workspaceLimit = 100, docLimit = 1000, - }: Jobs['backendRuntime.rebuildWorkspaceDocBlobRefsBySid']) { + }: Jobs['backendRuntime.reconcileWorkspaceStorageBySid']) { const workspaces = await this.db.workspace.findMany({ - where: { - sid: { - gt: lastSid, - }, - }, - orderBy: { - sid: 'asc', - }, - select: { - id: true, - sid: true, - }, + where: { sid: { gt: lastSid } }, + orderBy: { sid: 'asc' }, + select: { id: true, sid: true }, take: workspaceLimit, }); + const objectStorageConfigured = await this.hasObjectStorage( + 'storage reconciliation blob cleanup planning' + ); for (const workspace of workspaces) { try { + await this.rt.reconcileWorkspaceDocuments(workspace.id); await this.drainWorkspaceDocBlobRefs(workspace.id, docLimit, { sid: workspace.sid, }); + if (objectStorageConfigured) { + await this.drainBlobCleanupPlanning(workspace.id, 30, 1000, { + sid: workspace.sid, + }); + } } catch (err) { + metrics.storage + .counter('document_cleanup_workspace_failure_total') + .add(1); this.logger.error( - `doc blob refs rebuild failed workspace=${workspace.id} sid=${workspace.sid}`, + `storage reconciliation failed workspace=${workspace.id} sid=${workspace.sid}`, err ); } @@ -294,14 +273,71 @@ export class StorageBlobJob { const nextSid = workspaces.at(-1)?.sid; if (nextSid !== undefined && workspaces.length === workspaceLimit) { - await this.enqueueRebuildWorkspaceDocBlobRefsBySid( - nextSid, + await this.queue.add('backendRuntime.reconcileWorkspaceStorageBySid', { + lastSid: nextSid, workspaceLimit, - docLimit - ); + docLimit, + }); } } + @OnJob('backendRuntime.executeDocumentCleanupCandidates') + async executeDocumentCleanupCandidates({ + workspaceId, + gracePeriodDays = 30, + limit = 100, + }: Jobs['backendRuntime.executeDocumentCleanupCandidates']) { + const result = await this.rt.executeDocumentCleanupCandidates( + workspaceId, + gracePeriodDays, + limit + ); + metrics.storage + .counter('document_cleanup_serialization_retry_total') + .add(result.serializationRetries); + metrics.storage + .counter('document_cleanup_execute_failure_total') + .add(result.failed); + for (const effect of result.effects) { + if (!effect.searchDone) { + await this.queue.add('indexer.reconcileDocumentCleanup', effect, { + jobId: `document-cleanup:search:${effect.workspaceId}:${effect.docId}:${effect.cleanupVersion}`, + }); + } + if (!effect.copilotDone) { + await this.queue.add( + 'copilot.embedding.reconcileDocumentCleanup', + effect, + { + jobId: `document-cleanup:copilot:${effect.workspaceId}:${effect.docId}:${effect.cleanupVersion}`, + } + ); + } + if (effect.commentObjectsDone) { + await this.event.emitAsync('workspace.blobs.updated', { + workspaceId: effect.workspaceId, + }); + } + } + await this.recordDocumentCleanupHealth(); + return result; + } + + @OnJob('backendRuntime.ackDocumentCleanupEffect') + async ackDocumentCleanupEffect({ + workspaceId, + docId, + cleanupVersion, + effect, + }: Jobs['backendRuntime.ackDocumentCleanupEffect']) { + await this.rt.ackDocumentCleanupEffect( + workspaceId, + docId, + cleanupVersion, + effect + ); + } + @OnJob('backendRuntime.planUnreferencedWorkspaceBlobs') async planUnreferencedWorkspaceBlobs({ workspaceId, @@ -545,6 +581,82 @@ export class StorageBlobJob { return rows.map(row => row.runId); } + private async recordDocumentCleanupHealth() { + const [health] = await this.db.$queryRaw< + { + marked: bigint; + failed: bigint; + effectsPending: bigint; + failedWorkspaceCheckpoints: bigint; + rootFailureCheckpoints: bigint; + staleProjectionBlocks: bigint; + oldestFailedSeconds: number | null; + oldestEffectsPendingSeconds: number | null; + }[] + >` + SELECT + COUNT(*) FILTER (WHERE status = 'marked') AS marked, + COUNT(*) FILTER (WHERE status = 'failed') AS failed, + COUNT(*) FILTER (WHERE status = 'effects_pending') AS "effectsPending", + EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - + MIN(updated_at) FILTER (WHERE status = 'failed'))::double precision AS "oldestFailedSeconds", + EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - + MIN(updated_at) FILTER (WHERE status = 'effects_pending'))::double precision AS "oldestEffectsPendingSeconds", + (SELECT COUNT(*) FROM storage_reconciliation_checkpoints + WHERE kind = 'document_cleanup' AND status = 'failed') AS "failedWorkspaceCheckpoints", + (SELECT COUNT(*) FROM storage_reconciliation_checkpoints + WHERE kind = 'document_cleanup' AND status = 'failed' + AND metadata->>'failureKind' = 'root') AS "rootFailureCheckpoints", + (SELECT COUNT(*) FROM storage_reconciliation_runs + WHERE kind = 'blob_cleanup_plan' + AND started_at >= CURRENT_TIMESTAMP - INTERVAL '1 day' + AND jsonb_array_length(COALESCE(metadata->'staleOrFailedProjectionWorkspaces', '[]')) > 0 + ) AS "staleProjectionBlocks" + FROM document_cleanup_candidates + `; + if (!health) { + return; + } + for (const [name, value] of [ + ['document_cleanup_marked', health.marked], + ['document_cleanup_failed', health.failed], + ['document_cleanup_effects_pending', health.effectsPending], + [ + 'document_cleanup_failed_workspace_checkpoints', + health.failedWorkspaceCheckpoints, + ], + [ + 'document_cleanup_root_failure_checkpoints', + health.rootFailureCheckpoints, + ], + [ + 'document_cleanup_stale_projection_blocks', + health.staleProjectionBlocks, + ], + ] as const) { + metrics.storage.gauge(name).record(Number(value)); + } + metrics.storage + .gauge('document_cleanup_oldest_failed_seconds') + .record(health.oldestFailedSeconds ?? 0); + metrics.storage + .gauge('document_cleanup_oldest_effects_pending_seconds') + .record(health.oldestEffectsPendingSeconds ?? 0); + + if ( + health.failedWorkspaceCheckpoints > 0n || + health.rootFailureCheckpoints > 0n || + health.staleProjectionBlocks > 0n || + (health.oldestFailedSeconds ?? 0) >= 86_400 || + (health.oldestEffectsPendingSeconds ?? 0) >= 86_400 || + health.marked >= 10_000n + ) { + this.logger.warn( + `document cleanup health marked=${health.marked} failed=${health.failed} effectsPending=${health.effectsPending} failedWorkspaceCheckpoints=${health.failedWorkspaceCheckpoints} rootFailureCheckpoints=${health.rootFailureCheckpoints} staleProjectionBlocks=${health.staleProjectionBlocks} oldestFailedSeconds=${health.oldestFailedSeconds ?? 0} oldestEffectsPendingSeconds=${health.oldestEffectsPendingSeconds ?? 0}` + ); + } + } + private async hasMarkedBlobCleanupCandidates() { const rows = await this.db.$queryRaw<{ exists: boolean }[]>` SELECT EXISTS( diff --git a/packages/backend/server/src/core/utils/__tests__/blocksute.spec.ts b/packages/backend/server/src/core/utils/__tests__/blocksute.spec.ts index 15b606dc9..39e3ab69f 100644 --- a/packages/backend/server/src/core/utils/__tests__/blocksute.spec.ts +++ b/packages/backend/server/src/core/utils/__tests__/blocksute.spec.ts @@ -1,5 +1,6 @@ import test from 'ava'; import { omit } from 'lodash-es'; +import * as Y from 'yjs'; import { createModule } from '../../../__tests__/create-module'; import { Mockers } from '../../../__tests__/mocks'; @@ -38,6 +39,86 @@ test('can read all doc ids from workspace snapshot', async t => { t.snapshot(docIds); }); +test('merged root updates retain trash and exclude permanently removed docs', async t => { + const rootDoc = await models.doc.get(workspace.id, workspace.id); + t.truthy(rootDoc); + + const root = new Y.Doc(); + Y.applyUpdate(root, rootDoc!.blob); + const pending = new Y.Doc(); + Y.applyUpdate(pending, Y.encodeStateAsUpdate(root)); + const trashMeta = new Y.Map(); + trashMeta.set('id', 'trash-doc'); + trashMeta.set('title', 'Trash'); + trashMeta.set('trash', true); + (pending.getMap('meta').get('pages') as Y.Array>).push([ + trashMeta, + ]); + const pendingUpdate = Y.encodeStateAsUpdate( + pending, + Y.encodeStateVector(root) + ); + const merged = Y.mergeUpdates([rootDoc!.blob, pendingUpdate]); + + t.false(readAllDocIdsFromWorkspaceSnapshot(merged).includes('trash-doc')); + t.true( + readAllDocIdsFromWorkspaceSnapshot(merged, true).includes('trash-doc') + ); + + const removed = new Y.Doc(); + Y.applyUpdate(removed, merged); + const pages = removed.getMap('meta').get('pages') as Y.Array>; + const target = pages + .toArray() + .findIndex(meta => meta instanceof Y.Map && meta.get('id') === 'trash-doc'); + pages.delete(target, 1); + t.false( + readAllDocIdsFromWorkspaceSnapshot( + Y.encodeStateAsUpdate(removed), + true + ).includes('trash-doc') + ); +}); + +test('nested concurrent meta edits do not restore a deleted entry', t => { + const initial = new Y.Doc(); + const pages = new Y.Array>(); + const meta = new Y.Map(); + meta.set('id', 'doc-1'); + meta.set('title', 'Initial'); + pages.push([meta]); + initial.getMap('meta').set('pages', pages); + const state = Y.encodeStateAsUpdate(initial); + + const deletingClient = new Y.Doc(); + const editingClient = new Y.Doc(); + Y.applyUpdate(deletingClient, state); + Y.applyUpdate(editingClient, state); + ( + deletingClient.getMap('meta').get('pages') as Y.Array> + ).delete(0, 1); + (editingClient.getMap('meta').get('pages') as Y.Array>) + .get(0) + .set('title', 'Offline edit'); + + const merged = new Y.Doc(); + Y.applyUpdate(merged, Y.encodeStateAsUpdate(deletingClient)); + Y.applyUpdate(merged, Y.encodeStateAsUpdate(editingClient)); + t.is( + (merged.getMap('meta').get('pages') as Y.Array>).length, + 0 + ); + + (editingClient.getMap('meta').get('pages') as Y.Array>).push([ + meta.clone(), + ]); + Y.applyUpdate(merged, Y.encodeStateAsUpdate(editingClient)); + t.is( + (merged.getMap('meta').get('pages') as Y.Array>).length, + 1 + ); +}); + test('can read all blocks from doc snapshot', async t => { const rootDoc = await models.doc.get(workspace.id, workspace.id); t.truthy(rootDoc); diff --git a/packages/backend/server/src/core/utils/blocksuite.ts b/packages/backend/server/src/core/utils/blocksuite.ts index d0b66a368..241767c50 100644 --- a/packages/backend/server/src/core/utils/blocksuite.ts +++ b/packages/backend/server/src/core/utils/blocksuite.ts @@ -45,8 +45,11 @@ export function parsePageDoc( ); } -export function readAllDocIdsFromWorkspaceSnapshot(snapshot: Uint8Array) { - return readAllDocIdsFromRootDoc(Buffer.from(snapshot), false); +export function readAllDocIdsFromWorkspaceSnapshot( + snapshot: Uint8Array, + includeTrash = false +) { + return readAllDocIdsFromRootDoc(Buffer.from(snapshot), includeTrash); } function safeParseJson(str: string): T | undefined { diff --git a/packages/backend/server/src/models/copilot-context.ts b/packages/backend/server/src/models/copilot-context.ts index 4678bbdc9..4ea322302 100644 --- a/packages/backend/server/src/models/copilot-context.ts +++ b/packages/backend/server/src/models/copilot-context.ts @@ -337,10 +337,14 @@ export class CopilotContextModel extends BaseModel { } async deleteWorkspaceEmbedding(workspaceId: string, docId: string) { + await this.purgeWorkspaceEmbedding(workspaceId, docId); + await this.fulfillEmptyEmbedding(workspaceId, docId); + } + + async purgeWorkspaceEmbedding(workspaceId: string, docId: string) { await this.db.aiWorkspaceEmbedding.deleteMany({ where: { workspaceId, docId }, }); - await this.fulfillEmptyEmbedding(workspaceId, docId); } async matchWorkspaceEmbedding( diff --git a/packages/backend/server/src/plugins/copilot/embedding/job.ts b/packages/backend/server/src/plugins/copilot/embedding/job.ts index ee6635e07..5c6fd3f54 100644 --- a/packages/backend/server/src/plugins/copilot/embedding/job.ts +++ b/packages/backend/server/src/plugins/copilot/embedding/job.ts @@ -205,20 +205,49 @@ export class CopilotEmbeddingJob { ); } - @OnJob('copilot.embedding.deleteDoc') - async deleteDocEmbeddingQueueFromEvent( - doc: Jobs['copilot.embedding.deleteDoc'] - ) { + private async deleteDocEmbedding(doc: { + workspaceId: string; + docId: string; + }) { await this.queue.remove( `workspace:embedding:${doc.workspaceId}:${doc.docId}`, 'copilot.embedding.docs' ); - await this.models.copilotContext.deleteWorkspaceEmbedding( + await this.models.copilotContext.purgeWorkspaceEmbedding( doc.workspaceId, doc.docId ); } + @OnJob('copilot.embedding.reconcileDocumentCleanup') + async reconcileDocumentCleanup({ + workspaceId, + docId, + cleanupVersion, + }: Jobs['copilot.embedding.reconcileDocumentCleanup']) { + const root = await this.doc.getDoc(workspaceId, workspaceId); + if (!root) { + throw new Error(`workspace root ${workspaceId} not found`); + } + const live = readAllDocIdsFromWorkspaceSnapshot(root.bin, true).includes( + docId + ); + if (live) { + if (!(await this.doc.getDoc(workspaceId, docId))) { + throw new Error(`restored document ${workspaceId}/${docId} not found`); + } + await this.addDocEmbeddingQueueFromEvent({ workspaceId, docId }); + } else { + await this.deleteDocEmbedding({ workspaceId, docId }); + } + await this.queue.add('backendRuntime.ackDocumentCleanupEffect', { + workspaceId, + docId, + cleanupVersion, + effect: 'copilot', + }); + } + private async readCopilotBlob( userId: string, workspaceId: string, diff --git a/packages/backend/server/src/plugins/copilot/embedding/types.ts b/packages/backend/server/src/plugins/copilot/embedding/types.ts index 73b8cffd4..d1fe38838 100644 --- a/packages/backend/server/src/plugins/copilot/embedding/types.ts +++ b/packages/backend/server/src/plugins/copilot/embedding/types.ts @@ -68,9 +68,10 @@ declare global { docId: string; }; - 'copilot.embedding.deleteDoc': { + 'copilot.embedding.reconcileDocumentCleanup': { workspaceId: string; docId: string; + cleanupVersion: string; }; 'copilot.embedding.files': { diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index 863619d2e..d537bb405 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -39,10 +39,6 @@ declare global { 'copilot.session.generateTitle': { sessionId: string; }; - 'copilot.session.deleteDoc': { - workspaceId: string; - docId: string; - }; } } @@ -512,23 +508,6 @@ export class ChatSessionService { return null; } - @OnJob('copilot.session.deleteDoc') - async deleteDocSessions(doc: Jobs['copilot.session.deleteDoc']) { - const sessionIds = await this.models.copilotSession - .list({ - userId: undefined, - workspaceId: doc.workspaceId, - docId: doc.docId, - }) - .then(s => s.map(s => [s.userId, s.id])); - for (const [userId, sessionId] of sessionIds) { - await this.models.copilotSession.update( - { userId, sessionId, docId: null }, - true - ); - } - } - @OnJob('copilot.session.generateTitle') async generateSessionTitle(job: Jobs['copilot.session.generateTitle']) { const { sessionId } = job; diff --git a/packages/backend/server/src/plugins/indexer/__tests__/job.spec.ts b/packages/backend/server/src/plugins/indexer/__tests__/job.spec.ts index eacadd780..dff7e8ace 100644 --- a/packages/backend/server/src/plugins/indexer/__tests__/job.spec.ts +++ b/packages/backend/server/src/plugins/indexer/__tests__/job.spec.ts @@ -6,14 +6,17 @@ import Sinon from 'sinon'; import { createModule } from '../../../__tests__/create-module'; import { Mockers } from '../../../__tests__/mocks'; -import { JOB_SIGNAL } from '../../../base'; +import { Config, JOB_SIGNAL } from '../../../base'; import { ConfigModule } from '../../../base/config'; import { ServerConfigModule } from '../../../core/config'; +import { DocReader } from '../../../core/doc'; import { Models } from '../../../models'; +import { addDocToRootDoc } from '../../../native'; import { SearchProviderFactory } from '../factory'; import { IndexerModule, IndexerService } from '../index'; import { IndexerJob } from '../job'; import { ManticoresearchProvider } from '../providers'; +import { blockSQL, docSQL, SearchTable } from '../tables'; const module = await createModule({ imports: [ @@ -32,6 +35,8 @@ const indexerJob = module.get(IndexerJob); const searchProviderFactory = module.get(SearchProviderFactory); const manticoresearch = module.get(ManticoresearchProvider); const models = module.get(Models); +const docReader = module.get(DocReader); +const config = module.get(Config); const user = await module.create(Mockers.User); const workspace = await module.create(Mockers.Workspace, { @@ -39,6 +44,11 @@ const workspace = await module.create(Mockers.Workspace, { owner: user, }); +test.before(async () => { + await manticoresearch.recreateTable(SearchTable.block, blockSQL); + await manticoresearch.recreateTable(SearchTable.doc, docSQL); +}); + test.after.always(async () => { await module.close(); }); @@ -105,7 +115,7 @@ test('should not sync existing doc', async t => { t.is(module.queue.count('indexer.indexDoc'), count); }); -test('should delete doc from indexer when docId is not in workspace', async t => { +test('should delete dangling indexed docs absent from the root live set', async t => { const count = module.queue.count('indexer.deleteDoc'); mock.method(indexerService, 'listDocIds', async () => { return ['mock-doc-id1', 'mock-doc-id2']; @@ -121,6 +131,104 @@ test('should delete doc from indexer when docId is not in workspace', async t => t.is(module.queue.count('indexer.deleteDoc'), count + 2); }); +test('document cleanup reconcile deletes missing search state before ack', async t => { + const deleteSpy = Sinon.spy(indexerService, 'deleteDoc'); + const indexSpy = Sinon.spy(indexerService, 'indexDoc'); + const cleanupWorkspace = await module.create(Mockers.Workspace, { + owner: user, + }); + await module.create(Mockers.DocSnapshot, { + workspaceId: cleanupWorkspace.id, + docId: cleanupWorkspace.id, + user, + blob: addDocToRootDoc(Buffer.from([0, 0]), 'live-doc', 'Live'), + }); + + await indexerJob.reconcileDocumentCleanup({ + workspaceId: cleanupWorkspace.id, + docId: 'missing-doc', + cleanupVersion: 'version-1', + }); + + t.true(deleteSpy.calledOnceWith(cleanupWorkspace.id, 'missing-doc')); + t.false(indexSpy.called); + const { payload } = await module.queue.waitFor( + 'backendRuntime.ackDocumentCleanupEffect' + ); + t.deepEqual(payload, { + workspaceId: cleanupWorkspace.id, + docId: 'missing-doc', + cleanupVersion: 'version-1', + effect: 'search', + }); +}); + +test('document cleanup reconcile reindexes restored doc before ack', async t => { + const deleteSpy = Sinon.spy(indexerService, 'deleteDoc'); + const indexSpy = Sinon.spy(indexerService, 'indexDoc'); + const cleanupWorkspace = await module.create(Mockers.Workspace, { + owner: user, + }); + await module.create(Mockers.DocSnapshot, { + workspaceId: cleanupWorkspace.id, + docId: cleanupWorkspace.id, + user, + blob: addDocToRootDoc(Buffer.from([0, 0]), 'restored-doc', 'Restored'), + }); + await module.create(Mockers.DocSnapshot, { + workspaceId: cleanupWorkspace.id, + docId: 'restored-doc', + user, + }); + const getDocSpy = Sinon.spy(docReader, 'getDoc'); + + await indexerJob.reconcileDocumentCleanup({ + workspaceId: cleanupWorkspace.id, + docId: 'restored-doc', + cleanupVersion: 'version-2', + }); + + t.true(indexSpy.calledOnceWith(cleanupWorkspace.id, 'restored-doc')); + t.false(deleteSpy.called); + t.true(getDocSpy.calledWith(cleanupWorkspace.id, cleanupWorkspace.id)); + t.true(getDocSpy.calledWith(cleanupWorkspace.id, 'restored-doc')); + const { payload } = await module.queue.waitFor( + 'backendRuntime.ackDocumentCleanupEffect' + ); + t.deepEqual(payload, { + workspaceId: cleanupWorkspace.id, + docId: 'restored-doc', + cleanupVersion: 'version-2', + effect: 'search', + }); +}); + +test('document cleanup reconcile only acknowledges when indexer is disabled', async t => { + Sinon.stub(config.indexer, 'enabled').value(false); + const deleteSpy = Sinon.spy(indexerService, 'deleteDoc'); + const indexSpy = Sinon.spy(indexerService, 'indexDoc'); + const getDocSpy = Sinon.spy(docReader, 'getDoc'); + + await indexerJob.reconcileDocumentCleanup({ + workspaceId: workspace.id, + docId: 'disabled-doc', + cleanupVersion: 'version-disabled', + }); + + t.false(deleteSpy.called); + t.false(indexSpy.called); + t.false(getDocSpy.called); + const { payload } = await module.queue.waitFor( + 'backendRuntime.ackDocumentCleanupEffect' + ); + t.deepEqual(payload, { + workspaceId: workspace.id, + docId: 'disabled-doc', + cleanupVersion: 'version-disabled', + effect: 'search', + }); +}); + test('should handle indexer.deleteWorkspace job', async t => { const spy = Sinon.spy(indexerService, 'deleteWorkspace'); diff --git a/packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts b/packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts index e34ac55c1..fd2bd657e 100644 --- a/packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts +++ b/packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts @@ -1887,12 +1887,9 @@ test('should delete doc work', async t => { t.is(result4.nodes.length, 1); t.deepEqual(result4.nodes[0].fields.docId, [docId2]); - const count = module.queue.count('copilot.embedding.deleteDoc'); - await indexerService.deleteDoc(workspaceId, docId1, { refresh: true, }); - t.is(module.queue.count('copilot.embedding.deleteDoc'), count + 1); // make sure the docId1 is deleted result1 = await indexerService.search({ diff --git a/packages/backend/server/src/plugins/indexer/index.ts b/packages/backend/server/src/plugins/indexer/index.ts index 6e2afc817..44d655e97 100644 --- a/packages/backend/server/src/plugins/indexer/index.ts +++ b/packages/backend/server/src/plugins/indexer/index.ts @@ -3,6 +3,7 @@ import './config'; import { Module } from '@nestjs/common'; import { ServerConfigModule } from '../../core/config'; +import { DocStorageModule } from '../../core/doc'; import { PermissionModule } from '../../core/permission'; import { QuotaServiceModule } from '../../core/quota'; import { IndexerEvent } from './event'; @@ -13,7 +14,12 @@ import { IndexerResolver } from './resolver'; import { IndexerService } from './service'; @Module({ - imports: [ServerConfigModule, PermissionModule, QuotaServiceModule], + imports: [ + ServerConfigModule, + DocStorageModule, + PermissionModule, + QuotaServiceModule, + ], providers: [ IndexerResolver, IndexerService, diff --git a/packages/backend/server/src/plugins/indexer/job.ts b/packages/backend/server/src/plugins/indexer/job.ts index 9d1be70fb..0ce3d914a 100644 --- a/packages/backend/server/src/plugins/indexer/job.ts +++ b/packages/backend/server/src/plugins/indexer/job.ts @@ -1,6 +1,7 @@ import { Injectable, Logger } from '@nestjs/common'; import { Config, JOB_SIGNAL, JobQueue, OnJob } from '../../base'; +import { DocReader } from '../../core/doc'; import { readAllDocIdsFromWorkspaceSnapshot } from '../../core/utils/blocksuite'; import { Models } from '../../models'; import { IndexerService } from './service'; @@ -24,6 +25,11 @@ declare global { 'indexer.autoIndexWorkspaces': { lastIndexedWorkspaceSid?: number; }; + 'indexer.reconcileDocumentCleanup': { + workspaceId: string; + docId: string; + cleanupVersion: string; + }; } } @@ -35,7 +41,8 @@ export class IndexerJob { private readonly models: Models, private readonly service: IndexerService, private readonly queue: JobQueue, - private readonly config: Config + private readonly config: Config, + private readonly doc: DocReader ) {} @OnJob('indexer.indexDoc') @@ -66,6 +73,39 @@ export class IndexerJob { await this.service.deleteDoc(workspaceId, docId); } + @OnJob('indexer.reconcileDocumentCleanup') + async reconcileDocumentCleanup({ + workspaceId, + docId, + cleanupVersion, + }: Jobs['indexer.reconcileDocumentCleanup']) { + if (this.config.indexer.enabled) { + const root = await this.doc.getDoc(workspaceId, workspaceId); + if (!root) { + throw new Error(`workspace root ${workspaceId} not found`); + } + const live = readAllDocIdsFromWorkspaceSnapshot(root.bin, true).includes( + docId + ); + if (live) { + if (!(await this.doc.getDoc(workspaceId, docId))) { + throw new Error( + `restored document ${workspaceId}/${docId} not found` + ); + } + await this.service.indexDoc(workspaceId, docId); + } else { + await this.service.deleteDoc(workspaceId, docId); + } + } + await this.queue.add('backendRuntime.ackDocumentCleanupEffect', { + workspaceId, + docId, + cleanupVersion, + effect: 'search', + }); + } + @OnJob('indexer.indexWorkspace') async indexWorkspace({ workspaceId }: Jobs['indexer.indexWorkspace']) { if (!this.config.indexer.enabled) { @@ -79,16 +119,13 @@ export class IndexerJob { return; } - const snapshot = await this.models.doc.getSnapshot( - workspaceId, - workspaceId - ); - if (!snapshot) { + const root = await this.doc.getDoc(workspaceId, workspaceId); + if (!root) { this.logger.warn(`workspace snapshot ${workspaceId} not found`); return; } - const docIdsInWorkspace = readAllDocIdsFromWorkspaceSnapshot(snapshot.blob); + const docIdsInWorkspace = readAllDocIdsFromWorkspaceSnapshot(root.bin); const docIdsInIndexer = await this.service.listDocIds(workspaceId); const docIdsInWorkspaceSet = new Set(docIdsInWorkspace); diff --git a/packages/backend/server/src/plugins/indexer/service.ts b/packages/backend/server/src/plugins/indexer/service.ts index 2cf943d10..93885b722 100644 --- a/packages/backend/server/src/plugins/indexer/service.ts +++ b/packages/backend/server/src/plugins/indexer/service.ts @@ -374,14 +374,6 @@ export class IndexerService { ); await this.deleteBlocksByDocId(workspaceId, docId, options); - await this.queue.add('copilot.session.deleteDoc', { - workspaceId, - docId, - }); - await this.queue.add('copilot.embedding.deleteDoc', { - workspaceId, - docId, - }); this.logger.log(`deleted doc ${workspaceId}/${docId}`); }