diff --git a/src-tauri/crates/storage/src/lib.rs b/src-tauri/crates/storage/src/lib.rs index 69deb39e6f..57f98ce76e 100644 --- a/src-tauri/crates/storage/src/lib.rs +++ b/src-tauri/crates/storage/src/lib.rs @@ -601,6 +601,28 @@ impl FileStorage { self.replace_all_many_and_then(replacements, || Ok(())) } + pub fn append_many_uncached(&self, appends: Vec<(&str, Vec)>) -> AppResult { + self.ensure_writes_available()?; + let appends = appends + .into_iter() + .filter(|(_, rows)| !rows.is_empty()) + .collect::>(); + if appends.is_empty() { + return Ok(true); + } + + let _guard = self + .lock + .write() + .map_err(|_| AppError::new("lock_error", "Storage lock poisoned"))?; + self.ensure_writes_available()?; + if self.any_collection_dirty_cached(appends.iter().map(|(collection, _)| *collection))? { + return Ok(false); + } + + self.append_many_uncached_locked(appends) + } + pub fn update_collections_atomically( &self, collections: Vec<&str>, @@ -795,6 +817,45 @@ impl FileStorage { Ok(()) } + fn any_collection_dirty_cached<'a>( + &self, + collections: impl IntoIterator, + ) -> AppResult { + let cache = self + .cache + .read() + .map_err(|_| AppError::new("lock_error", "Storage cache lock poisoned"))?; + for collection in collections { + validate_collection_name(collection)?; + if cache + .collections + .get(collection) + .is_some_and(|cached| cached.dirty) + { + return Ok(true); + } + } + Ok(false) + } + + fn append_cached_collection_rows(&self, appends: &[(&str, Vec)]) -> AppResult<()> { + let mut cache = self + .cache + .write() + .map_err(|_| AppError::new("lock_error", "Storage cache lock poisoned"))?; + for (collection, rows) in appends { + validate_collection_name(collection)?; + cache + .projected_lists + .retain(|key, _| key.collection != *collection); + if let Some(cached) = cache.collections.get_mut(*collection) { + cached.rows.extend(rows.iter().cloned()); + cached.dirty = false; + } + } + Ok(()) + } + fn invalidate_projected_cache_for_collection(&self, collection: &str) -> AppResult<()> { validate_collection_name(collection)?; let mut cache = self @@ -1809,6 +1870,141 @@ impl FileStorage { Ok(()) } + fn append_many_uncached_locked(&self, appends: Vec<(&str, Vec)>) -> AppResult { + let transaction_id = storage_transaction_id(); + let mut pending = Vec::with_capacity(appends.len()); + let mut seen_paths = HashSet::new(); + let prepare_result = (|| -> AppResult { + for (index, (collection, rows)) in appends.iter().enumerate() { + let Some(item) = self.stage_appended_collection( + collection, + rows, + &transaction_id, + 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(false); + } + 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 append: {rollback_error}" + ), + )); + } + cleanup_pending_collection_transaction_files(&pending); + return Err(error); + } + + cleanup_pending_collection_transaction_files(&pending); + for (collection, _) in &appends { + self.invalidate_projected_cache_for_collection(collection)?; + } + self.append_cached_collection_rows(&appends)?; + Ok(true) + } + + fn stage_appended_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 append: {collection}" + ))); + } + if let Some(parent) = path.parent() { + fs::create_dir_all(parent)?; + } + 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()), + }; + let is_empty = !existed || fs::metadata(&path)?.len() == 0; + 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 = if is_empty { + (|| -> AppResult { + fs::write(&item.tmp, serde_json::to_vec_pretty(rows)?)?; + sync_file(&item.tmp)?; + Ok(true) + })() + } else { + stage_append_to_collection_file(&item.path, &item.tmp, rows) + }; + match staged { + Ok(true) => Ok(Some(item)), + Ok(false) => { + let _ = remove_path_if_exists(&item.tmp); + Ok(None) + } + Err(error) => { + let _ = remove_path_if_exists(&item.tmp); + Err(error) + } + } + } + fn recover_collection_after_read_error( &self, collection: &str, @@ -2034,6 +2230,67 @@ fn write_file_atomically(path: &Path, bytes: &[u8]) -> AppResult<()> { Ok(()) } +fn stage_append_to_collection_file(path: &Path, tmp: &Path, rows: &[Value]) -> AppResult { + if looks_nul_filled(path) { + return Ok(false); + } + + let mut file = fs::File::open(path)?; + let mut cursor = file.metadata()?.len(); + let mut byte = [0_u8; 1]; + let mut found_non_whitespace = false; + while cursor > 0 { + cursor -= 1; + file.seek(SeekFrom::Start(cursor))?; + file.read_exact(&mut byte)?; + if !byte[0].is_ascii_whitespace() { + found_non_whitespace = true; + break; + } + } + if !found_non_whitespace || byte[0] != b']' { + return Ok(false); + } + + let mut before_close = cursor; + let mut is_empty = false; + let mut found_array_prefix = false; + while before_close > 0 { + before_close -= 1; + file.seek(SeekFrom::Start(before_close))?; + file.read_exact(&mut byte)?; + if byte[0].is_ascii_whitespace() { + continue; + } + is_empty = byte[0] == b'['; + found_array_prefix = true; + break; + } + if !found_array_prefix { + return Ok(false); + } + + let mut source = fs::File::open(path)?; + let mut output = fs::File::create(tmp)?; + std::io::copy(&mut Read::by_ref(&mut source).take(cursor), &mut output)?; + for (index, row) in rows.iter().enumerate() { + let serialized = serde_json::to_string_pretty(row)?; + let indented = serialized + .lines() + .map(|line| format!(" {line}")) + .collect::>() + .join("\n"); + if is_empty && index == 0 { + output.write_all(format!("\n{indented}").as_bytes())?; + } else { + output.write_all(format!(",\n{indented}").as_bytes())?; + } + } + output.write_all(b"\n]\n")?; + output.sync_all()?; + Ok(true) +} + fn sync_file(path: &Path) -> AppResult<()> { fs::OpenOptions::new() .read(true) @@ -4059,6 +4316,155 @@ mod tests { fs::remove_dir_all(root).unwrap(); } + #[test] + fn append_many_uncached_appends_multiple_collections() { + let root = temp_storage_root("append-many-uncached"); + 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(); + fs::write( + collections.join("message-swipes.json"), + serde_json::to_vec_pretty(&json!([ + { "id": "message-1::swipe::0", "messageId": "message-1" } + ])) + .unwrap(), + ) + .unwrap(); + + let appended = storage + .append_many_uncached(vec![ + ("messages", vec![json!({ "id": "message-2" })]), + ( + "message-swipes", + vec![json!({ "id": "message-2::swipe::0", "messageId": "message-2" })], + ), + ]) + .unwrap(); + + assert!(appended); + assert_eq!( + parse_collection_file("messages", &collections.join("messages.json")) + .unwrap() + .len(), + 2 + ); + assert_eq!( + parse_collection_file("message-swipes", &collections.join("message-swipes.json")) + .unwrap() + .len(), + 2 + ); + + fs::remove_dir_all(root).unwrap(); + } + + #[test] + fn append_many_uncached_updates_clean_cached_collections() { + let root = temp_storage_root("append-many-cached"); + let storage = FileStorage::new(&root).unwrap(); + storage + .replace_all("messages", vec![json!({ "id": "message-1" })]) + .unwrap(); + storage + .replace_all( + "message-swipes", + vec![json!({ "id": "message-1::swipe::0", "messageId": "message-1" })], + ) + .unwrap(); + assert_eq!(storage.list("messages").unwrap().len(), 1); + assert_eq!(storage.list("message-swipes").unwrap().len(), 1); + + let appended = storage + .append_many_uncached(vec![ + ("messages", vec![json!({ "id": "message-2" })]), + ( + "message-swipes", + vec![json!({ "id": "message-2::swipe::0", "messageId": "message-2" })], + ), + ]) + .unwrap(); + + assert!(appended); + assert_eq!(storage.list("messages").unwrap().len(), 2); + assert_eq!(storage.list("message-swipes").unwrap().len(), 2); + assert_eq!( + parse_collection_file("messages", &root.join("collections").join("messages.json")) + .unwrap() + .len(), + 2 + ); + 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_uncached_refuses_dirty_cached_collections() { + let root = temp_storage_root("append-many-dirty-cached"); + let storage = FileStorage::new(&root).unwrap(); + storage + .cache_collection("messages", &[json!({ "id": "message-1" })], true) + .unwrap(); + + let appended = storage + .append_many_uncached(vec![("messages", vec![json!({ "id": "message-2" })])]) + .unwrap(); + + assert!(!appended); + assert_eq!(storage.list("messages").unwrap().len(), 1); + + 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"); + 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_uncached(vec![ + ("messages", vec![json!({ "id": "message-2" })]), + ("messages", vec![json!({ "id": "message-3" })]), + ]) + .expect_err("duplicate 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 update_collections_atomically_reads_and_replaces_multiple_collections() { let root = temp_storage_root("update-collections-atomically"); @@ -4579,7 +4985,10 @@ mod tests { storage .cache_collection("characters", &[json!({ "id": "pending" })], true) .unwrap(); - assert!(storage.dirty_collection_count() > 0, "write should be pending"); + assert!( + storage.dirty_collection_count() > 0, + "write should be pending" + ); storage.flush().unwrap(); diff --git a/src-tauri/src/commands/storage/message_swipes.rs b/src-tauri/src/commands/storage/message_swipes.rs index 57ee30bedc..928a1ba46c 100644 --- a/src-tauri/src/commands/storage/message_swipes.rs +++ b/src-tauri/src/commands/storage/message_swipes.rs @@ -307,9 +307,14 @@ fn normalize_message_rows_and_sidecars_inner( Ok((messages, sidecars, changed)) } -fn prepare_message_create_row(state: &AppState, value: Value) -> AppResult { +struct PreparedMessageCreate { + message: Value, + caller_supplied_id: bool, +} + +fn prepare_message_create_row(state: &AppState, value: Value) -> AppResult { let mut object = ensure_object(with_message_create_defaults(value)?)?; - let had_id = object + let caller_supplied_id = object .get("id") .and_then(Value::as_str) .is_some_and(|id| !id.trim().is_empty()); @@ -320,7 +325,7 @@ fn prepare_message_create_row(state: &AppState, value: Value) -> AppResult AppResult AppResult> { + let (_, stored_message) = message_row_for_write(message.clone(), false)?; + let sidecars = swipe_rows_for_message(&stored_message, swipes)?; + if !state.storage.append_many_uncached(vec![ + ("messages", vec![stored_message.clone()]), + (COLLECTION, sidecars.clone()), + ])? { + return Ok(None); + } + + let mut materialized = stored_message; + apply_sidecar_swipes( + &mut materialized, + &sidecars, + MessageSwipeMaterialization::full(), + ); + Ok(Some(materialized)) +} + +fn persist_created_message_with_swipes( + state: &AppState, + mut message: Value, + swipes: Vec, + caller_supplied_id: bool, +) -> AppResult { + clamp_message_active_swipe_index(&mut message, swipes.len()); + if !caller_supplied_id { + if let Some(updated) = + append_created_message_and_swipes_if_uncached(state, &message, &swipes)? + { + return Ok(updated); + } + } + + let mut updated = write_message_and_swipes(state, message, swipes, false)?; + materialize_message(state, &mut updated, true)?; + Ok(updated) +} + fn materialized_message_from_loaded_rows( message: &Value, message_id: &str, @@ -999,11 +1050,15 @@ pub(crate) fn delete_message_rows_for_chats_with_swipes( } pub(crate) fn create_message(state: &AppState, message: Value) -> AppResult { - let message = prepare_message_create_row(state, message)?; - persist_created_message_swipes(state, message) + let prepared = prepare_message_create_row(state, message)?; + persist_created_message_swipes(state, prepared.message, prepared.caller_supplied_id) } -fn persist_created_message_swipes(state: &AppState, mut message: Value) -> AppResult { +fn persist_created_message_swipes( + state: &AppState, + mut message: Value, + caller_supplied_id: bool, +) -> AppResult { if message.get("swipes").is_some() { let embedded_swipe_count = message .get("swipes") @@ -1017,16 +1072,10 @@ fn persist_created_message_swipes(state: &AppState, mut message: Value) -> AppRe if swipes.is_empty() { swipes.push(initial_swipe_for_message(&message)); } - clamp_message_active_swipe_index(&mut message, swipes.len()); - let mut updated = write_message_and_swipes(state, message, swipes, false)?; - materialize_message(state, &mut updated, true)?; - return Ok(updated); + return persist_created_message_with_swipes(state, message, swipes, caller_supplied_id); } let swipes = vec![initial_swipe_for_message(&message)]; - clamp_message_active_swipe_index(&mut message, swipes.len()); - let mut updated = write_message_and_swipes(state, message, swipes, false)?; - materialize_message(state, &mut updated, true)?; - Ok(updated) + persist_created_message_with_swipes(state, message, swipes, caller_supplied_id) } #[cfg(test)] @@ -1051,6 +1100,49 @@ mod tests { .expect("test app state should initialize") } + fn value_id(value: &Value) -> String { + value + .get("id") + .and_then(Value::as_str) + .expect("value should have an id") + .to_string() + } + + fn strip_generated_shape_fields(value: &mut Value) { + let Some(object) = value.as_object_mut() else { + return; + }; + for key in ["id", "messageId", "createdAt", "updatedAt"] { + object.remove(key); + } + } + + fn persisted_message_shape(state: &AppState, message_id: &str) -> Value { + let mut message = state + .storage + .get("messages", message_id) + .expect("message lookup should not fail") + .expect("message should exist"); + strip_generated_shape_fields(&mut message); + + let mut sidecars = state + .storage + .list(COLLECTION) + .expect("sidecars should list") + .into_iter() + .filter(|row| sidecar_matches_message_id(row, message_id)) + .collect::>(); + sort_sidecar_rows(&mut sidecars); + for sidecar in &mut sidecars { + strip_generated_shape_fields(sidecar); + } + + json!({ + "message": message, + "sidecars": sidecars + }) + } + #[test] fn migration_moves_nested_swipes_to_sidecar_and_strips_message_rows() { let root = temp_root("migrate"); @@ -1642,6 +1734,197 @@ mod tests { assert_eq!(sidecars[0]["messageId"], "message-2"); } + #[test] + fn create_message_with_caller_id_replaces_stale_sidecars() { + let state = test_state("create-caller-id-stale-sidecars"); + state + .storage + .replace_all( + COLLECTION, + vec![ + json!({ + "id": "message-1::swipe::0", + "messageId": "message-1", + "chatId": "chat-1", + "index": 0, + "content": "stale" + }), + json!({ + "id": "message-2::swipe::0", + "messageId": "message-2", + "chatId": "chat-1", + "index": 0, + "content": "unrelated" + }), + ], + ) + .expect("sidecars should seed"); + + let created = create_message( + &state, + json!({ + "id": "message-1", + "chatId": "chat-1", + "role": "assistant", + "content": "first", + "activeSwipeIndex": 1, + "swipes": [ + { "content": "first" }, + { "content": "second" } + ] + }), + ) + .expect("message should create"); + + assert_eq!(created["id"], json!("message-1")); + assert_eq!(created["content"], json!("second")); + + let messages = state + .storage + .list("messages") + .expect("messages should list"); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0]["id"], json!("message-1")); + + let sidecars = state + .storage + .list(COLLECTION) + .expect("sidecars should list"); + let replacement_sidecars = sidecars + .iter() + .filter(|row| sidecar_matches_message_id(row, "message-1")) + .collect::>(); + assert_eq!(replacement_sidecars.len(), 2); + assert_eq!(replacement_sidecars[0]["content"], json!("first")); + assert_eq!(replacement_sidecars[1]["content"], json!("second")); + assert!(!replacement_sidecars + .iter() + .any(|row| row.get("content") == Some(&json!("stale")))); + assert_eq!( + sidecars + .iter() + .filter(|row| sidecar_matches_message_id(row, "message-2")) + .count(), + 1 + ); + } + + #[test] + fn caller_id_persistence_replaces_existing_rows_after_precheck_window() { + let state = test_state("caller-id-post-precheck-collision"); + state + .storage + .replace_all( + "messages", + vec![json!({ + "id": "message-1", + "chatId": "chat-1", + "role": "assistant", + "content": "old", + "activeSwipeIndex": 0 + })], + ) + .expect("messages should seed"); + state + .storage + .replace_all( + COLLECTION, + vec![json!({ + "id": "message-1::swipe::0", + "messageId": "message-1", + "chatId": "chat-1", + "index": 0, + "content": "old" + })], + ) + .expect("sidecars should seed"); + + let updated = persist_created_message_swipes( + &state, + json!({ + "id": "message-1", + "chatId": "chat-1", + "role": "assistant", + "content": "new", + "activeSwipeIndex": 0, + "swipes": [{ "content": "new" }] + }), + true, + ) + .expect("post-precheck persistence should replace existing rows"); + + assert_eq!(updated["id"], json!("message-1")); + assert_eq!(updated["content"], json!("new")); + + let messages = state + .storage + .list("messages") + .expect("messages should list"); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0]["id"], json!("message-1")); + assert_eq!(messages[0]["content"], json!("new")); + + let sidecars = state + .storage + .list(COLLECTION) + .expect("sidecars should list"); + assert_eq!(sidecars.len(), 1); + assert_eq!(sidecars[0]["messageId"], json!("message-1")); + assert_eq!(sidecars[0]["content"], json!("new")); + } + + #[test] + fn fast_append_and_dirty_cache_fallback_persist_same_generated_message_shape() { + let input = json!({ + "chatId": "chat-shape", + "role": "assistant", + "content": "first", + "activeSwipeIndex": 3, + "extra": { + "hiddenFromAI": true, + "generationInfo": { "model": "shape-model" } + } + }); + + let fast_state = test_state("create-shape-fast"); + let prepared = + prepare_message_create_row(&fast_state, input.clone()).expect("message should prepare"); + let mut fast_message = prepared.message; + let fast_swipes = vec![initial_swipe_for_message(&fast_message)]; + clamp_message_active_swipe_index(&mut fast_message, fast_swipes.len()); + let fast_created = + append_created_message_and_swipes_if_uncached(&fast_state, &fast_message, &fast_swipes) + .expect("fast append should not fail") + .expect("clean generated message should use append fast path"); + let fast_id = value_id(&fast_created); + + let dirty_state = test_state("create-shape-dirty"); + dirty_state + .storage + .replace_all( + "messages", + vec![json!({ + "id": "dirty-seed", + "chatId": "chat-shape", + "role": "user", + "content": "seed" + })], + ) + .expect("dirty seed should write"); + dirty_state + .storage + .patch("messages", "dirty-seed", json!({ "content": "dirty seed" })) + .expect("dirty seed should patch"); + let dirty_created = + create_message(&dirty_state, input).expect("dirty fallback should create message"); + let dirty_id = value_id(&dirty_created); + + assert_eq!( + persisted_message_shape(&fast_state, &fast_id), + persisted_message_shape(&dirty_state, &dirty_id) + ); + } + #[test] fn create_message_keeps_response_compatible_but_persists_sidecar_swipes() { let state = test_state("create"); @@ -1890,20 +2173,16 @@ mod tests { assert_eq!(error.code, "invalid_input"); assert_eq!(error.message, "Message swipe content is required"); - assert!( - state - .storage - .list("messages") - .expect("messages should list") - .is_empty() - ); - assert!( - state - .storage - .list(COLLECTION) - .expect("sidecars should list") - .is_empty() - ); + assert!(state + .storage + .list("messages") + .expect("messages should list") + .is_empty()); + assert!(state + .storage + .list(COLLECTION) + .expect("sidecars should list") + .is_empty()); } #[test]