diff --git a/src/app.rs b/src/app.rs index 4a2fab5..fe502cd 100644 --- a/src/app.rs +++ b/src/app.rs @@ -1757,11 +1757,19 @@ impl App { .collect::>() }) .unwrap_or_default(); + for session_id in &checkpoint_ids { + if let Err(error) = + discard_session_checkpoint_files(&kv_cache_path(), *session_id) + { + self.error = Some(error); + return Task::none(); + } + } + self.finish_cache_change(); if let Some(database) = &mut self.database { match database.delete_project(project_id) { Ok(()) => { for session_id in checkpoint_ids { - let _ = fs::remove_file(session_checkpoint_path(session_id)); #[cfg(target_os = "macos")] self.background_chats.remove(&session_id); } @@ -1902,22 +1910,33 @@ impl App { self.error = Some("Stop the active generation before rebuilding context.".into()); } else { - match fs::remove_file(session_checkpoint_path(session_id)) { - Ok(()) => { + match discard_session_checkpoint_files(&kv_cache_path(), session_id) { + Ok(_) => { self.error = None; self.finish_cache_change(); } - Err(error) if error.kind() == std::io::ErrorKind::NotFound => { - self.error = None; - } Err(error) => { - self.error = Some(format!("Could not discard the checkpoint: {error}")); + self.error = Some(error); } } } } Message::SetSessionState(session_id, state) => { self.session_menu = None; + if state == SessionState::Archived { + if self.session_is_active(session_id) { + self.error = + Some("Stop the active generation before archiving its session.".into()); + return Task::none(); + } + if let Err(error) = + discard_session_checkpoint_files(&kv_cache_path(), session_id) + { + self.error = Some(error); + return Task::none(); + } + self.finish_cache_change(); + } if let Some(database) = &mut self.database { match database.set_session_state(session_id, state) { Ok(()) => { @@ -2047,13 +2066,16 @@ impl App { Some("Stop the active generation before deleting its session.".into()); return Task::none(); } + if let Err(error) = discard_session_checkpoint_files(&kv_cache_path(), session_id) { + self.error = Some(error); + return Task::none(); + } + self.finish_cache_change(); if let Some(database) = &mut self.database { match database.delete_session(session_id) { Ok(()) => { - let _ = fs::remove_file(session_checkpoint_path(session_id)); #[cfg(target_os = "macos")] self.background_chats.remove(&session_id); - self.finish_cache_change(); if self.session_menu == Some(session_id) { self.session_menu = None; } @@ -2723,6 +2745,24 @@ fn session_checkpoint_path(session_id: i32) -> PathBuf { kv_cache_path().join(format!("{session_id}.bin")) } +fn discard_session_checkpoint_files(directory: &Path, session_id: i32) -> Result { + let mut removed = false; + for extension in ["bin", "tmp", "compacting"] { + let path = directory.join(format!("{session_id}.{extension}")); + match fs::remove_file(&path) { + Ok(()) => removed = true, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(format!( + "Could not discard session checkpoint {}: {error}", + path.display() + )); + } + } + } + Ok(removed) +} + fn session_compaction_checkpoint_path(session_id: i32) -> PathBuf { kv_cache_path().join(format!("{session_id}.compacting")) } diff --git a/src/app/generation.rs b/src/app/generation.rs index 11ca25f..5098112 100644 --- a/src/app/generation.rs +++ b/src/app/generation.rs @@ -613,7 +613,13 @@ impl App { } } }; - let Some(service) = &self.generation_service else { + let archived_session = self + .projects + .iter() + .flat_map(|project| &project.sessions) + .find(|session| session.id == session_id) + .is_some_and(|session| session.state() == SessionState::Archived); + let Some(service) = self.generation_service.clone() else { self.error = Some("The model runtime is unavailable.".into()); return; }; @@ -633,6 +639,13 @@ impl App { return; } }; + if archived_session { + self.reload_projects(); + self.context_notice = Some( + "Rebuilding context: the session was archived and its checkpoint was discarded." + .into(), + ); + } let user_id = saved[saved.len() - 2].id; self.active_turn .get_or_insert_with(TurnSummary::new) @@ -1643,6 +1656,12 @@ impl App { } else { None }; + let archived_session = self + .projects + .iter() + .flat_map(|project| &project.sessions) + .find(|session| session.id == session_id) + .is_some_and(|session| session.state() == SessionState::Archived); let tail_start = message_ids.get(compacted.tail_start).copied(); let messages = self .database @@ -1659,7 +1678,9 @@ impl App { .map_err(|error| format!("Could not save compacted conversation: {error}"))?; self.conversation .extend(messages.into_iter().map(ChatMessage::from)); - if let Some(session) = self + if archived_session { + self.reload_projects(); + } else if let Some(session) = self .projects .iter_mut() .flat_map(|project| &mut project.sessions) diff --git a/src/app/projects.rs b/src/app/projects.rs index e3ed5dd..a601b72 100644 --- a/src/app/projects.rs +++ b/src/app/projects.rs @@ -515,6 +515,25 @@ mod tests { std::fs::remove_dir_all(directory).unwrap(); } + #[test] + fn discarding_a_session_removes_every_checkpoint_stage() { + let directory = + std::env::temp_dir().join(format!("ds4-server-discard-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + for name in ["4.bin", "4.tmp", "4.compacting", "5.bin"] { + std::fs::write(directory.join(name), b"payload").unwrap(); + } + + assert!(super::super::discard_session_checkpoint_files(&directory, 4).unwrap()); + assert!(!directory.join("4.bin").exists()); + assert!(!directory.join("4.tmp").exists()); + assert!(!directory.join("4.compacting").exists()); + assert!(directory.join("5.bin").exists()); + assert!(!super::super::discard_session_checkpoint_files(&directory, 4).unwrap()); + + std::fs::remove_dir_all(directory).unwrap(); + } + #[test] fn git_state_lists_and_switches_local_branches() { let nonce = SystemTime::now() diff --git a/src/database.rs b/src/database.rs index 4483e75..d8503fa 100644 --- a/src/database.rs +++ b/src/database.rs @@ -448,6 +448,7 @@ impl Database { ) -> Result, String> { self.connection .transaction(|connection| { + reactivate_archived_session(connection, session_id)?; touch_session(connection, session_id)?; let mut stored = Vec::with_capacity(system_messages.len() + 2); for content in system_messages { @@ -654,6 +655,7 @@ impl Database { .map_err(|_| "Context limit is too large to save".to_owned())?; self.connection .transaction(|connection| { + reactivate_archived_session(connection, session_id)?; diesel::update(sessions::table.find(session_id)) .set(( sessions::compacted_summary.eq(Some(summary)), @@ -715,6 +717,20 @@ fn touch_session( .map(|_| ()) } +fn reactivate_archived_session( + connection: &mut SqliteConnection, + session_id: i32, +) -> Result<(), diesel::result::Error> { + diesel::update( + sessions::table + .find(session_id) + .filter(sessions::state.eq(SessionState::Archived.as_id())), + ) + .set(sessions::state.eq(SessionState::Normal.as_id())) + .execute(connection) + .map(|_| ()) +} + #[cfg(test)] mod tests { use super::*; @@ -827,6 +843,44 @@ mod tests { fs::remove_file(path).unwrap(); } + #[test] + fn context_work_reactivates_an_archived_session() { + let id = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + let path = std::env::temp_dir().join(format!("ds4-reactivate-{id}.sqlite3")); + let mut database = Database::open(&path).unwrap(); + let project = database + .create_project("DS4", "/tmp/ds4-reactivate") + .unwrap(); + let session = database.create_session(project.id, "Archived").unwrap(); + database + .set_session_state(session.id, SessionState::Archived) + .unwrap(); + + database + .start_chat_turn(session.id, "Resume", None, &[], false) + .unwrap(); + + assert_eq!( + database.load_projects().unwrap()[0].sessions[0].state(), + SessionState::Normal + ); + database + .set_session_state(session.id, SessionState::Archived) + .unwrap(); + database + .record_compaction(session.id, "Summary", None, None, 100, 1_000) + .unwrap(); + assert_eq!( + database.load_projects().unwrap()[0].sessions[0].state(), + SessionState::Normal + ); + drop(database); + fs::remove_file(path).unwrap(); + } + #[test] fn a2ui_dismissal_persists_a_fresh_surface_boundary() { let id = SystemTime::now() diff --git a/src/engine.rs b/src/engine.rs index d19187c..4cff18e 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -962,6 +962,15 @@ impl Generator { matches: impl Fn([u8; 32]) -> bool, ) -> Result { let resident_hit = self.activate_resident(checkpoint.to_owned())?; + if !checkpoint.is_file() { + self.executor.reset()?; + self.checkpoint = None; + self.metrics.kv_lookup(KvLookup::Miss); + return Ok(CheckpointSelection { + found: false, + incompatible: false, + }); + } if resident_hit && matches(self.executor.checkpoint_tag()) { self.checkpoint = Some(checkpoint.to_owned()); self.metrics.kv_lookup(KvLookup::MemoryHit); @@ -971,15 +980,6 @@ impl Generator { }); } if self.checkpoint.as_deref() == Some(checkpoint) { - if !checkpoint.is_file() { - self.executor.reset()?; - self.checkpoint = None; - self.metrics.kv_lookup(KvLookup::Miss); - return Ok(CheckpointSelection { - found: false, - incompatible: false, - }); - } let found = matches(self.executor.checkpoint_tag()); self.metrics.kv_lookup(if found { KvLookup::MemoryHit