fix: Preserve legacy embedding on rollback

Author: nguyenthdat

src/daemon/mod.rs

@@ -1193,6 +1193,23 @@ fn config_from_acquire(request: &AcquireProjectRequest) -> Result<MemoryConfig>
                 crate::model::embedding_config_for_profile(&active.profile_id, config.embedding())?;
             config.set_embedding(persisted);
         } else {
+            let store = crate::embedding_generation::ModelSwitchStore::load(
+                &config.model_switch_path(),
+                config.project_id(),
+            )?;
+            if let Some(embedding) = store
+                .current
+                .iter()
+                .chain(store.history.iter().rev())
+                .find(|job| {
+                    job.phase == crate::embedding_generation::SwitchPhase::Succeeded
+                        && job.target_generation_id == active.generation_id
+                        && job.target_profile_id == active.profile_id
+                })
+                .and_then(|job| job.target_embedding.clone())
+            {
+                config.set_embedding(embedding);
+            }
             anyhow::ensure!(
                 active.profile_fingerprint == config.embedding_profile_fingerprint()?,
                 "requested custom embedding does not match the persisted active profile"

src/embedding_generation.rs

@@ -280,6 +280,8 @@ pub(crate) struct ModelSwitchJournal {
     pub source_profile_id: String,
     #[serde(default)]
     pub source_embedding: Option<EmbeddingConfig>,
+    #[serde(default)]
+    pub target_embedding: Option<EmbeddingConfig>,
     pub target_generation_id: String,
     pub target_profile_id: String,
     pub phase: SwitchPhase,
@@ -442,6 +444,7 @@ mod tests {
             source_generation_id: "legacy".to_string(),
             source_profile_id: "profile-a".to_string(),
             source_embedding: Some(EmbeddingConfig::default()),
+            target_embedding: None,
             target_generation_id: format!("gen_{id}"),
             target_profile_id: "profile-b".to_string(),
             phase: SwitchPhase::Succeeded,

src/engine/mod.rs

@@ -101,14 +101,32 @@ impl MemoryEngine {
         let state = MemoryState::load(&config.state_path())?;
         let document_index = DocumentIndexManifest::load(&config.document_index_path())?;
         let graph = GraphStore::load(&config.graph_state_path(), &config.graph_pending_path())?;
+        let switch_store =
+            ModelSwitchStore::load(&config.model_switch_path(), config.project_id())?;
         let persisted_active =
             ActiveEmbedding::load(&config.active_embedding_path(), config.project_id())?;
-        if let Some(active) = &persisted_active
-            && active.profile_id != "legacy-custom"
-        {
-            let embedding =
-                crate::model::embedding_config_for_profile(&active.profile_id, config.embedding())?;
-            config.set_embedding(embedding);
+        if let Some(active) = &persisted_active {
+            if active.profile_id == "legacy-custom" {
+                if let Some(embedding) = switch_store
+                    .current
+                    .iter()
+                    .chain(switch_store.history.iter().rev())
+                    .find(|job| {
+                        job.phase == crate::embedding_generation::SwitchPhase::Succeeded
+                            && job.target_generation_id == active.generation_id
+                            && job.target_profile_id == active.profile_id
+                    })
+                    .and_then(|job| job.target_embedding.clone())
+                {
+                    config.set_embedding(embedding);
+                }
+            } else {
+                let embedding = crate::model::embedding_config_for_profile(
+                    &active.profile_id,
+                    config.embedding(),
+                )?;
+                config.set_embedding(embedding);
+            }
         }
         let embedder = LlamaCppEmbedder::load(config.embedding(), config.model_cache())?;
         let active_embedding = match persisted_active {
@@ -157,9 +175,6 @@ impl MemoryEngine {
                 now_ms()?,
             )?
         };
-        let switch_store =
-            ModelSwitchStore::load(&config.model_switch_path(), config.project_id())?;
-
         let mut engine = Self {
             config,
             collection,

src/engine/model_switch.rs

@@ -100,6 +100,12 @@ impl MemoryEngine {
         });
         validate_generation_id(&target_generation_id)?;
         let direct_rollback = request.target_generation_id.is_some();
+        let target_embedding = if direct_rollback && target_generation_id == "legacy" {
+            self.legacy_snapshot(&request.target_profile_id)
+                .and_then(|snapshot| snapshot.source_embedding.clone())
+        } else {
+            None
+        };
         if direct_rollback {
             if target_generation_id == "legacy" {
                 let snapshot = self
@@ -137,6 +143,7 @@ impl MemoryEngine {
             source_generation_id: self.active_embedding.generation_id.clone(),
             source_profile_id: self.active_embedding.profile_id.clone(),
             source_embedding: Some(self.config.embedding().clone()),
+            target_embedding,
             target_generation_id,
             target_profile_id: request.target_profile_id.clone(),
             phase: if direct_rollback {
@@ -656,8 +663,12 @@ impl MemoryEngine {
         }
         if self.switch_target_config.is_none() {
             self.switch_target_config = Some(if job.target_generation_id == "legacy" {
-                self.legacy_snapshot(&job.target_profile_id)
-                    .and_then(|snapshot| snapshot.source_embedding.clone())
+                job.target_embedding
+                    .clone()
+                    .or_else(|| {
+                        self.legacy_snapshot(&job.target_profile_id)
+                            .and_then(|snapshot| snapshot.source_embedding.clone())
+                    })
                     .ok_or_else(|| {
                         anyhow!("retained legacy embedding configuration is unavailable")
                     })?

src/storage/atomic.rs

@@ -37,16 +37,16 @@ pub(crate) fn write_json_atomic<T: Serialize>(
 pub(crate) fn remove_file_durable(path: &Path) -> Result<()> {
     if path.exists() {
         fs::remove_file(path).with_context(|| format!("cannot remove {}", path.display()))?;
-        sync_parent(path)?;
     }
+    sync_parent(path)?;
     Ok(())
 }
 
 pub(crate) fn remove_dir_all_durable(path: &Path) -> Result<()> {
     if path.exists() {
         fs::remove_dir_all(path).with_context(|| format!("cannot remove {}", path.display()))?;
-        sync_parent(path)?;
     }
+    sync_parent(path)?;
     Ok(())
 }