Skip to main content

nexus_memory_hooks/
injection.rs

1//! Reference injection system for agent configuration files.
2
3use nexus_core::fsutil::atomic_write;
4use std::fs;
5use std::io::{self, Write};
6use std::path::{Path, PathBuf};
7use tracing::{debug, info};
8
9/// Target for injecting Nexus references into an agent's configuration.
10#[derive(Debug, Clone)]
11pub struct AgentInjectionTarget {
12    pub agent_type: String,
13    pub global_config: Option<PathBuf>,
14    pub project_config_filename: String,
15}
16
17/// Sentinel markers for Nexus-managed blocks.
18pub const NEXUS_BLOCK_START: &str = "<!-- NEXUS:START -->";
19pub const NEXUS_BLOCK_END: &str = "<!-- NEXUS:END -->";
20
21impl AgentInjectionTarget {
22    /// Get known agent injection targets.
23    pub fn known_agents() -> Vec<Self> {
24        let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
25        vec![
26            Self {
27                agent_type: "claude-code".to_string(),
28                global_config: Some(home.join(".claude").join("CLAUDE.md")),
29                project_config_filename: "CLAUDE.md".to_string(),
30            },
31            Self {
32                agent_type: "amp".to_string(),
33                global_config: Some(home.join(".config").join("amp").join("AGENTS.md")),
34                project_config_filename: "AGENTS.md".to_string(),
35            },
36            Self {
37                agent_type: "codex".to_string(),
38                global_config: Some(home.join(".config").join("codex").join("AGENTS.md")),
39                project_config_filename: "AGENTS.md".to_string(),
40            },
41            Self {
42                agent_type: "gemini".to_string(),
43                global_config: Some(home.join(".gemini").join("GEMINI.md")),
44                project_config_filename: "GEMINI.md".to_string(),
45            },
46            Self {
47                agent_type: "pi-mono".to_string(),
48                global_config: Some(home.join(".pi").join("agent").join("AGENTS.md")),
49                project_config_filename: ".pi/AGENTS.md".to_string(),
50            },
51        ]
52    }
53
54    /// Find injection target by agent type.
55    pub fn find(agent_type: &str) -> Option<Self> {
56        Self::known_agents()
57            .into_iter()
58            .find(|t| t.agent_type == agent_type || agent_type.contains(&t.agent_type))
59    }
60}
61
62/// Inject Nexus references into a config file.
63pub fn inject_reference(
64    config_file: &Path,
65    soul_path: &Path,
66    context_path: &Path,
67) -> io::Result<()> {
68    if !config_file.exists() {
69        return Ok(());
70    }
71
72    let content = fs::read_to_string(config_file)?;
73    let original_content = content.clone();
74
75    // Build the Nexus block
76    let block = format!(
77        "{}\n\
78        ## Nexus Memory Substrate\n\
79        - Identity: [{soul_name}]({soul_path})\n\
80        - Project Context: [{context_name}]({context_path})\n\
81        {}",
82        NEXUS_BLOCK_START,
83        NEXUS_BLOCK_END,
84        soul_name = "Soul",
85        soul_path = soul_path.to_string_lossy(),
86        context_name = "Project Context",
87        context_path = context_path.to_string_lossy(),
88    );
89
90    let new_content = if let (Some(start), Some(end)) = (
91        content.find(NEXUS_BLOCK_START),
92        content.find(NEXUS_BLOCK_END),
93    ) {
94        if start >= end {
95            // Malformed markers (end before start) — strip both markers from
96            // content, then append a fresh block. Simple replace avoids the
97            // slicing bugs that kept stale markers around.
98            let stripped = content
99                .replace(NEXUS_BLOCK_START, "")
100                .replace(NEXUS_BLOCK_END, "");
101            let mut updated = stripped.trim_end().to_string();
102            updated.push('\n');
103            updated.push_str(&block);
104            if !updated.ends_with('\n') {
105                updated.push('\n');
106            }
107            updated
108        } else {
109            // Replace existing block
110            let mut updated = content[..start].to_string();
111            updated.push_str(&block);
112            updated.push_str(&content[end + NEXUS_BLOCK_END.len()..]);
113            updated
114        }
115    } else {
116        // Append to end
117        let mut updated = content;
118        if !updated.is_empty() && !updated.ends_with('\n') {
119            updated.push('\n');
120        }
121        updated.push_str(&block);
122        if !updated.ends_with('\n') {
123            updated.push('\n');
124        }
125        updated
126    };
127
128    if new_content != original_content {
129        atomic_write(config_file, &new_content)?;
130        debug!("Injected Nexus reference into {}", config_file.display());
131    }
132
133    Ok(())
134}
135
136/// Inject only the soul identity reference into a config file (no project context).
137/// Used for global config files that should not reference project-specific context.md.
138pub fn inject_soul_only(config_file: &Path, soul_path: &Path) -> io::Result<()> {
139    if !config_file.exists() {
140        return Ok(());
141    }
142
143    let content = fs::read_to_string(config_file)?;
144    let original_content = content.clone();
145
146    let block = format!(
147        "{}\n\
148        ## Nexus Memory Substrate\n\
149        - Identity: [Soul]({soul_path_val})\n\
150        {}",
151        NEXUS_BLOCK_START,
152        NEXUS_BLOCK_END,
153        soul_path_val = soul_path.to_string_lossy(),
154    );
155    let new_content = if let (Some(start), Some(end)) = (
156        content.find(NEXUS_BLOCK_START),
157        content.find(NEXUS_BLOCK_END),
158    ) {
159        if start >= end {
160            // Malformed markers (end before start) — strip both markers from
161            // content, then append a fresh block. Simple replace avoids the
162            // slicing bugs that kept stale markers around.
163            let stripped = content
164                .replace(NEXUS_BLOCK_START, "")
165                .replace(NEXUS_BLOCK_END, "");
166            let mut updated = stripped.trim_end().to_string();
167            updated.push('\n');
168            updated.push_str(&block);
169            if !updated.ends_with('\n') {
170                updated.push('\n');
171            }
172            updated
173        } else {
174            let mut updated = content[..start].to_string();
175            updated.push_str(&block);
176            updated.push_str(&content[end + NEXUS_BLOCK_END.len()..]);
177            updated
178        }
179    } else {
180        let mut updated = content;
181        if !updated.is_empty() && !updated.ends_with('\n') {
182            updated.push('\n');
183        }
184        updated.push_str(&block);
185        if !updated.ends_with('\n') {
186            updated.push('\n');
187        }
188        updated
189    };
190
191    if new_content != original_content {
192        atomic_write(config_file, &new_content)?;
193        debug!(
194            "Injected soul-only Nexus reference into {}",
195            config_file.display()
196        );
197    }
198
199    Ok(())
200}
201/// Remove Nexus references from a config file.
202pub fn remove_reference(config_file: &Path) -> io::Result<()> {
203    if !config_file.exists() {
204        return Ok(());
205    }
206
207    let content = fs::read_to_string(config_file)?;
208    if let (Some(start), Some(end)) = (
209        content.find(NEXUS_BLOCK_START),
210        content.find(NEXUS_BLOCK_END),
211    ) {
212        let mut updated = content[..start].to_string();
213        let remaining = &content[end + NEXUS_BLOCK_END.len()..];
214        updated.push_str(remaining);
215
216        // Final cleanup: if we're left with multiple newlines at the end, collapse them
217        while updated.ends_with("\n\n") {
218            updated.pop();
219        }
220
221        atomic_write(config_file, &updated)?;
222    }
223
224    Ok(())
225}
226
227/// Pipeline executed when a new agent session starts.
228pub async fn on_session_start(
229    cwd: &Path,
230    agent_type: &str,
231    session_id: &str,
232) -> anyhow::Result<()> {
233    let start_time = std::time::Instant::now();
234    info!(
235        "Starting Nexus session start pipeline for {} ({})",
236        agent_type, session_id
237    );
238
239    // 1. Resolve Project Identity
240    let project = nexus_core::ProjectIdentity::resolve(cwd);
241    let nexus_dir = project.root_dir.join(".nexus");
242    fs::create_dir_all(&nexus_dir)?;
243    fs::create_dir_all(nexus_dir.join("cache"))?;
244    fs::create_dir_all(nexus_dir.join("sessions"))?;
245
246    // 2. Initialize Repositories (for Morning Recall)
247    let config = nexus_core::Config::from_env().unwrap_or_default();
248    // SQLite create_if_missing won't create parent dirs
249    if let Some(parent) = config.database.path.parent() {
250        let _ = fs::create_dir_all(parent);
251    }
252    let mut storage = nexus_storage::StorageManager::from_url(&config.database_url()).await?;
253    storage.initialize().await?;
254    let memory_repo = nexus_storage::repository::MemoryRepository::new(storage.pool().clone());
255    let ns_repo = nexus_storage::repository::NamespaceRepository::new(storage.pool().clone());
256    let namespace = ns_repo.get_or_create(agent_type, agent_type).await?;
257
258    // 3. Load Cache and Perform Morning Recall
259    let cache = nexus_agent::cognitive_cache::CognitiveCache::load_or_init(&nexus_dir);
260
261    // Attempt to get embedder if available
262    let embedder = if config.embedding.enabled {
263        nexus_agent::runtime::create_embedding_service(&config).await
264    } else {
265        None
266    };
267
268    let recalls = cache
269        .morning_recall(
270            &project,
271            namespace.id,
272            &memory_repo,
273            embedder
274                .as_ref()
275                .map(|e| e.as_ref() as &dyn nexus_core::EmbeddingService),
276        )
277        .await;
278
279    // 4. Build and Write context.md
280    let window_size = nexus_agent::TokenBudget::estimate_window(agent_type) as f32;
281    let max_context_tokens =
282        (window_size * config.cognitive_system.context_allocation_pct) as usize;
283    let context_md = nexus_agent::context_builder::build_context_md(
284        &cache.hot_cache,
285        &recalls,
286        max_context_tokens,
287    );
288
289    let context_path = nexus_dir.join("context.md");
290    atomic_write(&context_path, &context_md)?;
291
292    // 5. Compute soul path for injection reference
293    // (soul.md is only created/modified during deep dream cycles, per spec)
294    let soul_path = dirs::config_dir()
295        .unwrap_or_else(|| PathBuf::from("."))
296        .join("nexus")
297        .join("soul.md");
298
299    // 6. Inject references into agent configs
300    if let Some(target) = AgentInjectionTarget::find(agent_type) {
301        // Project config
302        let project_config = project.root_dir.join(&target.project_config_filename);
303        inject_reference(&project_config, &soul_path, &context_path)?;
304
305        // Global config
306        if let Some(global_config) = target.global_config {
307            inject_soul_only(&global_config, &soul_path)?;
308        }
309    }
310
311    // 7. Start session scratch file
312    let session_manager = nexus_agent::session_manager::SessionManager::new(&project.root_dir);
313    session_manager.start_session(session_id, agent_type)?;
314
315    // 8. Hardening: .gitignore — ensure .nexus/ is always ignored
316    let gitignore = project.root_dir.join(".gitignore");
317    let gitignore_content = fs::read_to_string(&gitignore).unwrap_or_default();
318    let has_nexus_entry = gitignore_content.lines().any(|line| {
319        let trimmed = line.trim();
320        trimmed == ".nexus" || trimmed == ".nexus/" || trimmed == "/.nexus" || trimmed == "/.nexus/"
321    });
322    if !has_nexus_entry {
323        let mut f = fs::OpenOptions::new()
324            .create(true)
325            .append(true)
326            .open(&gitignore)?;
327        if !gitignore_content.is_empty() && !gitignore_content.ends_with('\n') {
328            writeln!(f)?;
329        }
330        writeln!(f, ".nexus/")?;
331    }
332
333    info!(
334        "Nexus session start pipeline completed in {:?}",
335        start_time.elapsed()
336    );
337    Ok(())
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343    use tempfile::tempdir;
344
345    #[test]
346    fn test_inject_reference_idempotency() {
347        let dir = tempdir().unwrap();
348        let config = dir.path().join("CLAUDE.md");
349        fs::write(&config, "# Existing Content\n").unwrap();
350
351        let soul = PathBuf::from("/tmp/soul.md");
352        let context = PathBuf::from("/tmp/context.md");
353
354        // 1. First injection
355        inject_reference(&config, &soul, &context).unwrap();
356        let content1 = fs::read_to_string(&config).unwrap();
357        assert!(content1.contains(NEXUS_BLOCK_START));
358
359        // 2. Second injection (idempotent)
360        inject_reference(&config, &soul, &context).unwrap();
361        let content2 = fs::read_to_string(&config).unwrap();
362        assert_eq!(content1, content2);
363    }
364
365    #[test]
366    fn test_remove_reference() {
367        let dir = tempdir().unwrap();
368        let config = dir.path().join("AGENTS.md");
369        fs::write(
370            &config,
371            "# Top\n<!-- NEXUS:START -->\n- Ref\n<!-- NEXUS:END -->\n# Bottom",
372        )
373        .unwrap();
374
375        remove_reference(&config).unwrap();
376        let content = fs::read_to_string(&config).unwrap();
377        assert!(!content.contains("NEXUS:START"));
378        assert!(content.contains("# Top"));
379        assert!(content.contains("# Bottom"));
380    }
381
382    #[tokio::test]
383    async fn test_on_session_start_creates_structure() {
384        let dir = tempdir().unwrap();
385        // Use NEXUS_DATABASE_PATH to isolate DB instead of HOME manipulation
386        let db_path = dir.path().join("test.db");
387        let original_db = std::env::var("NEXUS_DATABASE_PATH").ok();
388        std::env::set_var("NEXUS_DATABASE_PATH", &db_path);
389
390        let result = on_session_start(dir.path(), "claude-code", "test-session").await;
391
392        // Restore before assertions
393        if let Some(orig) = original_db {
394            std::env::set_var("NEXUS_DATABASE_PATH", orig);
395        } else {
396            std::env::remove_var("NEXUS_DATABASE_PATH");
397        }
398
399        result.unwrap();
400
401        assert!(dir.path().join(".nexus").exists());
402        assert!(dir.path().join(".nexus/context.md").exists());
403        assert!(dir.path().join(".nexus/sessions/test-session.md").exists());
404    }
405
406    #[test]
407    fn test_pi_mono_injection_target_exists() {
408        let target = AgentInjectionTarget::find("pi-mono");
409        assert!(target.is_some(), "pi-mono must be in known_agents()");
410
411        let target = target.unwrap();
412        assert_eq!(target.agent_type, "pi-mono");
413        assert!(target.global_config.is_some());
414        assert_eq!(target.project_config_filename, ".pi/AGENTS.md");
415
416        let global = target.global_config.unwrap();
417        assert!(
418            global.ends_with(".pi/agent/AGENTS.md")
419                || global.to_string_lossy().contains(".pi/agent/AGENTS.md")
420        );
421    }
422}