mi6_core/framework/

pi.rs

//! Pi Coding Agent framework adapter.
//!
//! This adapter handles integration with the Pi Coding Agent
//! (<https://github.com/badlogic/pi-mono>).
//!
//! Pi uses a TypeScript extension system rather than JSON/TOML hooks.
//! mi6 integrates by installing a TypeScript extension that subscribes
//! to Pi's event bus and forwards events to `mi6 ingest event`.
//!
//! # Installation Flow
//!
//! 1. `mi6 enable pi` writes the embedded TypeScript extension to
//!    `~/.pi/agent/extensions/mi6.ts` (user scope) or
//!    `.pi/extensions/mi6.ts` (project scope)
//! 2. Pi automatically discovers and loads extensions from these directories
//! 3. The extension sets `PI_SESSION_ID` env var for framework detection
//!
//! # Uninstallation
//!
//! `mi6 disable pi` removes the extension file.

use super::{FrameworkAdapter, ParsedHookInput, common};
use crate::model::EventType;
use crate::model::error::InitError;
use std::path::{Path, PathBuf};

/// Embedded mi6.ts extension content.
const MI6_EXTENSION_TS: &str = include_str!("pi-extension/mi6.ts");

/// Pi Coding Agent framework adapter.
///
/// Handles extension installation and event parsing for Pi.
pub struct PiAdapter;

impl FrameworkAdapter for PiAdapter {
    fn name(&self) -> &'static str {
        "pi"
    }

    fn display_name(&self) -> &'static str {
        "Pi Coding Agent"
    }

    fn project_config_path(&self) -> PathBuf {
        PathBuf::from(".pi/extensions/mi6.ts")
    }

    fn user_config_path(&self) -> Option<PathBuf> {
        dirs::home_dir().map(|h| h.join(".pi/agent/extensions/mi6.ts"))
    }

    fn generate_hooks_config(
        &self,
        _enabled_events: &[EventType],
        _mi6_bin: &str,
        _otel_enabled: bool,
        _otel_port: u16,
    ) -> serde_json::Value {
        // Pi uses TypeScript extensions, not JSON config.
        // Return a placeholder that describes what will be installed.
        serde_json::json!({
            "_note": "Pi uses TypeScript extensions, not JSON hooks",
            "extension_path": "~/.pi/agent/extensions/mi6.ts",
            "extension_content": "See embedded mi6.ts"
        })
    }

    fn merge_config(
        &self,
        generated: serde_json::Value,
        _existing: Option<serde_json::Value>,
    ) -> serde_json::Value {
        // No merging needed - we write a standalone file
        generated
    }

    fn parse_hook_input(
        &self,
        _event_type: &str,
        stdin_json: &serde_json::Value,
    ) -> ParsedHookInput {
        // Pi extension sends JSON with these fields:
        // {
        //   "session_id": "uuid",
        //   "tool_name": "Bash",
        //   "tool_use_id": "...",
        //   "cwd": "/path/to/project",
        //   "trigger": "auto" (for PreCompact),
        //   "model": "anthropic/claude-sonnet-4-20250514",
        //   "user_prompt": "hello there",
        //   "transcript_path": "/path/to/session.jsonl"
        // }
        ParsedHookInput {
            session_id: stdin_json
                .get("session_id")
                .and_then(|v| v.as_str())
                .map(String::from),
            tool_use_id: stdin_json
                .get("tool_use_id")
                .and_then(|v| v.as_str())
                .map(String::from),
            tool_name: stdin_json
                .get("tool_name")
                .and_then(|v| v.as_str())
                .map(String::from),
            cwd: stdin_json
                .get("cwd")
                .and_then(|v| v.as_str())
                .map(String::from),
            compact_trigger: stdin_json
                .get("trigger")
                .and_then(|v| v.as_str())
                .map(String::from),
            model: stdin_json
                .get("model")
                .and_then(|v| v.as_str())
                .map(String::from),
            transcript_path: stdin_json
                .get("transcript_path")
                .and_then(|v| v.as_str())
                .map(String::from),
            prompt: stdin_json
                .get("user_prompt")
                .and_then(|v| v.as_str())
                .map(String::from),
            // Token metrics from ApiRequest events
            tokens_input: stdin_json.get("tokens_input").and_then(|v| v.as_i64()),
            tokens_output: stdin_json.get("tokens_output").and_then(|v| v.as_i64()),
            tokens_cache_read: stdin_json.get("tokens_cache_read").and_then(|v| v.as_i64()),
            tokens_cache_write: stdin_json
                .get("tokens_cache_write")
                .and_then(|v| v.as_i64()),
            cost_usd: stdin_json.get("cost_usd").and_then(|v| v.as_f64()),
            // Pi doesn't have these Claude-specific features
            subagent_type: None,
            spawned_agent_id: None,
            permission_mode: None,
            session_source: None,
            agent_id: None,
            agent_transcript_path: None,
            duration_ms: None,
        }
    }

    fn map_event_type(&self, framework_event: &str) -> EventType {
        // Pi extension sends canonical event names
        framework_event
            .parse()
            .unwrap_or_else(|_| EventType::Custom(framework_event.to_string()))
    }

    fn supported_events(&self) -> Vec<&'static str> {
        vec![
            "SessionStart",
            "SessionEnd",
            "PreToolUse",
            "PostToolUse",
            "UserPromptSubmit",
            "Stop",
            "PreCompact",
            "ApiRequest",
        ]
    }

    fn detection_env_vars(&self) -> &[&'static str] {
        // The mi6 extension sets this when Pi starts
        &["PI_SESSION_ID"]
    }

    fn is_installed(&self) -> bool {
        common::is_framework_installed(dirs::home_dir().map(|h| h.join(".pi/agent")), "pi")
    }

    fn remove_hooks(&self, _existing: serde_json::Value) -> Option<serde_json::Value> {
        // Pi uses extension files, not JSON hooks
        // Removal is handled by uninstall_hooks
        None
    }

    // ========================================================================
    // Extension-based installation
    // ========================================================================

    fn settings_path(&self, local: bool, _settings_local: bool) -> Result<PathBuf, InitError> {
        if local {
            Ok(self.project_config_path())
        } else {
            self.user_config_path()
                .ok_or_else(|| InitError::Config("could not determine home directory".into()))
        }
    }

    fn has_mi6_hooks(&self, local: bool, settings_local: bool) -> bool {
        self.settings_path(local, settings_local)
            .map(|p| p.exists())
            .unwrap_or(false)
    }

    fn install_hooks(
        &self,
        path: &Path,
        _hooks: &serde_json::Value,
        _otel_env: Option<serde_json::Value>,
        _remove_otel: bool,
    ) -> Result<(), InitError> {
        // Create parent directory if needed
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent).map_err(|e| {
                InitError::Config(format!("failed to create {}: {e}", parent.display()))
            })?;
        }

        // Write the TypeScript extension
        std::fs::write(path, MI6_EXTENSION_TS)
            .map_err(|e| InitError::Config(format!("failed to write {}: {e}", path.display())))?;

        Ok(())
    }

    fn uninstall_hooks(&self, local: bool, settings_local: bool) -> Result<bool, InitError> {
        let path = self.settings_path(local, settings_local)?;

        if path.exists() {
            std::fs::remove_file(&path).map_err(|e| {
                InitError::Config(format!("failed to remove {}: {e}", path.display()))
            })?;
            Ok(true)
        } else {
            Ok(false)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_name() {
        let adapter = PiAdapter;
        assert_eq!(adapter.name(), "pi");
        assert_eq!(adapter.display_name(), "Pi Coding Agent");
    }

    #[test]
    fn test_project_config_path() {
        let adapter = PiAdapter;
        assert_eq!(
            adapter.project_config_path(),
            PathBuf::from(".pi/extensions/mi6.ts")
        );
    }

    #[test]
    fn test_user_config_path() {
        let adapter = PiAdapter;
        let path = adapter.user_config_path();
        assert!(path.is_some());
        let path = path.unwrap();
        assert!(path.ends_with(".pi/agent/extensions/mi6.ts"));
    }

    #[test]
    fn test_map_event_type() {
        let adapter = PiAdapter;

        // Canonical events should parse correctly
        assert_eq!(
            adapter.map_event_type("SessionStart"),
            EventType::SessionStart
        );
        assert_eq!(adapter.map_event_type("PreToolUse"), EventType::PreToolUse);
        assert_eq!(
            adapter.map_event_type("PostToolUse"),
            EventType::PostToolUse
        );
        assert_eq!(
            adapter.map_event_type("UserPromptSubmit"),
            EventType::UserPromptSubmit
        );
        assert_eq!(adapter.map_event_type("Stop"), EventType::Stop);
        assert_eq!(adapter.map_event_type("PreCompact"), EventType::PreCompact);

        // Unknown events become Custom
        assert_eq!(
            adapter.map_event_type("UnknownEvent"),
            EventType::Custom("UnknownEvent".to_string())
        );
    }

    #[test]
    fn test_parse_hook_input() {
        let adapter = PiAdapter;
        let input = serde_json::json!({
            "session_id": "pi-session-123",
            "tool_use_id": "tool-456",
            "tool_name": "Bash",
            "cwd": "/projects/test"
        });

        let parsed = adapter.parse_hook_input("PreToolUse", &input);

        assert_eq!(parsed.session_id, Some("pi-session-123".to_string()));
        assert_eq!(parsed.tool_use_id, Some("tool-456".to_string()));
        assert_eq!(parsed.tool_name, Some("Bash".to_string()));
        assert_eq!(parsed.cwd, Some("/projects/test".to_string()));
        // Pi doesn't provide Claude-specific fields
        assert_eq!(parsed.permission_mode, None);
        assert_eq!(parsed.subagent_type, None);
        assert_eq!(parsed.spawned_agent_id, None);
    }

    #[test]
    fn test_parse_hook_input_with_model_and_prompt() {
        let adapter = PiAdapter;
        let input = serde_json::json!({
            "session_id": "pi-session-123",
            "cwd": "/projects/test",
            "model": "anthropic/claude-sonnet-4-20250514",
            "user_prompt": "hello there",
            "transcript_path": "/home/user/.pi/agent/sessions/test/session.jsonl"
        });

        let parsed = adapter.parse_hook_input("SessionStart", &input);

        assert_eq!(parsed.session_id, Some("pi-session-123".to_string()));
        assert_eq!(
            parsed.model,
            Some("anthropic/claude-sonnet-4-20250514".to_string())
        );
        assert_eq!(parsed.prompt, Some("hello there".to_string()));
        assert_eq!(
            parsed.transcript_path,
            Some("/home/user/.pi/agent/sessions/test/session.jsonl".to_string())
        );
    }

    #[test]
    fn test_parse_hook_input_with_trigger() {
        let adapter = PiAdapter;
        let input = serde_json::json!({
            "session_id": "pi-123",
            "trigger": "auto"
        });

        let parsed = adapter.parse_hook_input("PreCompact", &input);

        assert_eq!(parsed.session_id, Some("pi-123".to_string()));
        assert_eq!(parsed.compact_trigger, Some("auto".to_string()));
    }

    #[test]
    fn test_parse_hook_input_with_tokens() {
        let adapter = PiAdapter;
        let input = serde_json::json!({
            "session_id": "pi-123",
            "model": "gpt-4.1",
            "tokens_input": 942,
            "tokens_output": 64,
            "tokens_cache_read": 0,
            "tokens_cache_write": 0,
            "cost_usd": 0.00239
        });

        let parsed = adapter.parse_hook_input("ApiRequest", &input);

        assert_eq!(parsed.session_id, Some("pi-123".to_string()));
        assert_eq!(parsed.model, Some("gpt-4.1".to_string()));
        assert_eq!(parsed.tokens_input, Some(942));
        assert_eq!(parsed.tokens_output, Some(64));
        assert_eq!(parsed.tokens_cache_read, Some(0));
        assert_eq!(parsed.tokens_cache_write, Some(0));
        assert!((parsed.cost_usd.unwrap() - 0.00239).abs() < 0.0001);
    }

    #[test]
    fn test_supported_events() {
        let adapter = PiAdapter;
        let events = adapter.supported_events();

        assert!(events.contains(&"SessionStart"));
        assert!(events.contains(&"SessionEnd"));
        assert!(events.contains(&"PreToolUse"));
        assert!(events.contains(&"PostToolUse"));
        assert!(events.contains(&"UserPromptSubmit"));
        assert!(events.contains(&"Stop"));
        assert!(events.contains(&"PreCompact"));
        assert!(events.contains(&"ApiRequest"));
        assert_eq!(events.len(), 8);
    }

    #[test]
    fn test_detection_env_vars() {
        let adapter = PiAdapter;
        let vars = adapter.detection_env_vars();

        assert!(vars.contains(&"PI_SESSION_ID"));
        assert_eq!(vars.len(), 1);
    }

    #[test]
    fn test_embedded_extension() {
        // Verify the embedded extension uses correct Pi API patterns
        assert!(MI6_EXTENSION_TS.contains("mi6"));
        assert!(MI6_EXTENSION_TS.contains("ingestEvent"));
        // Uses Pi's snake_case event names
        assert!(MI6_EXTENSION_TS.contains("session_start"));
        assert!(MI6_EXTENSION_TS.contains("session_shutdown"));
        assert!(MI6_EXTENSION_TS.contains("tool_call"));
        assert!(MI6_EXTENSION_TS.contains("tool_result"));
        assert!(MI6_EXTENSION_TS.contains("turn_start"));
        assert!(MI6_EXTENSION_TS.contains("turn_end"));
        assert!(MI6_EXTENSION_TS.contains("session_before_compact"));
        // Uses correct Pi import
        assert!(MI6_EXTENSION_TS.contains("@mariozechner/pi-coding-agent"));
        // Uses factory function pattern
        assert!(MI6_EXTENSION_TS.contains("export default function"));
        // Uses correct event handler signature (event, ctx)
        assert!(MI6_EXTENSION_TS.contains("ExtensionContext"));
    }

    #[test]
    fn test_generate_hooks_config() {
        let adapter = PiAdapter;
        let events = vec![EventType::SessionStart, EventType::PreToolUse];

        let config = adapter.generate_hooks_config(&events, "mi6", false, 4318);

        // Should return a placeholder JSON since Pi uses TypeScript
        assert!(config.get("_note").is_some());
        assert!(config.get("extension_path").is_some());
    }
}