From c2c4b924f258149dd43a88cde10475b773a0daef Mon Sep 17 00:00:00 2001 From: Xel Date: Sun, 7 Jun 2026 03:32:19 +1000 Subject: [PATCH 1/2] Add chat swipe append fast path --- src-tauri/crates/storage/src/lib.rs | 324 ++++++++++++++++++ src-tauri/src/commands/storage/chats.rs | 211 +++++++++++- .../src/commands/storage/message_swipes.rs | 143 ++++++++ 3 files changed, 671 insertions(+), 7 deletions(-) diff --git a/src-tauri/crates/storage/src/lib.rs b/src-tauri/crates/storage/src/lib.rs index 57f98ce76e..251421b401 100644 --- a/src-tauri/crates/storage/src/lib.rs +++ b/src-tauri/crates/storage/src/lib.rs @@ -623,6 +623,37 @@ impl FileStorage { self.append_many_uncached_locked(appends) } + pub fn append_many_and_update_collections_uncached( + &self, + appends: Vec<(&str, Vec)>, + update_collections: Vec<&str>, + update: F, + ) -> AppResult> + where + F: FnOnce(&mut [AtomicCollectionRows]) -> AppResult>, + { + self.ensure_writes_available()?; + let appends = appends + .into_iter() + .filter(|(_, rows)| !rows.is_empty()) + .collect::>(); + let _guard = self + .lock + .write() + .map_err(|_| AppError::new("lock_error", "Storage lock poisoned"))?; + self.ensure_writes_available()?; + if self.any_collection_dirty_cached( + update_collections + .iter() + .copied() + .chain(appends.iter().map(|(collection, _)| *collection)), + )? { + return Ok(None); + } + + self.append_many_and_update_collections_uncached_locked(appends, update_collections, update) + } + pub fn update_collections_atomically( &self, collections: Vec<&str>, @@ -1944,6 +1975,171 @@ impl FileStorage { Ok(true) } + fn append_many_and_update_collections_uncached_locked( + &self, + appends: Vec<(&str, Vec)>, + update_collections: Vec<&str>, + update: F, + ) -> AppResult> + where + F: FnOnce(&mut [AtomicCollectionRows]) -> AppResult>, + { + let mut loaded = Vec::with_capacity(update_collections.len()); + let mut seen_update_paths = HashSet::new(); + for collection in update_collections { + let path = self.collection_path(collection)?; + if !seen_update_paths.insert(path) { + return Err(AppError::invalid_input(format!( + "Duplicate collection update: {collection}" + ))); + } + loaded.push(AtomicCollectionRows { + collection: collection.to_string(), + rows: self.read_collection_no_recovery(collection)?, + }); + } + + let Some(output) = update(&mut loaded)? else { + return Ok(None); + }; + + let transaction_id = storage_transaction_id(); + let mut pending = Vec::with_capacity(loaded.len() + appends.len()); + let mut seen_paths = HashSet::new(); + let prepare_result = (|| -> AppResult { + for (index, entry) in loaded.iter().enumerate() { + let item = self.stage_replaced_collection( + entry.collection(), + entry.rows(), + &transaction_id, + index, + &mut seen_paths, + )?; + pending.push(item); + } + let append_offset = loaded.len(); + for (index, (collection, rows)) in appends.iter().enumerate() { + let Some(item) = self.stage_appended_collection( + collection, + rows, + &transaction_id, + append_offset + index, + &mut seen_paths, + )? + else { + return Ok(false); + }; + pending.push(item); + } + Ok(true) + })(); + match prepare_result { + Ok(true) => {} + Ok(false) => { + cleanup_pending_collection_temps(&pending); + return Ok(None); + } + Err(error) => { + cleanup_pending_collection_temps(&pending); + return Err(error); + } + } + + let mut backed_up = Vec::new(); + let mut installed = Vec::new(); + let result = (|| -> AppResult<()> { + for (index, item) in pending.iter().enumerate() { + if !item.existed { + continue; + } + refresh_collection_backup(&item.path)?; + fs::rename(&item.path, &item.backup)?; + backed_up.push(index); + } + for (index, item) in pending.iter().enumerate() { + fs::rename(&item.tmp, &item.path)?; + installed.push(index); + } + Ok(()) + })(); + + if let Err(error) = result { + if let Err(rollback_error) = + rollback_collection_replacements(&pending, &backed_up, &installed) + { + cleanup_pending_collection_temps(&pending); + return Err(AppError::new( + "storage_rollback_failed", + format!( + "{error}; additionally failed to roll back collection update append: {rollback_error}" + ), + )); + } + cleanup_pending_collection_transaction_files(&pending); + return Err(error); + } + + cleanup_pending_collection_transaction_files(&pending); + for entry in &loaded { + self.invalidate_projected_cache_for_collection(entry.collection())?; + self.cache_collection(entry.collection(), entry.rows(), false)?; + } + self.append_cached_collection_rows(&appends)?; + Ok(Some(output)) + } + + fn stage_replaced_collection( + &self, + collection: &str, + rows: &[Value], + transaction_id: &str, + index: usize, + seen_paths: &mut HashSet, + ) -> AppResult { + let path = self.collection_path(collection)?; + if !seen_paths.insert(path.clone()) { + return Err(AppError::invalid_input(format!( + "Duplicate collection update: {collection}" + ))); + } + let existed = match fs::symlink_metadata(&path) { + Ok(metadata) => { + if !metadata.file_type().is_file() { + return Err(AppError::io(std::io::Error::other(format!( + "Collection path is not a regular file: {}", + path.display() + )))); + } + true + } + Err(error) if error.kind() == ErrorKind::NotFound => false, + Err(error) => return Err(error.into()), + }; + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + let tmp = collection_transaction_path(&path, transaction_id, index, "tmp")?; + let backup = collection_transaction_path(&path, transaction_id, index, "backup")?; + let item = PendingCollectionReplacement { + path, + tmp, + backup, + existed, + }; + let staged = (|| -> AppResult<()> { + fs::write(&item.tmp, serde_json::to_vec_pretty(rows)?)?; + sync_file(&item.tmp)?; + Ok(()) + })(); + match staged { + Ok(()) => Ok(item), + Err(error) => { + let _ = remove_path_if_exists(&item.tmp); + Err(error) + } + } + } + fn stage_appended_collection( &self, collection: &str, @@ -4428,6 +4624,134 @@ mod tests { fs::remove_dir_all(root).unwrap(); } + #[test] + fn append_many_and_update_collections_uncached_updates_and_appends_clean_cache() { + let root = temp_storage_root("append-update-clean-cache"); + let storage = FileStorage::new(&root).unwrap(); + storage + .replace_all( + "messages", + vec![ + json!({ "id": "message-1", "content": "first" }), + json!({ "id": "message-2", "content": "untouched" }), + ], + ) + .unwrap(); + storage + .replace_all( + "message-swipes", + vec![json!({ "id": "message-1::swipe::0", "messageId": "message-1" })], + ) + .unwrap(); + assert_eq!(storage.list("messages").unwrap().len(), 2); + assert_eq!(storage.list("message-swipes").unwrap().len(), 1); + + let output = storage + .append_many_and_update_collections_uncached( + vec![( + "message-swipes", + vec![json!({ "id": "message-1::swipe::1", "messageId": "message-1" })], + )], + vec!["messages"], + |collections| { + let message = collections[0] + .rows_mut() + .iter_mut() + .find(|row| row.get("id").and_then(Value::as_str) == Some("message-1")) + .expect("message row should exist"); + message["content"] = json!("second"); + Ok(Some(message.clone())) + }, + ) + .unwrap() + .expect("clean collections should use fast path"); + + assert_eq!(output["content"], json!("second")); + assert_eq!( + storage.list("messages").unwrap()[0]["content"], + json!("second") + ); + assert_eq!(storage.list("message-swipes").unwrap().len(), 2); + assert_eq!( + parse_collection_file("messages", &root.join("collections").join("messages.json")) + .unwrap()[0]["content"], + json!("second") + ); + assert_eq!( + parse_collection_file( + "message-swipes", + &root.join("collections").join("message-swipes.json") + ) + .unwrap() + .len(), + 2 + ); + + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn append_many_and_update_collections_uncached_refuses_dirty_cached_collections() { + let root = temp_storage_root("append-update-dirty-cache"); + let storage = FileStorage::new(&root).unwrap(); + storage + .cache_collection("messages", &[json!({ "id": "message-1" })], true) + .unwrap(); + + let updated = storage + .append_many_and_update_collections_uncached( + vec![("message-swipes", vec![json!({ "id": "swipe-1" })])], + vec!["messages"], + |_| -> AppResult> { + panic!("dirty-cache fallback should skip update closure") + }, + ) + .unwrap(); + + assert!(updated.is_none()); + assert_eq!(storage.list("messages").unwrap().len(), 1); + + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn append_many_and_update_collections_uncached_cleans_prepared_temps_on_stage_error() { + let root = temp_storage_root("append-update-stage-error-cleanup"); + let storage = FileStorage::new(&root).unwrap(); + let collections = root.join("collections"); + fs::write( + collections.join("messages.json"), + serde_json::to_vec_pretty(&json!([{ "id": "message-1" }])).unwrap(), + ) + .unwrap(); + + let error = storage + .append_many_and_update_collections_uncached( + vec![("messages", vec![json!({ "id": "message-2" })])], + vec!["messages"], + |collections| Ok(Some(collections[0].rows()[0].clone())), + ) + .expect_err("duplicate touched collection should fail staging"); + + assert!(error.message.contains("Duplicate collection append")); + let leftover_transaction_files = fs::read_dir(&collections) + .unwrap() + .filter_map(Result::ok) + .filter(|entry| { + entry + .file_name() + .to_string_lossy() + .contains(".profile-import-") + }) + .collect::>(); + assert!( + leftover_transaction_files.is_empty(), + "stage error should remove pending transaction files" + ); + + fs::remove_dir_all(root).unwrap(); + } + #[test] fn append_many_uncached_cleans_prepared_temps_on_stage_error() { let root = temp_storage_root("append-many-stage-error-cleanup"); diff --git a/src-tauri/src/commands/storage/chats.rs b/src-tauri/src/commands/storage/chats.rs index 4f0bbe0ee1..0c961bc891 100644 --- a/src-tauri/src/commands/storage/chats.rs +++ b/src-tauri/src/commands/storage/chats.rs @@ -717,7 +717,17 @@ pub(crate) fn message_swipes( body: Value, ) -> AppResult { let mut message = get_required(state, "messages", message_id)?; + let stored_has_embedded_swipes = message.get("swipes").and_then(Value::as_array).is_some(); message_swipe_storage::materialize_message(state, &mut message, true)?; + let existing_sidecar_swipe_count = if stored_has_embedded_swipes { + 0 + } else { + message + .get("swipes") + .and_then(Value::as_array) + .map(Vec::len) + .unwrap_or(0) + }; let owner_chat_id = owned_message_chat_id(&message, chat_id)?; if body.is_null() { return Ok(message.get("swipes").cloned().unwrap_or_else(|| json!([]))); @@ -756,7 +766,14 @@ pub(crate) fn message_swipes( .map(|value| value as usize) .unwrap_or(0); let activate_new_swipe = should_activate_new_swipe(&body); - let (active_index, swipe_count, active_content, active_extra, active_character_id) = { + let ( + active_index, + swipe_count, + active_content, + active_extra, + active_character_id, + previous_swipe_count, + ) = { let swipes = object .entry("swipes".to_string()) .or_insert_with(|| json!([])) @@ -803,6 +820,7 @@ pub(crate) fn message_swipes( swipes[active_index]["content"].clone(), swipes[active_index]["extra"].clone(), swipes[active_index].get("characterId").cloned(), + previous_swipe_count, ) }; let visible_content_changed = @@ -818,6 +836,33 @@ pub(crate) fn message_swipes( object.insert("characterId".to_string(), character_id); } let swipes = message_swipe_storage::take_swipes_for_storage(&mut message)?.unwrap_or_default(); + if !stored_has_embedded_swipes { + let mut extra_collections = Vec::new(); + if visible_content_changed { + extra_collections.push("chats"); + } + if let Some(updated) = + message_swipe_storage::append_message_swipes_and_update_collections_if_uncached( + state, + message.clone(), + swipes.clone(), + existing_sidecar_swipe_count.min(previous_swipe_count), + extra_collections, + |collections, written_message| { + if visible_content_changed { + apply_message_memory_invalidation_in_collections( + collections, + &owner_chat_id, + written_message, + )?; + } + Ok(()) + }, + )? + { + return Ok(updated); + } + } let updated = replace_message_with_swipes_and_chat_cleanup( state, &owner_chat_id, @@ -1580,12 +1625,17 @@ fn apply_message_memory_invalidation_in_collections( chat_id: &str, message: &Value, ) -> AppResult<()> { - let Some(chat) = collections.get_mut(2).and_then(|collection| { - collection - .rows_mut() - .iter_mut() - .find(|row| row.get("id").and_then(Value::as_str) == Some(chat_id)) - }) else { + let Some(chat_collection) = collections + .iter_mut() + .find(|collection| collection.collection() == "chats") + else { + return Ok(()); + }; + let Some(chat) = chat_collection + .rows_mut() + .iter_mut() + .find(|row| row.get("id").and_then(Value::as_str) == Some(chat_id)) + else { return Ok(()); }; apply_chat_memory_invalidation_from_message(chat, message) @@ -4247,6 +4297,153 @@ mod tests { ); } + #[test] + fn message_swipes_active_append_updates_sidecar_backed_message_and_prunes_memories() { + let state = test_state("swipe-active-sidecar-backed-memory-prune"); + state + .storage + .create( + "chats", + json!({ + "id": "chat-1", + "name": "Memory chat", + "memories": [ + { + "id": "keep-before", + "messageIds": ["message-before"], + "lastMessageAt": "2026-06-01T09:00:00.000Z" + }, + { + "id": "drop-edited", + "messageIds": ["message-1"], + "lastMessageAt": "2026-06-01T10:00:00.000Z" + }, + { + "id": "drop-newer", + "lastMessageAt": "2026-06-01T10:01:00.000Z" + } + ] + }), + ) + .expect("chat should seed"); + message_swipe_storage::create_message( + &state, + json!({ + "id": "message-1", + "chatId": "chat-1", + "role": "assistant", + "content": "first", + "createdAt": "2026-06-01T10:00:00.000Z", + "activeSwipeIndex": 0, + "swipes": [{ + "content": "first", + "extra": { "generationInfo": { "model": "first-model" } } + }] + }), + ) + .expect("message should seed"); + + let updated = message_swipes( + &state, + "POST", + "chat-1", + "message-1", + json!({ + "content": "second", + "extra": { "generationInfo": { "model": "second-model" } } + }), + ) + .expect("swipe should be added"); + + assert_eq!(updated["activeSwipeIndex"], json!(1)); + assert_eq!(updated["swipeCount"], json!(2)); + assert_eq!(updated["content"], json!("second")); + assert_eq!( + updated["extra"]["generationInfo"]["model"], + json!("second-model") + ); + + let persisted = stored_message(&state); + assert!(persisted.get("swipes").is_none()); + assert_eq!(persisted["content"], json!("second")); + assert_eq!(persisted["activeSwipeIndex"], json!(1)); + assert!(persisted.get("swipeCount").is_none()); + + let persisted_swipes = message_swipe_storage::swipes_for_message(&state, "message-1") + .expect("message sidecar swipes should read"); + assert_eq!(persisted_swipes.len(), 2); + assert_eq!(persisted_swipes[0]["content"], json!("first")); + assert_eq!(persisted_swipes[1]["content"], json!("second")); + assert_eq!( + persisted_swipes[1]["extra"]["generationInfo"]["model"], + json!("second-model") + ); + assert_eq!( + memory_ids(&stored_chat(&state)["memories"]), + vec!["keep-before"] + ); + } + + #[test] + fn message_swipes_append_cleans_trimmed_legacy_sidecar_rows() { + let state = test_state("swipe-append-clean-trimmed-sidecar"); + state + .storage + .replace_all( + "messages", + vec![json!({ + "id": "message-1", + "chatId": "chat-1", + "role": "assistant", + "content": "old parent", + "createdAt": "2026-06-01T10:00:00.000Z", + "activeSwipeIndex": 0 + })], + ) + .expect("message should seed"); + state + .storage + .replace_all( + message_swipe_storage::COLLECTION, + vec![json!({ + "id": "message-1::swipe::0", + "chatId": "chat-1", + "messageId": " message-1 ", + "index": 0, + "content": "legacy first" + })], + ) + .expect("legacy sidecar should seed"); + + let updated = message_swipes( + &state, + "POST", + "chat-1", + "message-1", + json!({ "content": "second" }), + ) + .expect("swipe should append"); + + assert_eq!(updated["activeSwipeIndex"], json!(1)); + assert_eq!(updated["swipeCount"], json!(2)); + assert_eq!(updated["swipes"][0]["content"], json!("legacy first")); + assert_eq!(updated["swipes"][1]["content"], json!("second")); + + let sidecars = state + .storage + .list(message_swipe_storage::COLLECTION) + .expect("sidecars should list"); + assert_eq!(sidecars.len(), 2); + assert_eq!(sidecars[0]["id"], json!("message-1::swipe::0")); + assert_eq!(sidecars[0]["messageId"], json!("message-1")); + assert_eq!(sidecars[0]["index"], json!(0)); + assert_eq!(sidecars[0]["content"], json!("legacy first")); + assert_eq!(sidecars[1]["id"], json!("message-1::swipe::1")); + assert_eq!(sidecars[1]["messageId"], json!("message-1")); + assert_eq!(sidecars[1]["index"], json!(1)); + assert_eq!(sidecars[1]["content"], json!("second")); + } + #[test] fn message_swipes_read_rejects_message_from_another_chat() { let state = test_state("swipe-read-cross-chat-owner"); diff --git a/src-tauri/src/commands/storage/message_swipes.rs b/src-tauri/src/commands/storage/message_swipes.rs index 928a1ba46c..fc3310d5c3 100644 --- a/src-tauri/src/commands/storage/message_swipes.rs +++ b/src-tauri/src/commands/storage/message_swipes.rs @@ -216,6 +216,39 @@ fn sort_sidecar_rows(rows: &mut [Value]) { }); } +fn sidecars_are_canonical_for_append( + rows: &mut [Value], + message_id: &str, + expected_count: usize, +) -> bool { + if rows.len() != expected_count { + return false; + } + sort_swipes(rows); + rows.iter().enumerate().all(|(index, row)| { + let expected_id = sidecar_row_id(message_id, index); + row.get("messageId").and_then(Value::as_str) == Some(message_id) + && row.get("id").and_then(Value::as_str) == Some(expected_id.as_str()) + && row.get("index").and_then(Value::as_u64) == Some(index as u64) + }) +} + +fn stored_sidecars_are_canonical_for_append( + state: &AppState, + message_id: &str, + expected_count: usize, +) -> AppResult { + let filter_values = HashSet::from([message_id.to_string()]); + let mut rows = state + .storage + .list_where_in(COLLECTION, "messageId", &filter_values)?; + Ok(sidecars_are_canonical_for_append( + &mut rows, + message_id, + expected_count, + )) +} + pub(crate) fn normalize_message_rows_and_sidecars( messages: Vec, sidecars: Vec, @@ -458,6 +491,60 @@ fn append_created_message_and_swipes_if_uncached( Ok(Some(materialized)) } +pub(crate) fn append_message_swipes_and_update_collections_if_uncached( + state: &AppState, + message: Value, + swipes: Vec, + append_start_index: usize, + extra_collections: Vec<&str>, + update_collections: F, +) -> AppResult> +where + F: FnOnce(&mut [AtomicCollectionRows], &Value) -> AppResult<()>, +{ + let (message_id, stored_message) = message_row_for_write(message, true)?; + let replacement = swipe_rows_for_message(&stored_message, &swipes)?; + if append_start_index > replacement.len() { + return Ok(None); + } + if !stored_sidecars_are_canonical_for_append(state, &message_id, append_start_index)? { + return Ok(None); + } + let appended = replacement[append_start_index..].to_vec(); + if appended.is_empty() { + return Ok(None); + } + + let mut collections = vec!["messages"]; + collections.extend(extra_collections); + let stored_message = state.storage.append_many_and_update_collections_uncached( + vec![(COLLECTION, appended)], + collections, + move |collections| { + let messages = collections[0].rows_mut(); + let Some(row) = messages + .iter_mut() + .find(|row| row.get("id").and_then(Value::as_str) == Some(message_id.as_str())) + else { + return Ok(None); + }; + *row = stored_message.clone(); + update_collections(collections, &stored_message)?; + Ok(Some(stored_message)) + }, + )?; + + let Some(mut materialized) = stored_message else { + return Ok(None); + }; + apply_sidecar_swipes( + &mut materialized, + &replacement, + MessageSwipeMaterialization::full(), + ); + Ok(Some(materialized)) +} + fn persist_created_message_with_swipes( state: &AppState, mut message: Value, @@ -1143,6 +1230,62 @@ mod tests { }) } + #[test] + fn append_fast_path_requires_canonical_contiguous_sidecars() { + let mut canonical = vec![ + json!({ + "id": "message-1::swipe::1", + "messageId": "message-1", + "index": 1, + "content": "second", + "customField": "preserved" + }), + json!({ + "id": "message-1::swipe::0", + "messageId": "message-1", + "index": 0, + "content": "first" + }), + ]; + assert!(sidecars_are_canonical_for_append( + &mut canonical, + "message-1", + 2 + )); + + let mut trimmed_message_id = vec![json!({ + "id": "message-1::swipe::0", + "messageId": " message-1 ", + "index": 0, + "content": "first" + })]; + assert!(!sidecars_are_canonical_for_append( + &mut trimmed_message_id, + "message-1", + 1 + )); + + let mut wrong_id = vec![json!({ + "id": "legacy-random", + "messageId": "message-1", + "index": 0, + "content": "first" + })]; + assert!(!sidecars_are_canonical_for_append( + &mut wrong_id, + "message-1", + 1 + )); + + let mut gap = vec![json!({ + "id": "message-1::swipe::1", + "messageId": "message-1", + "index": 1, + "content": "second" + })]; + assert!(!sidecars_are_canonical_for_append(&mut gap, "message-1", 1)); + } + #[test] fn migration_moves_nested_swipes_to_sidecar_and_strips_message_rows() { let root = temp_root("migrate"); From d4571a0f18c6fcde9b7c7bc28184916eb7835837 Mon Sep 17 00:00:00 2001 From: Xel Date: Sun, 7 Jun 2026 03:51:30 +1000 Subject: [PATCH 2/2] Harden swipe append transaction checks --- src-tauri/crates/storage/src/lib.rs | 365 ++++++++++++++++-- .../src/commands/storage/message_swipes.rs | 73 ++-- 2 files changed, 380 insertions(+), 58 deletions(-) diff --git a/src-tauri/crates/storage/src/lib.rs b/src-tauri/crates/storage/src/lib.rs index 251421b401..07d3aa38c3 100644 --- a/src-tauri/crates/storage/src/lib.rs +++ b/src-tauri/crates/storage/src/lib.rs @@ -631,27 +631,98 @@ impl FileStorage { ) -> AppResult> where F: FnOnce(&mut [AtomicCollectionRows]) -> AppResult>, + { + self.append_many_and_update_collections_uncached_with_append_check( + appends, + update_collections, + |_| Ok(true), + update, + ) + } + + pub fn append_many_and_update_collections_uncached_with_append_check( + &self, + appends: Vec<(&str, Vec)>, + update_collections: Vec<&str>, + validate_append_collections: C, + update: F, + ) -> AppResult> + where + F: FnOnce(&mut [AtomicCollectionRows]) -> AppResult>, + C: FnOnce(&mut [AtomicCollectionRows]) -> AppResult, { self.ensure_writes_available()?; let appends = appends .into_iter() .filter(|(_, rows)| !rows.is_empty()) .collect::>(); + let update_collections = update_collections + .into_iter() + .map(ToOwned::to_owned) + .collect::>(); + let _atomic_update = self.begin_atomic_update()?; + let (mut loaded, original_stamps) = { + let _guard = self + .lock + .write() + .map_err(|_| AppError::new("lock_error", "Storage lock poisoned"))?; + if self.any_collection_dirty_cached( + update_collections + .iter() + .map(String::as_str) + .chain(appends.iter().map(|(collection, _)| *collection)), + )? { + return Ok(None); + } + + let mut loaded = Vec::with_capacity(update_collections.len()); + let mut original_stamps = Vec::with_capacity(update_collections.len()); + let mut seen_update_paths = HashSet::new(); + for collection in &update_collections { + let path = self.collection_path(collection)?; + if !seen_update_paths.insert(path.clone()) { + return Err(AppError::invalid_input(format!( + "Duplicate collection update: {collection}" + ))); + } + loaded.push(AtomicCollectionRows { + collection: collection.to_string(), + rows: self.read_collection_no_recovery(collection)?, + }); + original_stamps.push((collection.to_string(), collection_file_stamp(&path)?)); + } + (loaded, original_stamps) + }; + + let Some(output) = update(&mut loaded)? else { + return Ok(None); + }; + let _guard = self .lock .write() .map_err(|_| AppError::new("lock_error", "Storage lock poisoned"))?; - self.ensure_writes_available()?; if self.any_collection_dirty_cached( update_collections .iter() - .copied() + .map(String::as_str) .chain(appends.iter().map(|(collection, _)| *collection)), )? { return Ok(None); } + for (collection, original_stamp) in &original_stamps { + let path = self.collection_path(collection)?; + if collection_file_stamp(&path)? != *original_stamp { + return Ok(None); + } + } - self.append_many_and_update_collections_uncached_locked(appends, update_collections, update) + self.append_many_and_update_collections_uncached_locked( + appends, + loaded, + validate_append_collections, + output, + ) } pub fn update_collections_atomically( @@ -1975,33 +2046,22 @@ impl FileStorage { Ok(true) } - fn append_many_and_update_collections_uncached_locked( + fn append_many_and_update_collections_uncached_locked( &self, appends: Vec<(&str, Vec)>, - update_collections: Vec<&str>, - update: F, + loaded: Vec, + validate_append_collections: C, + output: T, ) -> AppResult> where - F: FnOnce(&mut [AtomicCollectionRows]) -> AppResult>, + C: FnOnce(&mut [AtomicCollectionRows]) -> AppResult, { - let mut loaded = Vec::with_capacity(update_collections.len()); - let mut seen_update_paths = HashSet::new(); - for collection in update_collections { - let path = self.collection_path(collection)?; - if !seen_update_paths.insert(path) { - return Err(AppError::invalid_input(format!( - "Duplicate collection update: {collection}" - ))); - } - loaded.push(AtomicCollectionRows { - collection: collection.to_string(), - rows: self.read_collection_no_recovery(collection)?, - }); - } - - let Some(output) = update(&mut loaded)? else { + let Some(mut append_collections) = self.load_unique_append_collections(&appends)? else { return Ok(None); }; + if !validate_append_collections(&mut append_collections)? { + return Ok(None); + } let transaction_id = storage_transaction_id(); let mut pending = Vec::with_capacity(loaded.len() + appends.len()); @@ -2084,10 +2144,39 @@ impl FileStorage { self.invalidate_projected_cache_for_collection(entry.collection())?; self.cache_collection(entry.collection(), entry.rows(), false)?; } + for (collection, _) in &appends { + self.invalidate_projected_cache_for_collection(collection)?; + } self.append_cached_collection_rows(&appends)?; Ok(Some(output)) } + fn load_unique_append_collections( + &self, + appends: &[(&str, Vec)], + ) -> AppResult>> { + let mut loaded = Vec::new(); + let mut seen_paths = HashSet::new(); + for (collection, _) in appends { + let path = self.collection_path(collection)?; + if !seen_paths.insert(path) { + continue; + } + let rows = match self.read_collection_no_recovery(collection) { + Ok(rows) => rows, + Err(error) if matches!(error.code.as_str(), "invalid_input" | "json_error") => { + return Ok(None); + } + Err(error) => return Err(error), + }; + loaded.push(AtomicCollectionRows { + collection: (*collection).to_string(), + rows, + }); + } + Ok(Some(loaded)) + } + fn stage_replaced_collection( &self, collection: &str, @@ -2430,6 +2519,9 @@ fn stage_append_to_collection_file(path: &Path, tmp: &Path, rows: &[Value]) -> A if looks_nul_filled(path) { return Ok(false); } + if !collection_file_is_json_array(path)? { + return Ok(false); + } let mut file = fs::File::open(path)?; let mut cursor = file.metadata()?.len(); @@ -2487,6 +2579,17 @@ fn stage_append_to_collection_file(path: &Path, tmp: &Path, rows: &[Value]) -> A Ok(true) } +fn collection_file_is_json_array(path: &Path) -> AppResult { + let raw = fs::read_to_string(path)?; + if raw.trim().is_empty() { + return Ok(false); + } + match serde_json::from_str::>(&raw) { + Ok(_) => Ok(true), + Err(_) => Ok(false), + } +} + fn sync_file(path: &Path) -> AppResult<()> { fs::OpenOptions::new() .read(true) @@ -4690,6 +4793,177 @@ mod tests { fs::remove_dir_all(root).unwrap(); } + #[test] + fn append_many_and_update_collections_uncached_callback_reads_and_rejects_reentrant_writes() { + let root = temp_storage_root("append-update-callback-reentrant"); + let storage = FileStorage::new(&root).unwrap(); + storage + .replace_all( + "messages", + vec![json!({ "id": "message-1", "content": "first" })], + ) + .unwrap(); + storage.replace_all("message-swipes", Vec::new()).unwrap(); + + let error = storage + .append_many_and_update_collections_uncached( + vec![( + "message-swipes", + vec![json!({ "id": "message-1::swipe::0", "messageId": "message-1" })], + )], + vec!["messages"], + |collections| { + assert_eq!(storage.list("messages")?.len(), 1); + storage.create( + "personas", + json!({ "id": "persona-1", "name": "reentrant" }), + )?; + collections[0].rows_mut()[0]["content"] = json!("second"); + Ok(Some(())) + }, + ) + .expect_err("reentrant writes should fail instead of deadlocking"); + + assert_eq!(error.code, "storage_transaction_active"); + assert_eq!( + storage.get("messages", "message-1").unwrap().unwrap()["content"], + json!("first") + ); + assert!(storage.list("message-swipes").unwrap().is_empty()); + assert!(storage.list("personas").unwrap().is_empty()); + + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn append_many_and_update_collections_uncached_checks_append_rows_before_staging() { + let root = temp_storage_root("append-update-checked-appends"); + let storage = FileStorage::new(&root).unwrap(); + storage + .replace_all( + "messages", + vec![json!({ "id": "message-1", "content": "first" })], + ) + .unwrap(); + storage + .replace_all( + "message-swipes", + vec![json!({ + "id": "message-1::swipe::0", + "messageId": " message-1 ", + "index": 0 + })], + ) + .unwrap(); + + let output = storage + .append_many_and_update_collections_uncached_with_append_check( + vec![( + "message-swipes", + vec![json!({ + "id": "message-1::swipe::1", + "messageId": "message-1", + "index": 1 + })], + )], + vec!["messages"], + |append_collections| { + let sidecars = append_collections + .iter() + .find(|collection| collection.collection() == "message-swipes") + .expect("append collection should be loaded"); + Ok(sidecars.rows().iter().all(|row| { + row.get("messageId").and_then(Value::as_str) == Some("message-1") + })) + }, + |collections| { + collections[0].rows_mut()[0]["content"] = json!("second"); + Ok(Some(json!({ "updated": true }))) + }, + ) + .unwrap(); + + assert!(output.is_none()); + assert_eq!( + storage.get("messages", "message-1").unwrap().unwrap()["content"], + json!("first") + ); + let sidecars = storage.list("message-swipes").unwrap(); + assert_eq!(sidecars.len(), 1); + assert_eq!(sidecars[0]["messageId"], json!(" message-1 ")); + + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn append_many_and_update_collections_uncached_invalidates_append_projected_cache() { + let root = temp_storage_root("append-update-projected-cache"); + let storage = FileStorage::new(&root).unwrap(); + storage + .replace_all( + "messages", + vec![json!({ "id": "message-1", "content": "first" })], + ) + .unwrap(); + storage + .replace_all( + "message-swipes", + vec![json!({ + "id": "message-1::swipe::0", + "messageId": "message-1", + "index": 0, + "content": "first" + })], + ) + .unwrap(); + let fields = ["messageId", "index", "content"] + .into_iter() + .map(ToOwned::to_owned) + .collect::>(); + let values = HashSet::from(["message-1".to_string()]); + assert_eq!( + storage + .list_projected_where_in( + "message-swipes", + "messageId", + &values, + &fields, + &Map::new() + ) + .unwrap() + .len(), + 1 + ); + + storage + .append_many_and_update_collections_uncached( + vec![( + "message-swipes", + vec![json!({ + "id": "message-1::swipe::1", + "messageId": "message-1", + "index": 1, + "content": "second" + })], + )], + vec!["messages"], + |collections| { + collections[0].rows_mut()[0]["content"] = json!("second"); + Ok(Some(())) + }, + ) + .unwrap() + .expect("clean hybrid append should commit"); + + let projected = storage + .list_projected_where_in("message-swipes", "messageId", &values, &fields, &Map::new()) + .unwrap(); + assert_eq!(projected.len(), 2); + assert_eq!(projected[1]["content"], json!("second")); + + fs::remove_dir_all(root).unwrap(); + } + #[test] fn append_many_and_update_collections_uncached_refuses_dirty_cached_collections() { let root = temp_storage_root("append-update-dirty-cache"); @@ -4752,6 +5026,51 @@ mod tests { fs::remove_dir_all(root).unwrap(); } + #[test] + fn append_many_and_update_collections_uncached_falls_back_on_malformed_append_file() { + let root = temp_storage_root("append-update-malformed-append-file"); + let storage = FileStorage::new(&root).unwrap(); + storage + .replace_all( + "messages", + vec![json!({ "id": "message-1", "content": "first" })], + ) + .unwrap(); + let collections = root.join("collections"); + let sidecar_path = collections.join("message-swipes.json"); + fs::write( + &sidecar_path, + b"[{\"id\":\"message-1::swipe::0\",\"messageId\":\"message-1\"} trailing]", + ) + .unwrap(); + + let output = storage + .append_many_and_update_collections_uncached( + vec![( + "message-swipes", + vec![json!({ "id": "message-1::swipe::1", "messageId": "message-1" })], + )], + vec!["messages"], + |collections| { + collections[0].rows_mut()[0]["content"] = json!("second"); + Ok(Some(())) + }, + ) + .unwrap(); + + assert!(output.is_none()); + assert_eq!( + storage.get("messages", "message-1").unwrap().unwrap()["content"], + json!("first") + ); + assert_eq!( + fs::read_to_string(&sidecar_path).unwrap(), + "[{\"id\":\"message-1::swipe::0\",\"messageId\":\"message-1\"} trailing]" + ); + + fs::remove_dir_all(root).unwrap(); + } + #[test] fn append_many_uncached_cleans_prepared_temps_on_stage_error() { let root = temp_storage_root("append-many-stage-error-cleanup"); diff --git a/src-tauri/src/commands/storage/message_swipes.rs b/src-tauri/src/commands/storage/message_swipes.rs index fc3310d5c3..2d43c0b8af 100644 --- a/src-tauri/src/commands/storage/message_swipes.rs +++ b/src-tauri/src/commands/storage/message_swipes.rs @@ -233,22 +233,6 @@ fn sidecars_are_canonical_for_append( }) } -fn stored_sidecars_are_canonical_for_append( - state: &AppState, - message_id: &str, - expected_count: usize, -) -> AppResult { - let filter_values = HashSet::from([message_id.to_string()]); - let mut rows = state - .storage - .list_where_in(COLLECTION, "messageId", &filter_values)?; - Ok(sidecars_are_canonical_for_append( - &mut rows, - message_id, - expected_count, - )) -} - pub(crate) fn normalize_message_rows_and_sidecars( messages: Vec, sidecars: Vec, @@ -507,9 +491,6 @@ where if append_start_index > replacement.len() { return Ok(None); } - if !stored_sidecars_are_canonical_for_append(state, &message_id, append_start_index)? { - return Ok(None); - } let appended = replacement[append_start_index..].to_vec(); if appended.is_empty() { return Ok(None); @@ -517,22 +498,44 @@ where let mut collections = vec!["messages"]; collections.extend(extra_collections); - let stored_message = state.storage.append_many_and_update_collections_uncached( - vec![(COLLECTION, appended)], - collections, - move |collections| { - let messages = collections[0].rows_mut(); - let Some(row) = messages - .iter_mut() - .find(|row| row.get("id").and_then(Value::as_str) == Some(message_id.as_str())) - else { - return Ok(None); - }; - *row = stored_message.clone(); - update_collections(collections, &stored_message)?; - Ok(Some(stored_message)) - }, - )?; + let append_check_message_id = message_id.clone(); + let update_message_id = message_id.clone(); + let stored_message = state + .storage + .append_many_and_update_collections_uncached_with_append_check( + vec![(COLLECTION, appended)], + collections, + move |append_collections| { + let Some(sidecars) = append_collections + .iter_mut() + .find(|collection| collection.collection() == COLLECTION) + else { + return Ok(false); + }; + let mut matching_sidecars = sidecars + .rows() + .iter() + .filter(|row| sidecar_matches_message_id(row, &append_check_message_id)) + .cloned() + .collect::>(); + Ok(sidecars_are_canonical_for_append( + &mut matching_sidecars, + &append_check_message_id, + append_start_index, + )) + }, + move |collections| { + let messages = collections[0].rows_mut(); + let Some(row) = messages.iter_mut().find(|row| { + row.get("id").and_then(Value::as_str) == Some(update_message_id.as_str()) + }) else { + return Ok(None); + }; + *row = stored_message.clone(); + update_collections(collections, &stored_message)?; + Ok(Some(stored_message)) + }, + )?; let Some(mut materialized) = stored_message else { return Ok(None);