fix: record agent episodes reliably

rust/src/cli/session_cmd.rs

@@ -176,6 +176,7 @@ fn default_opts() -> SessionToolOptions<'static> {
         write: false,
         privacy: None,
         terse: None,
+        agent_id: None,
     }
 }
 

rust/src/core/episodic_memory.rs

@@ -5,7 +5,11 @@
 
 use chrono::{DateTime, Utc};
 use serde::{Deserialize, Serialize};
-use std::path::PathBuf;
+use std::{
+    collections::HashMap,
+    path::{Path, PathBuf},
+    sync::{Arc, Mutex, OnceLock},
+};
 
 use crate::core::memory_policy::EpisodicPolicy;
 
@@ -27,6 +31,8 @@ pub struct Episode {
     pub summary: String,
     pub duration_secs: u64,
     pub tokens_used: u64,
+    #[serde(default)]
+    pub agent_id: Option<String>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -57,6 +63,35 @@ impl Outcome {
     }
 }
 
+fn episodic_lock(project_hash: &str) -> Arc<Mutex<()>> {
+    static LOCKS: OnceLock<Mutex<HashMap<String, Arc<Mutex<()>>>>> = OnceLock::new();
+    let mut locks = LOCKS
+        .get_or_init(|| Mutex::new(HashMap::new()))
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    locks.entry(project_hash.to_string()).or_default().clone()
+}
+
+fn acquire_file_lock(path: &Path) -> Option<std::fs::File> {
+    use fs2::FileExt;
+    let parent = path.parent()?;
+    let name = path.file_name()?.to_string_lossy();
+    let lock_path = parent.join(format!(".{name}.lock"));
+    let file = std::fs::OpenOptions::new()
+        .create(true)
+        .truncate(false)
+        .write(true)
+        .open(&lock_path)
+        .ok()?;
+    #[cfg(unix)]
+    {
+        use std::os::unix::fs::PermissionsExt;
+        let _ = std::fs::set_permissions(&lock_path, std::fs::Permissions::from_mode(0o600));
+    }
+    file.lock_exclusive().ok()?;
+    Some(file)
+}
+
 impl EpisodicStore {
     pub fn new(project_hash: &str) -> Self {
         Self {
@@ -241,14 +276,36 @@ impl EpisodicStore {
         Self::load(project_hash).unwrap_or_else(|| Self::new(project_hash))
     }
 
+    pub fn mutate_locked<T>(
+        project_hash: &str,
+        mutate: impl FnOnce(&mut Self) -> T,
+    ) -> Result<(Self, T), String> {
+        let lock = episodic_lock(project_hash);
+        let _guard = lock
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+
+        let path = Self::store_path(project_hash)
+            .ok_or_else(|| "Cannot determine data directory".to_string())?;
+        if let Some(dir) = path.parent() {
+            std::fs::create_dir_all(dir).map_err(|e| format!("{e}"))?;
+        }
+        let _file_lock = acquire_file_lock(&path);
+
+        let mut store = Self::load_or_create(project_hash);
+        let result = mutate(&mut store);
+        store.save()?;
+        Ok((store, result))
+    }
+
     pub fn save(&self) -> Result<(), String> {
         let path = Self::store_path(&self.project_hash)
             .ok_or_else(|| "Cannot determine data directory".to_string())?;
         if let Some(dir) = path.parent() {
             std::fs::create_dir_all(dir).map_err(|e| format!("{e}"))?;
         }
         let json = serde_json::to_string_pretty(self).map_err(|e| format!("{e}"))?;
-        std::fs::write(path, json).map_err(|e| format!("{e}"))
+        crate::core::atomic_fs::write_bytes_with_fallback(&path, json.as_bytes(), None)
     }
 }
 
@@ -324,9 +381,46 @@ pub fn create_episode_from_session(
         // Cumulative session counter at record time; the caller converts
         // this into a per-task delta (see `finalize_episode_metrics`).
         tokens_used: session.stats.total_tokens_saved,
+        agent_id: None,
     }
 }
 
+pub fn record_session_episode(
+    project_hash: &str,
+    session: &super::session::SessionState,
+    tool_calls: &[(String, u64)],
+    agent_id: Option<&str>,
+    policy: &EpisodicPolicy,
+    deduplicate: bool,
+) -> Result<Option<String>, String> {
+    let normalized_agent_id = agent_id
+        .map(str::trim)
+        .filter(|id| !id.is_empty())
+        .map(str::to_string);
+
+    let (_, episode_id) = EpisodicStore::mutate_locked(project_hash, |store| {
+        let mut episode = create_episode_from_session(session, tool_calls);
+        episode.agent_id.clone_from(&normalized_agent_id);
+
+        if deduplicate
+            && store.episodes.iter().any(|existing| {
+                existing.session_id == episode.session_id
+                    && existing.agent_id == episode.agent_id
+                    && existing.task_description == episode.task_description
+            })
+        {
+            return None;
+        }
+
+        finalize_episode_metrics(&mut episode, store, session.started_at);
+        let id = episode.id.clone();
+        store.record_episode(episode, policy);
+        Some(id)
+    })?;
+
+    Ok(episode_id)
+}
+
 /// Converts the cumulative session counters captured by
 /// [`create_episode_from_session`] into per-task values.
 ///
@@ -449,9 +543,97 @@ mod tests {
             summary: String::new(),
             duration_secs: 60,
             tokens_used: 5000,
+            agent_id: None,
         }
     }
 
+    #[test]
+    fn mutate_locked_preserves_successive_agent_episodes() {
+        let tmp = tempfile::tempdir().unwrap();
+        crate::test_env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
+        let policy = EpisodicPolicy::default();
+        let project_hash = "episodic-locked-writes";
+
+        EpisodicStore::mutate_locked(project_hash, |store| {
+            let mut episode = make_episode("Task from agent A", Outcome::Unknown);
+            episode.agent_id = Some("agent-a".to_string());
+            store.record_episode(episode, &policy);
+        })
+        .unwrap();
+        EpisodicStore::mutate_locked(project_hash, |store| {
+            let mut episode = make_episode("Task from agent B", Outcome::Unknown);
+            episode.agent_id = Some("agent-b".to_string());
+            store.record_episode(episode, &policy);
+        })
+        .unwrap();
+
+        let store = EpisodicStore::load_or_create(project_hash);
+        assert_eq!(store.episodes.len(), 2);
+        assert!(
+            store
+                .episodes
+                .iter()
+                .any(|episode| episode.agent_id.as_deref() == Some("agent-a"))
+        );
+        assert!(
+            store
+                .episodes
+                .iter()
+                .any(|episode| episode.agent_id.as_deref() == Some("agent-b"))
+        );
+    }
+
+    #[test]
+    fn record_session_episode_deduplicates_per_agent() {
+        let tmp = tempfile::tempdir().unwrap();
+        crate::test_env::set_var("LEAN_CTX_DATA_DIR", tmp.path());
+        let policy = EpisodicPolicy::default();
+        let project_hash = "episodic-agent-dedup";
+        let mut session = super::super::session::SessionState::new();
+        session.set_task("same task", None);
+        let tool_calls = vec![("ctx_read".to_string(), 10)];
+
+        assert!(
+            record_session_episode(
+                project_hash,
+                &session,
+                &tool_calls,
+                Some("agent-a"),
+                &policy,
+                true,
+            )
+            .unwrap()
+            .is_some()
+        );
+        assert!(
+            record_session_episode(
+                project_hash,
+                &session,
+                &tool_calls,
+                Some("agent-a"),
+                &policy,
+                true,
+            )
+            .unwrap()
+            .is_none()
+        );
+        assert!(
+            record_session_episode(
+                project_hash,
+                &session,
+                &tool_calls,
+                Some("agent-b"),
+                &policy,
+                true,
+            )
+            .unwrap()
+            .is_some()
+        );
+
+        let store = EpisodicStore::load_or_create(project_hash);
+        assert_eq!(store.episodes.len(), 2);
+    }
+
     #[test]
     fn record_and_search() {
         let policy = EpisodicPolicy::default();

rust/src/core/procedural_memory.rs

@@ -363,6 +363,7 @@ mod tests {
             summary: String::new(),
             duration_secs: 60,
             tokens_used: 1000,
+            agent_id: None,
         }
     }
 

rust/src/tools/ctx_session.rs

@@ -8,6 +8,7 @@ pub struct SessionToolOptions<'a> {
     pub privacy: Option<&'a str>,
     /// For `action=configure`: set terse output mode when `Some`.
     pub terse: Option<bool>,
+    pub agent_id: Option<&'a str>,
 }
 
 pub fn handle(
@@ -235,7 +236,7 @@ stale_files: {}\n",
                 desc.contains("[100%]") || lower.contains("[done]") || lower.contains("[complete]");
             let mut note = String::new();
             if completed {
-                match auto_record_episode(session, tool_calls) {
+                match auto_record_episode(session, tool_calls, opts.agent_id) {
                     Ok(Some(id)) => {
                         note = format!("\nEpisode auto-recorded: {id}");
                     }
@@ -562,18 +563,20 @@ stale_files: {}\n",
                 }
             };
             let hash = crate::core::project_hash::hash_project_root(&project_root);
-            let mut store = crate::core::episodic_memory::EpisodicStore::load_or_create(&hash);
-
             match value {
                 Some("record") => {
-                    let ep = crate::core::episodic_memory::create_episode_from_session(
-                        session, tool_calls,
-                    );
-                    let id = ep.id.clone();
-                    store.record_episode(ep, &policy.episodic);
-                    if let Err(e) = store.save() {
-                        return format!("Episode record failed: {e}");
-                    }
+                    let id = match crate::core::episodic_memory::record_session_episode(
+                        &hash,
+                        session,
+                        tool_calls,
+                        opts.agent_id,
+                        &policy.episodic,
+                        false,
+                    ) {
+                        Ok(Some(id)) => id,
+                        Ok(None) => return "Episode already recorded.".to_string(),
+                        Err(e) => return format!("Episode record failed: {e}"),
+                    };
                     crate::core::events::emit(crate::core::events::EventKind::KnowledgeUpdate {
                         category: "episodic".to_string(),
                         key: id.clone(),
@@ -596,6 +599,7 @@ stale_files: {}\n",
                 }
                 Some(v) if v.starts_with("search ") => {
                     let q = v.trim_start_matches("search ").trim();
+                    let store = crate::core::episodic_memory::EpisodicStore::load_or_create(&hash);
                     let hits = store.search(q);
                     if hits.is_empty() {
                         return "No episodes matched.".to_string();
@@ -615,6 +619,7 @@ stale_files: {}\n",
                 }
                 Some(v) if v.starts_with("file ") => {
                     let f = v.trim_start_matches("file ").trim();
+                    let store = crate::core::episodic_memory::EpisodicStore::load_or_create(&hash);
                     let hits = store.by_file(f);
                     let mut out = format!("Episodes for file match '{f}' ({}):", hits.len());
                     for ep in hits.into_iter().take(10) {
@@ -631,6 +636,7 @@ stale_files: {}\n",
                 }
                 Some(v) if v.starts_with("outcome ") => {
                     let label = v.trim_start_matches("outcome ").trim();
+                    let store = crate::core::episodic_memory::EpisodicStore::load_or_create(&hash);
                     let hits = store.by_outcome(label);
                     let mut out = format!("Episodes outcome '{label}' ({}):", hits.len());
                     for ep in hits.into_iter().take(10) {
@@ -640,6 +646,7 @@ stale_files: {}\n",
                     out
                 }
                 _ => {
+                    let store = crate::core::episodic_memory::EpisodicStore::load_or_create(&hash);
                     let stats = store.stats();
                     let recent = store.recent(10);
                     let mut out = format!(
@@ -776,6 +783,7 @@ stale_files: {}\n",
 fn auto_record_episode(
     session: &SessionState,
     tool_calls: &[(String, u64)],
+    agent_id: Option<&str>,
 ) -> Result<Option<String>, String> {
     let project_root = session.project_root.clone().unwrap_or_else(|| {
         std::env::current_dir().map_or_else(
@@ -787,20 +795,17 @@ fn auto_record_episode(
         .memory_policy_effective()
         .map_err(|e| format!("invalid memory policy: {e}"))?;
     let hash = crate::core::project_hash::hash_project_root(&project_root);
-    let mut store = crate::core::episodic_memory::EpisodicStore::load_or_create(&hash);
-
-    let mut ep = crate::core::episodic_memory::create_episode_from_session(session, tool_calls);
-    if let Some(last) = store.recent(1).first()
-        && last.task_description == ep.task_description
-    {
+    let Some(id) = crate::core::episodic_memory::record_session_episode(
+        &hash,
+        session,
+        tool_calls,
+        agent_id,
+        &policy.episodic,
+        true,
+    )?
+    else {
         return Ok(None);
-    }
-    // Convert cumulative session counters into per-task delta + duration.
-    crate::core::episodic_memory::finalize_episode_metrics(&mut ep, &store, session.started_at);
-
-    let id = ep.id.clone();
-    store.record_episode(ep, &policy.episodic);
-    store.save()?;
+    };
     crate::core::events::emit(crate::core::events::EventKind::KnowledgeUpdate {
         category: "episodic".to_string(),
         key: id.clone(),
@@ -810,6 +815,7 @@ fn auto_record_episode(
     // Each new episode is a chance to learn a procedure: mine the episode
     // history for repeated tool sequences. Best-effort — pattern detection
     // must never fail the task update itself (#478).
+    let store = crate::core::episodic_memory::EpisodicStore::load_or_create(&hash);
     let episodes: Vec<crate::core::episodic_memory::Episode> =
         store.recent(50).into_iter().cloned().collect();
     let mut procs = crate::core::procedural_memory::ProceduralStore::load_or_create(&hash);