diff --git a/apps/codex-plus-launcher/src/main.rs b/apps/codex-plus-launcher/src/main.rs index 81b9bb134..700433a14 100644 --- a/apps/codex-plus-launcher/src/main.rs +++ b/apps/codex-plus-launcher/src/main.rs @@ -487,10 +487,12 @@ impl LauncherDataService { } fn storage_adapter(&self) -> codex_plus_data::SQLiteStorageAdapter { + let allowed_db_paths = self.candidate_db_paths(); codex_plus_data::SQLiteStorageAdapter::new( self.db_path.clone(), codex_plus_data::BackupStore::new(self.backup_dir.clone()), ) + .with_allowed_db_paths(allowed_db_paths) } } diff --git a/assets/inject/renderer-inject.js b/assets/inject/renderer-inject.js index 7f763c971..df8e3c50a 100644 --- a/assets/inject/renderer-inject.js +++ b/assets/inject/renderer-inject.js @@ -6681,6 +6681,7 @@ undo.addEventListener("click", async () => { const result = await postJson("/undo", { undo_token: undoToken }); toast.textContent = result.message || "撤销完成"; + if (result.status === "undone") window.location.reload(); setTimeout(() => toast.remove(), 5000); }); toast.appendChild(undo); @@ -7730,7 +7731,7 @@ const shouldReload = isCurrentSessionRow(row, ref); row.remove(); if (shouldReload) { - window.location.reload(); + setTimeout(() => window.location.reload(), 10000); } } diff --git a/crates/codex-plus-data/src/storage.rs b/crates/codex-plus-data/src/storage.rs index 41e3cc472..ab790e816 100644 --- a/crates/codex-plus-data/src/storage.rs +++ b/crates/codex-plus-data/src/storage.rs @@ -1,7 +1,7 @@ use crate::BackupStore; use codex_plus_core::models::{DeleteResult, DeleteStatus, SessionRef}; use rusqlite::types::{ToSqlOutput, Value as SqlValue, ValueRef}; -use rusqlite::{Connection, OptionalExtension, ToSql}; +use rusqlite::{Connection, OpenFlags, OptionalExtension, ToSql}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value, json}; use std::collections::HashSet; @@ -20,11 +20,15 @@ pub fn delete_local_from_paths( "Thread not found in local storage".to_string(), ); let mut deleted_count = 0usize; + let mut backup_tokens = Vec::new(); for db_path in db_paths { let adapter = SQLiteStorageAdapter::new(db_path, backup_store.clone()); let candidate_result = adapter.delete_local(session); if matches!(candidate_result.status, DeleteStatus::LocalDeleted) { deleted_count += 1; + if let Some(token) = candidate_result.undo_token.as_ref() { + backup_tokens.push(token.clone()); + } result = candidate_result; } else if deleted_count == 0 { result = candidate_result; @@ -32,6 +36,8 @@ pub fn delete_local_from_paths( } if deleted_count > 1 { result.message = format!("已从 {deleted_count} 个本地存储删除"); + result.undo_token = Some(json!(backup_tokens).to_string()); + result.backup_path = None; } result } @@ -58,6 +64,7 @@ pub fn move_codex_thread_workspace_from_paths( pub struct SQLiteStorageAdapter { db_path: PathBuf, backup_store: BackupStore, + allowed_db_paths: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -95,12 +102,23 @@ impl ToSql for OwnedSqlValue { impl SQLiteStorageAdapter { pub fn new(db_path: impl Into, backup_store: BackupStore) -> Self { + let db_path = db_path.into(); Self { - db_path: db_path.into(), + allowed_db_paths: vec![db_path.clone()], + db_path, backup_store, } } + pub fn with_allowed_db_paths(mut self, db_paths: impl IntoIterator) -> Self { + for db_path in db_paths { + if !self.allowed_db_paths.contains(&db_path) { + self.allowed_db_paths.push(db_path); + } + } + self + } + pub fn delete_local(&self, session: &SessionRef) -> DeleteResult { if !self.db_path.exists() { return failed( @@ -229,52 +247,9 @@ impl SQLiteStorageAdapter { pub fn undo(&self, token: &str) -> DeleteResult { let result = (|| -> anyhow::Result { - let backup = self.backup_store.read_backup(token)?; - let session_id = backup["session_id"].as_str().unwrap_or("").to_string(); - let mut db = Connection::open(&self.db_path)?; - if let Some(tables) = backup["tables"].as_object() { - validate_restore_tables(tables)?; - detect_restore_conflicts(&db, tables)?; - detect_file_restore_conflicts(tables)?; - let tx = db.transaction()?; - for (table, rows) in tables { - if table.starts_with("__") { - continue; - } - let Some(rows) = rows.as_array() else { - continue; - }; - for row in rows { - if let Some(row) = row.as_object() { - if table == "agent_job_items" - && update_existing_agent_job_item(&tx, row)? - { - continue; - } - insert_row(&tx, table, row)?; - } - } - } - tx.commit()?; - if let Some(files) = tables.get("__files").and_then(Value::as_array) { - for file in files { - let Some(path) = file.get("path").and_then(Value::as_str) else { - continue; - }; - let Some(content) = file.get("content_b64").and_then(Value::as_str) else { - continue; - }; - let bytes = base64::Engine::decode( - &base64::engine::general_purpose::STANDARD, - content, - )?; - if let Some(parent) = Path::new(path).parent() { - fs::create_dir_all(parent)?; - } - fs::write(path, bytes)?; - } - } - } + let backups = undo_backups(&self.backup_store, token)?; + let session_id = backups[0]["session_id"].as_str().unwrap_or("").to_string(); + restore_backups(&backups, &self.db_path, &self.allowed_db_paths)?; Ok(DeleteResult { status: DeleteStatus::Undone, session_id, @@ -896,6 +871,122 @@ fn normalize_codex_thread_id(session_id: &str) -> String { .to_string() } +fn undo_backups(backup_store: &BackupStore, token: &str) -> anyhow::Result> { + let tokens = + serde_json::from_str::>(token).unwrap_or_else(|_| vec![token.to_string()]); + if tokens.is_empty() { + anyhow::bail!("empty undo token"); + } + tokens + .into_iter() + .map(|token| backup_store.read_backup(&token)) + .collect() +} + +fn restore_backups( + backups: &[Value], + fallback_db_path: &Path, + allowed_db_paths: &[PathBuf], +) -> anyhow::Result<()> { + for backup in backups { + let Some(tables) = backup["tables"].as_object() else { + continue; + }; + let source_db = backup_source_db(backup, fallback_db_path, allowed_db_paths)?; + let db = Connection::open_with_flags(&source_db, OpenFlags::SQLITE_OPEN_READ_WRITE)?; + validate_restore_tables(tables)?; + detect_restore_conflicts(&db, tables)?; + detect_file_restore_conflicts(tables)?; + preflight_restore_rows(&db, tables)?; + } + + for backup in backups { + let Some(tables) = backup["tables"].as_object() else { + continue; + }; + let source_db = backup_source_db(backup, fallback_db_path, allowed_db_paths)?; + let mut db = Connection::open_with_flags(&source_db, OpenFlags::SQLITE_OPEN_READ_WRITE)?; + let tx = db.transaction()?; + restore_rows(&tx, tables)?; + tx.commit()?; + if let Some(files) = tables.get("__files").and_then(Value::as_array) { + for file in files { + let Some(path) = file.get("path").and_then(Value::as_str) else { + continue; + }; + let Some(content) = file.get("content_b64").and_then(Value::as_str) else { + continue; + }; + let bytes = + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, content)?; + if let Some(parent) = Path::new(path).parent() { + fs::create_dir_all(parent)?; + } + fs::write(path, bytes)?; + } + } + } + Ok(()) +} + +fn preflight_restore_rows(db: &Connection, tables: &Map) -> anyhow::Result<()> { + db.execute_batch("SAVEPOINT codex_plus_restore_preflight")?; + let restore_result = restore_rows(db, tables); + let rollback_result = db.execute_batch( + "ROLLBACK TO codex_plus_restore_preflight; RELEASE codex_plus_restore_preflight", + ); + restore_result?; + rollback_result?; + Ok(()) +} + +fn restore_rows(db: &Connection, tables: &Map) -> anyhow::Result<()> { + for (table, rows) in tables { + if table.starts_with("__") { + continue; + } + let Some(rows) = rows.as_array() else { + continue; + }; + for row in rows { + if let Some(row) = row.as_object() { + if table == "agent_job_items" && update_existing_agent_job_item(db, row)? { + continue; + } + insert_row(db, table, row)?; + } + } + } + Ok(()) +} + +fn backup_source_db( + backup: &Value, + fallback_db_path: &Path, + allowed_db_paths: &[PathBuf], +) -> anyhow::Result { + let source_db = backup["source_db"] + .as_str() + .filter(|path| !path.trim().is_empty()) + .map(PathBuf::from) + .unwrap_or_else(|| fallback_db_path.to_path_buf()); + if !source_db.is_file() { + anyhow::bail!( + "Backup source database not found: {}", + source_db.to_string_lossy() + ); + } + let source_db = fs::canonicalize(source_db)?; + let allowed = allowed_db_paths + .iter() + .filter_map(|path| fs::canonicalize(path).ok()) + .any(|path| path == source_db); + if !allowed { + anyhow::bail!("Backup source database is not an allowed local storage path"); + } + Ok(source_db) +} + fn schema_kind(db: &Connection) -> anyhow::Result> { if has_table(db, "sessions")? && has_columns(db, "sessions", &["id", "title"])? { if has_table(db, "messages")? && !has_columns(db, "messages", &["session_id"])? { @@ -1064,6 +1155,9 @@ fn detect_file_restore_conflicts(tables: &Map) -> anyhow::Result< if Path::new(path).exists() { anyhow::bail!("restore conflict: file already exists: {path}"); } + if let Some(content) = file.get("content_b64").and_then(Value::as_str) { + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, content)?; + } } } Ok(()) diff --git a/crates/codex-plus-data/tests/storage_adapter.rs b/crates/codex-plus-data/tests/storage_adapter.rs index 28b64a152..9fd1dfc70 100644 --- a/crates/codex-plus-data/tests/storage_adapter.rs +++ b/crates/codex-plus-data/tests/storage_adapter.rs @@ -451,6 +451,131 @@ fn delete_local_from_paths_removes_duplicate_threads_from_all_databases() { assert!(!second_rollout.exists()); } +#[test] +fn delete_local_from_paths_undo_restores_duplicate_threads_and_shared_rollout_to_source_databases() +{ + let tmp = tempdir().unwrap(); + let home = tmp.path().join("codex-home"); + let sqlite_dir = home.join("sqlite"); + fs::create_dir_all(&sqlite_dir).unwrap(); + let old_db = sqlite_dir.join("state_5.sqlite"); + let new_db = home.join("state_5.sqlite"); + let rollout = home.join("rollout.jsonl"); + let rollout_text = "{\"type\":\"message\",\"payload\":\"original\"}\n"; + fs::write(&rollout, rollout_text).unwrap(); + create_codex_thread_db(&old_db, &rollout); + create_codex_thread_db(&new_db, &rollout); + let db = Connection::open(&new_db).unwrap(); + db.execute("ALTER TABLE threads ADD COLUMN recency_at INTEGER", []) + .unwrap(); + db.execute("UPDATE threads SET recency_at = 42 WHERE id = 't1'", []) + .unwrap(); + drop(db); + + let backups = BackupStore::new(tmp.path().join("backups")); + let deleted = delete_local_from_paths( + vec![old_db.clone(), new_db.clone()], + backups.clone(), + &session("t1", "Codex Thread"), + ); + let token = deleted.undo_token.as_deref().unwrap(); + + assert_eq!(deleted.status, DeleteStatus::LocalDeleted); + assert_eq!(deleted.message, "已从 2 个本地存储删除"); + assert_eq!(thread_count(&old_db, "t1"), 0); + assert_eq!(thread_count(&new_db, "t1"), 0); + assert!(!rollout.exists()); + + let restored = SQLiteStorageAdapter::new(&old_db, backups) + .with_allowed_db_paths(vec![old_db.clone(), new_db.clone()]) + .undo(token); + + assert_eq!(restored.status, DeleteStatus::Undone); + assert_eq!(restored.undo_token.as_deref(), Some(token)); + assert_eq!(thread_count(&old_db, "t1"), 1); + assert_eq!(thread_count(&new_db, "t1"), 1); + assert_eq!(fs::read_to_string(&rollout).unwrap(), rollout_text); + let db = Connection::open(&new_db).unwrap(); + assert_eq!( + db.query_row( + "SELECT recency_at FROM threads WHERE id = 't1'", + [], + |row| row.get::<_, i64>(0), + ) + .unwrap(), + 42 + ); +} + +#[test] +fn grouped_undo_preflights_all_databases_before_restoring_any() { + let tmp = tempdir().unwrap(); + let first_db = tmp.path().join("first.sqlite"); + let second_db = tmp.path().join("second.sqlite"); + let rollout = tmp.path().join("rollout.jsonl"); + fs::write(&rollout, "{\"type\":\"message\"}\n").unwrap(); + create_codex_thread_db(&first_db, &rollout); + create_codex_thread_db(&second_db, &rollout); + let backups = BackupStore::new(tmp.path().join("backups")); + let deleted = delete_local_from_paths( + vec![first_db.clone(), second_db.clone()], + backups.clone(), + &session("t1", "Codex Thread"), + ); + let token = deleted.undo_token.as_deref().unwrap(); + Connection::open(&second_db) + .unwrap() + .execute( + "ALTER TABLE threads RENAME COLUMN title TO renamed_title", + [], + ) + .unwrap(); + + let restored = SQLiteStorageAdapter::new(&first_db, backups) + .with_allowed_db_paths(vec![first_db.clone(), second_db.clone()]) + .undo(token); + + assert_eq!(restored.status, DeleteStatus::Failed); + assert!(restored.message.contains("no column named title")); + assert_eq!(thread_count(&first_db, "t1"), 0); + assert_eq!(thread_count(&second_db, "t1"), 0); + assert!(!rollout.exists()); +} + +#[test] +fn undo_rejects_source_database_outside_allowed_paths() { + let tmp = tempdir().unwrap(); + let allowed_dir = tmp.path().join("home").join("sqlite"); + let outside_dir = tmp.path().join("outside"); + fs::create_dir_all(&allowed_dir).unwrap(); + fs::create_dir_all(&outside_dir).unwrap(); + let allowed_db = allowed_dir.join("codex.sqlite"); + let outside_db = outside_dir.join("codex.sqlite"); + create_supported_db(&allowed_db); + create_supported_db(&outside_db); + let backups = BackupStore::new(tmp.path().join("backups")); + let deleted = SQLiteStorageAdapter::new(&outside_db, backups.clone()) + .delete_local(&session("s1", "First")); + let token = deleted.undo_token.as_deref().unwrap(); + + let restored = SQLiteStorageAdapter::new(&allowed_db, backups).undo(token); + + assert_eq!(restored.status, DeleteStatus::Failed); + assert!( + restored + .message + .contains("not an allowed local storage path") + ); + let outside = Connection::open(&outside_db).unwrap(); + assert_eq!( + outside + .query_row("SELECT COUNT(*) FROM sessions", [], |row| row + .get::<_, i64>(0)) + .unwrap(), + 0 + ); +} + #[test] fn move_thread_workspace_from_paths_uses_database_that_contains_thread() { let tmp = tempdir().unwrap();