diff --git a/rust/src/cli/session_cmd.rs b/rust/src/cli/session_cmd.rs index 88979fc841..1704877d72 100644 --- a/rust/src/cli/session_cmd.rs +++ b/rust/src/cli/session_cmd.rs @@ -176,6 +176,7 @@ fn default_opts() -> SessionToolOptions<'static> { write: false, privacy: None, terse: None, + agent_id: None, } } diff --git a/rust/src/core/episodic_memory.rs b/rust/src/core/episodic_memory.rs index 0daba415b3..c4cd50222e 100644 --- a/rust/src/core/episodic_memory.rs +++ b/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, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -57,6 +63,35 @@ impl Outcome { } } +fn episodic_lock(project_hash: &str) -> Arc> { + static LOCKS: OnceLock>>>> = 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 { + 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,6 +276,28 @@ impl EpisodicStore { Self::load(project_hash).unwrap_or_else(|| Self::new(project_hash)) } + pub fn mutate_locked( + 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())?; @@ -248,7 +305,7 @@ impl EpisodicStore { 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, 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(); diff --git a/rust/src/core/procedural_memory.rs b/rust/src/core/procedural_memory.rs index 275d592a94..8d82be7bf6 100644 --- a/rust/src/core/procedural_memory.rs +++ b/rust/src/core/procedural_memory.rs @@ -363,6 +363,7 @@ mod tests { summary: String::new(), duration_secs: 60, tokens_used: 1000, + agent_id: None, } } diff --git a/rust/src/tools/ctx_session.rs b/rust/src/tools/ctx_session.rs index c34392c124..7e7de69018 100644 --- a/rust/src/tools/ctx_session.rs +++ b/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, + 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, 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 = store.recent(50).into_iter().cloned().collect(); let mut procs = crate::core::procedural_memory::ProceduralStore::load_or_create(&hash); diff --git a/rust/src/tools/registered/ctx_session.rs b/rust/src/tools/registered/ctx_session.rs index b6ee91a7a4..da9e1ae9d2 100644 --- a/rust/src/tools/registered/ctx_session.rs +++ b/rust/src/tools/registered/ctx_session.rs @@ -56,6 +56,10 @@ impl McpTool for CtxSessionTool { let tc = tool_calls_handle.blocking_read(); tc.iter().map(|c| (c.tool.clone(), c.duration_ms)).collect() }; + let agent_id = ctx + .agent_id + .as_ref() + .and_then(|agent_id| agent_id.blocking_read().clone()); let session_handle = ctx .session @@ -74,6 +78,7 @@ impl McpTool for CtxSessionTool { write, privacy: privacy.as_deref(), terse, + agent_id: agent_id.as_deref(), }, ); diff --git a/rust/src/tools/registered/ctx_transcript_compact.rs b/rust/src/tools/registered/ctx_transcript_compact.rs index 1e19369006..9d4fff9902 100644 --- a/rust/src/tools/registered/ctx_transcript_compact.rs +++ b/rust/src/tools/registered/ctx_transcript_compact.rs @@ -90,6 +90,7 @@ impl McpTool for CtxTranscriptCompactTool { write: false, privacy: None, terse: Some(true), + agent_id: None, }, ); } diff --git a/rust/src/tools/server_lifecycle.rs b/rust/src/tools/server_lifecycle.rs index 07598357c0..fe9204dca3 100644 --- a/rust/src/tools/server_lifecycle.rs +++ b/rust/src/tools/server_lifecycle.rs @@ -242,8 +242,53 @@ impl LeanCtxServer { *self.last_call.write().await = Instant::now(); } + async fn record_shutdown_episode(&self) { + let tool_calls: Vec<(String, u64)> = self + .tool_calls + .read() + .await + .iter() + .map(|call| (call.tool.clone(), call.duration_ms)) + .collect(); + if tool_calls.is_empty() { + return; + } + + let session = self.session.read().await.clone(); + let Some(project_root) = session + .project_root + .clone() + .or_else(|| self.startup_project_root.clone()) + else { + return; + }; + let agent_id = self.agent_id.read().await.clone(); + let agent_id = match agent_id { + Some(agent_id) => Some(agent_id), + None => self.presence_agent_id.read().await.clone(), + }; + let policy = crate::core::config::Config::load() + .memory_policy_effective() + .unwrap_or_default(); + let project_hash = crate::core::project_hash::hash_project_root(&project_root); + + match crate::core::episodic_memory::record_session_episode( + &project_hash, + &session, + &tool_calls, + agent_id.as_deref(), + &policy.episodic, + true, + ) { + Ok(Some(id)) => tracing::info!("lean-ctx: recorded shutdown episode {id}"), + Ok(None) => {} + Err(error) => tracing::warn!("lean-ctx: failed to record shutdown episode: {error}"), + } + } + /// Aggressive cleanup on connection drop: save session, consolidate knowledge, clear caches. pub async fn shutdown(&self) { + self.record_shutdown_episode().await; if let Some(agent_id) = self.presence_agent_id.read().await.clone() && let Err(error) = crate::core::agents::AgentRegistry::finish_persistent(&agent_id) { diff --git a/rust/tests/suite/compaction_survival_tests.rs b/rust/tests/suite/compaction_survival_tests.rs index 3f3860475d..81d174de67 100644 --- a/rust/tests/suite/compaction_survival_tests.rs +++ b/rust/tests/suite/compaction_survival_tests.rs @@ -157,6 +157,7 @@ fn session_resume_action() { write: false, privacy: None, terse: None, + agent_id: None, }, ); assert!(