nexus_memory_hooks/agents/
pi_mono.rs

1//! Pi-Mono hook implementation
2//!
3//! Pi-mono is a TypeScript/Bun-based coding agent with a rich extension system.
4//! Integration uses a native TypeScript extension installed at
5//! `~/.pi/agent/extensions/nexus-memory.ts` that hooks into pi's
6//! lifecycle events (session_start, session_shutdown, agent_end, etc.)
7//! and forwards them to the Nexus CLI for memory capture.
8//!
9//! Repository: https://github.com/badlogic/pi-mono
10//! Stack: TypeScript, Node.js/Bun runtime
11//! Config: ~/.pi/agent/extensions/
12//! Detection: `pi` or `pi-coding-agent` process
13
14use async_trait::async_trait;
15use std::path::PathBuf;
16
17use crate::base::{AgentHook, BaseHook, LifecycleCapabilities, SessionEndCallback};
18use crate::error::{HookError, Result};
19use crate::monitor::ProcessMonitor;
20use crate::session::{
21    FileAction, FileInfo, SessionContext, SubagentExecution, TaskInfo, TaskStatus,
22};
23use crate::types::{AgentType, SessionActivity, SupportTier};
24
25/// Pi-Mono hook for extracting memory from pi-mono session execution.
26///
27/// Uses a native TypeScript extension (not SKILL.md) that integrates
28/// with pi-mono's extension API for full lifecycle event coverage.
29///
30/// # Integration Points
31///
32/// - **Extension:** `~/.pi/agent/extensions/nexus-memory.ts`
33/// - **Transport:** CLI (`nexus session start`, `nexus ingest-hook-event`)
34/// - **Fallback:** Process monitoring + `.pi/sessions/` file scanning
35///
36/// # Lifecycle Coverage
37///
38/// All five lifecycle events are supported:
39/// - `session_start` — via extension `session_start` event
40/// - `session_end` — via extension `session_shutdown` event
41/// - `checkpoint` — via debounced `agent_end` and explicit triggers
42/// - `error_hook` — synthetic from failed `tool_result` / abnormal `agent_end`
43/// - `compact` — via extension `session_compact` event
44pub struct PiMonoHook {
45    /// Base hook functionality
46    base: BaseHook,
47
48    /// Config directory
49    config_dir: PathBuf,
50
51    /// Session directory
52    session_dir: PathBuf,
53
54    /// Extensions directory
55    extensions_dir: PathBuf,
56
57    /// Process monitor
58    process_monitor: ProcessMonitor,
59
60    /// Whether extension is installed
61    extension_installed: bool,
62}
63
64impl PiMonoHook {
65    /// Agent type string
66    pub const AGENT_TYPE: &'static str = "pi-mono";
67
68    /// Config directory name
69    pub const CONFIG_DIR_NAME: &'static str = ".pi";
70
71    /// Extensions subdirectory
72    pub const EXTENSIONS_SUBDIR: &'static str = "agent/extensions";
73
74    /// Sessions subdirectory
75    pub const SESSIONS_SUBDIR: &'static str = "sessions";
76
77    /// Logs subdirectory
78    pub const LOGS_SUBDIR: &'static str = "logs";
79
80    /// Create a new Pi-Mono hook
81    pub fn new() -> Self {
82        Self::new_with_install(true)
83    }
84
85    /// Create a new Pi-Mono hook without mutating user state.
86    pub fn new_readonly() -> Self {
87        Self::new_with_install(false)
88    }
89
90    fn new_with_install(auto_install: bool) -> Self {
91        let config_dir = dirs::home_dir()
92            .unwrap_or_else(|| PathBuf::from("."))
93            .join(Self::CONFIG_DIR_NAME);
94
95        let session_dir = config_dir.join(Self::SESSIONS_SUBDIR);
96        let extensions_dir = config_dir.join(Self::EXTENSIONS_SUBDIR);
97        let extension_installed = Self::extension_file_path(&extensions_dir).exists();
98
99        let mut hook = Self {
100            base: BaseHook::new(Self::AGENT_TYPE),
101            config_dir,
102            session_dir,
103            extensions_dir,
104            process_monitor: ProcessMonitor::new(),
105            extension_installed,
106        };
107
108        if auto_install {
109            // Migrate from legacy SKILL.md
110            hook.migrate_from_skill();
111
112            if !hook.extension_installed {
113                if let Err(e) = hook.install_extension() {
114                    tracing::warn!("Failed to install pi-mono extension: {}", e);
115                }
116            }
117        }
118
119        hook
120    }
121
122    fn extension_file_path(extensions_dir: &std::path::Path) -> PathBuf {
123        extensions_dir.join("nexus-memory.ts")
124    }
125
126    /// Install the nexus-memory TypeScript extension
127    fn install_extension(&mut self) -> Result<()> {
128        std::fs::create_dir_all(&self.extensions_dir).map_err(|e| {
129            HookError::InstallationFailed(format!("Failed to create extensions dir: {}", e))
130        })?;
131
132        let extension_path = Self::extension_file_path(&self.extensions_dir);
133
134        let extension_content = include_str!("../extension_ts/nexus_memory_pi.ts");
135
136        std::fs::write(&extension_path, extension_content).map_err(|e| {
137            HookError::InstallationFailed(format!("Failed to write extension: {}", e))
138        })?;
139
140        self.extension_installed = true;
141        tracing::info!("Pi-mono extension installed at: {:?}", extension_path);
142
143        Ok(())
144    }
145
146    /// Migrate from legacy SKILL.md installation
147    fn migrate_from_skill(&mut self) {
148        let legacy_skill_dir = self
149            .config_dir
150            .join("agent")
151            .join("skills")
152            .join("nexus-memory-extraction");
153
154        if legacy_skill_dir.exists() {
155            tracing::info!("Migrating pi-mono from SKILL.md to TypeScript extension");
156            if let Err(e) = std::fs::remove_dir_all(&legacy_skill_dir) {
157                tracing::warn!("Failed to remove legacy skill dir: {}", e);
158            }
159        }
160    }
161
162    /// Read session files
163    fn read_session_files(&self) -> Vec<serde_json::Value> {
164        let mut sessions = Vec::new();
165
166        if !self.session_dir.exists() {
167            return sessions;
168        }
169
170        if let Ok(entries) = std::fs::read_dir(&self.session_dir) {
171            let mut session_files: Vec<_> = entries
172                .filter_map(|e| e.ok())
173                .filter(|e| {
174                    e.path()
175                        .extension()
176                        .map(|ext| ext == "json")
177                        .unwrap_or(false)
178                })
179                .collect();
180
181            // Sort by modification time, most recent first
182            session_files.sort_by(|a, b| {
183                let time_a = a.metadata().ok().and_then(|m| m.modified().ok());
184                let time_b = b.metadata().ok().and_then(|m| m.modified().ok());
185                time_b.cmp(&time_a)
186            });
187
188            // Take up to 10 most recent sessions
189            for entry in session_files.into_iter().take(10) {
190                if let Ok(content) = std::fs::read_to_string(entry.path()) {
191                    if let Ok(data) = serde_json::from_str(&content) {
192                        sessions.push(data);
193                    }
194                }
195            }
196        }
197
198        sessions
199    }
200
201    /// Read log files
202    fn read_log_files(&self) -> Vec<String> {
203        let mut commands = Vec::new();
204        let logs_dir = self.config_dir.join(Self::LOGS_SUBDIR);
205
206        if !logs_dir.exists() {
207            return commands;
208        }
209
210        if let Ok(entries) = std::fs::read_dir(&logs_dir) {
211            let mut log_files: Vec<_> = entries
212                .filter_map(|e| e.ok())
213                .filter(|e| {
214                    e.path()
215                        .extension()
216                        .map(|ext| ext == "log")
217                        .unwrap_or(false)
218                })
219                .collect();
220
221            log_files.sort_by(|a, b| {
222                let time_a = a.metadata().ok().and_then(|m| m.modified().ok());
223                let time_b = b.metadata().ok().and_then(|m| m.modified().ok());
224                time_b.cmp(&time_a)
225            });
226
227            for entry in log_files.into_iter().take(5) {
228                if let Ok(content) = std::fs::read_to_string(entry.path()) {
229                    for line in content.lines() {
230                        if line.contains("Executing:") || line.contains("Command:") {
231                            commands.push(line.to_string());
232                        }
233                    }
234                }
235            }
236        }
237
238        commands
239    }
240}
241
242impl Default for PiMonoHook {
243    fn default() -> Self {
244        Self::new()
245    }
246}
247
248#[async_trait]
249impl AgentHook for PiMonoHook {
250    fn agent_type(&self) -> &str {
251        &self.base.agent_type
252    }
253
254    async fn install_session_start_hook(&mut self, callback: SessionEndCallback) -> Result<()> {
255        self.base.add_session_start_callback(callback);
256        Ok(())
257    }
258
259    async fn install_session_end_hook(&mut self, callback: SessionEndCallback) -> Result<()> {
260        self.base.add_callback(callback);
261        Ok(())
262    }
263
264    async fn install_checkpoint_hook(&mut self, callback: SessionEndCallback) -> Result<()> {
265        self.base.add_checkpoint_callback(callback);
266        Ok(())
267    }
268
269    async fn install_compact_hook(&mut self, callback: SessionEndCallback) -> Result<()> {
270        self.base.add_callback(callback);
271        Ok(())
272    }
273
274    async fn install_error_hook(&mut self, callback: SessionEndCallback) -> Result<()> {
275        self.base.add_error_callback(callback);
276        Ok(())
277    }
278
279    async fn detect_session_activity(&self) -> Result<SessionActivity> {
280        let mut monitor = self.process_monitor.clone();
281        let processes = monitor.find_agent_processes(AgentType::PiMono);
282
283        let mut activity = SessionActivity::new(AgentType::PiMono);
284
285        if !processes.is_empty() {
286            activity.is_active = true;
287            activity.processes = processes;
288        }
289
290        // Check for recent session file activity
291        if self.session_dir.exists() {
292            if let Ok(entries) = std::fs::read_dir(&self.session_dir) {
293                if let Some(most_recent) = entries
294                    .filter_map(|e| e.ok())
295                    .filter(|e| {
296                        e.path()
297                            .extension()
298                            .map(|ext| ext == "json")
299                            .unwrap_or(false)
300                    })
301                    .max_by_key(|e| e.metadata().ok().and_then(|m| m.modified().ok()))
302                {
303                    if let Ok(metadata) = most_recent.metadata() {
304                        if let Ok(modified) = metadata.modified() {
305                            let age = std::time::SystemTime::now()
306                                .duration_since(modified)
307                                .unwrap_or(std::time::Duration::MAX);
308
309                            // Consider active if modified in last 5 minutes
310                            if age.as_secs() < 300 {
311                                activity.is_active = true;
312                                activity.session_id = Some(
313                                    most_recent
314                                        .path()
315                                        .file_stem()
316                                        .unwrap()
317                                        .to_string_lossy()
318                                        .to_string(),
319                                );
320                            }
321                        }
322                    }
323                }
324            }
325        }
326
327        // Check for subagent processes
328        let subagent_check = std::process::Command::new("pgrep")
329            .arg("-f")
330            .arg("subagent|skill")
331            .output()
332            .ok();
333
334        if let Some(output) = subagent_check {
335            if output.status.success() && !output.stdout.is_empty() {
336                activity.is_active = true;
337            }
338        }
339
340        Ok(activity)
341    }
342
343    async fn extract_session_context(&self) -> Result<SessionContext> {
344        let mut context = SessionContext::new("pi-mono")
345            .with_source("native")
346            .with_reliability(1.0);
347
348        // Track role usage
349        let mut role_usage: std::collections::HashMap<String, i32> =
350            std::collections::HashMap::new();
351
352        // Extract from session files
353        for session_data in self.read_session_files() {
354            // Extract session info
355            if let Some(timestamp) = session_data.get("timestamp").and_then(|t| t.as_str()) {
356                context.add_custom(
357                    "session_timestamp",
358                    serde_json::Value::String(timestamp.to_string()),
359                );
360            }
361
362            // Extract tasks
363            if let Some(tasks) = session_data.get("tasks").and_then(|t| t.as_array()) {
364                for task in tasks {
365                    let description = task
366                        .get("description")
367                        .and_then(|d| d.as_str())
368                        .unwrap_or("");
369                    let role = task
370                        .get("role")
371                        .and_then(|r| r.as_str())
372                        .unwrap_or("unknown");
373
374                    let mut task_info = TaskInfo::new(description);
375                    task_info.subagent = Some(role.to_string());
376
377                    if let Some(status) = task.get("status").and_then(|s| s.as_str()) {
378                        task_info.status = match status {
379                            "completed" => TaskStatus::Completed,
380                            "failed" => TaskStatus::Failed,
381                            "in_progress" => TaskStatus::InProgress,
382                            _ => TaskStatus::Pending,
383                        };
384                    }
385
386                    context.tasks.push(task_info);
387
388                    // Track role usage
389                    *role_usage.entry(role.to_string()).or_insert(0) += 1;
390
391                    // Add as subagent execution
392                    context.subagent_executions.push(SubagentExecution {
393                        subagent_type: role.to_string(),
394                        task: description.to_string(),
395                        status: "completed".to_string(),
396                        started_at: chrono::Utc::now(),
397                        completed_at: Some(chrono::Utc::now()),
398                        result_summary: None,
399                    });
400                }
401            }
402
403            // Extract files modified
404            if let Some(files) = session_data
405                .get("files_modified")
406                .and_then(|f| f.as_array())
407            {
408                for file in files {
409                    if let Some(path) = file.as_str() {
410                        context.add_file(FileInfo::new(path, FileAction::Modified));
411                    }
412                }
413            }
414        }
415
416        // Extract commands from logs
417        for cmd in self.read_log_files() {
418            context.add_command(cmd);
419        }
420
421        // Store role usage stats
422        context.add_custom(
423            "role_usage",
424            serde_json::to_value(&role_usage).unwrap_or(serde_json::Value::Null),
425        );
426
427        // Get git status
428        let git_status = std::process::Command::new("git")
429            .args(["status", "--porcelain"])
430            .output()
431            .ok();
432
433        if let Some(output) = git_status {
434            if output.status.success() {
435                let status = String::from_utf8_lossy(&output.stdout);
436                for line in status.lines() {
437                    if line.len() > 3 {
438                        let file_path = &line[3..];
439                        context.add_file(FileInfo::new(file_path, FileAction::Modified));
440                    }
441                }
442            }
443        }
444
445        context.complete();
446        Ok(context)
447    }
448
449    fn is_hook_installed(&self) -> bool {
450        self.extension_installed
451    }
452
453    fn reliability_score(&self) -> f32 {
454        if self.extension_installed {
455            1.0
456        } else {
457            0.95
458        }
459    }
460
461    fn lifecycle_capabilities(&self) -> LifecycleCapabilities {
462        LifecycleCapabilities {
463            session_start: true,
464            session_end: true,
465            checkpoint: true,
466            error_hook: true,
467            compact: true,
468        }
469    }
470
471    fn support_tier(&self) -> SupportTier {
472        SupportTier::NativeLifecycle
473    }
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479    use std::sync::Arc;
480
481    #[test]
482    fn test_pi_mono_hook_new() {
483        let hook = PiMonoHook::new();
484        assert_eq!(hook.agent_type(), "pi-mono");
485    }
486
487    #[tokio::test]
488    async fn test_pi_mono_hook_detect_activity() {
489        let hook = PiMonoHook::new();
490        let activity = hook.detect_session_activity().await.unwrap();
491
492        assert_eq!(activity.agent_type, AgentType::PiMono);
493    }
494
495    #[test]
496    fn test_pi_mono_hook_constants() {
497        assert_eq!(PiMonoHook::AGENT_TYPE, "pi-mono");
498        assert_eq!(PiMonoHook::CONFIG_DIR_NAME, ".pi");
499        assert_eq!(PiMonoHook::EXTENSIONS_SUBDIR, "agent/extensions");
500    }
501
502    #[test]
503    fn test_pi_mono_extension_file_path() {
504        let dir = tempfile::tempdir().unwrap();
505        let path = PiMonoHook::extension_file_path(&dir.path().join("extensions"));
506        assert_eq!(
507            path.file_name().unwrap().to_str().unwrap(),
508            "nexus-memory.ts"
509        );
510    }
511
512    #[test]
513    fn test_pi_mono_hook_lifecycle_capabilities_full() {
514        let hook = PiMonoHook::new();
515        let caps = hook.lifecycle_capabilities();
516
517        assert!(caps.session_start, "pi-mono should support session_start");
518        assert!(caps.session_end, "pi-mono should support session_end");
519        assert!(caps.checkpoint, "pi-mono should support checkpoint");
520        assert!(caps.error_hook, "pi-mono should support error_hook");
521        assert!(caps.compact, "pi-mono should support compact");
522    }
523
524    #[tokio::test]
525    async fn test_pi_mono_hook_install_session_start() {
526        let mut hook = PiMonoHook::new();
527        let cb: SessionEndCallback = Arc::new(|_ctx| ());
528        let result = hook.install_session_start_hook(cb).await;
529        assert!(result.is_ok(), "pi-mono should accept session_start hook");
530    }
531
532    #[tokio::test]
533    async fn test_pi_mono_hook_install_checkpoint() {
534        let mut hook = PiMonoHook::new();
535        let cb: SessionEndCallback = Arc::new(|_ctx| ());
536        let result = hook.install_checkpoint_hook(cb).await;
537        assert!(result.is_ok(), "pi-mono should accept checkpoint hook");
538    }
539
540    #[tokio::test]
541    async fn test_pi_mono_hook_install_error() {
542        let mut hook = PiMonoHook::new();
543        let cb: SessionEndCallback = Arc::new(|_ctx| ());
544        let result = hook.install_error_hook(cb).await;
545        assert!(result.is_ok(), "pi-mono should accept error hook");
546    }
547
548    #[test]
549    fn test_pi_mono_legacy_migration() {
550        let dir = tempfile::tempdir().unwrap();
551        let legacy_dir = dir
552            .path()
553            .join(".pi")
554            .join("agent")
555            .join("skills")
556            .join("nexus-memory-extraction");
557        std::fs::create_dir_all(&legacy_dir).unwrap();
558        std::fs::write(legacy_dir.join("SKILL.md"), "legacy").unwrap();
559
560        // Simulate migration: remove directory
561        assert!(legacy_dir.exists());
562        std::fs::remove_dir_all(&legacy_dir).unwrap();
563        assert!(!legacy_dir.exists());
564    }
565}
nexus_memory_hooks/agents/pi_mono.rs

nexus_memory_hooks/agents/
pi_mono.rs
