Vectrum
Sign In
410 lines · 13.9 KB · Rust
//! Spawns the long-running agent process behind a thread: the `claude` CLI
//! speaking stream-json (the transport the Claude Agent SDKs use), or Pi in
//! its JSONL RPC mode. Both read commands on stdin and stream events on stdout.

use std::{
    io::{BufRead, BufReader, Write},
    os::unix::process::CommandExt,
    path::{Path, PathBuf},
    process::{Child, Command, Stdio},
    sync::{Arc, Mutex, mpsc},
    thread,
};

use anyhow::{Context as _, Result, anyhow};
use serde_json::{Value, json};

use crate::models::Provider;

/// Environment variables that mark a process as running *inside* another
/// Claude Code session. They leak in when this app is launched from a Claude
/// Code terminal and make the child CLI behave like a nested session.
const NESTED_SESSION_ENV: &[&str] = &[
    "CLAUDECODE",
    "CLAUDE_CODE_ENTRYPOINT",
    "CLAUDE_CODE_SESSION_ID",
    "CLAUDE_CODE_CHILD_SESSION",
    "CLAUDE_CODE_SESSION_ATTENDED",
    "CLAUDE_CODE_EXECPATH",
    "CLAUDE_CODE_MESSAGING_SOCKET",
    "CLAUDE_CODE_MESSAGING_TOKEN",
    "CLAUDE_AGENT_SDK_VERSION",
    "CLAUDE_PID",
];

#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum PermissionMode {
    #[default]
    Default,
    AcceptEdits,
    /// Claude's classifier approves routine actions and asks about the rest.
    Auto,
    Plan,
    BypassPermissions,
}

impl PermissionMode {
    pub const ALL: [PermissionMode; 5] = [
        PermissionMode::Default,
        PermissionMode::AcceptEdits,
        PermissionMode::Auto,
        PermissionMode::Plan,
        PermissionMode::BypassPermissions,
    ];

    pub fn as_cli(self) -> &'static str {
        match self {
            PermissionMode::Default => "default",
            PermissionMode::AcceptEdits => "acceptEdits",
            PermissionMode::Auto => "auto",
            PermissionMode::Plan => "plan",
            PermissionMode::BypassPermissions => "bypassPermissions",
        }
    }

    pub fn from_cli(value: &str) -> Option<Self> {
        Self::ALL.into_iter().find(|mode| mode.as_cli() == value)
    }

    pub fn label(self) -> &'static str {
        match self {
            PermissionMode::Default => "Supervised",
            PermissionMode::AcceptEdits => "Auto-accept edits",
            PermissionMode::Auto => "Auto",
            PermissionMode::Plan => "Plan",
            PermissionMode::BypassPermissions => "Full access",
        }
    }

    pub fn description(self) -> &'static str {
        match self {
            PermissionMode::Default => "Ask before commands and file changes.",
            PermissionMode::AcceptEdits => "Auto-approve edits, ask before other actions.",
            PermissionMode::Auto => "Claude approves routine actions; risky ones still ask.",
            PermissionMode::Plan => "Read-only: Claude plans before it changes anything.",
            PermissionMode::BypassPermissions => "Allow commands and edits without prompts.",
        }
    }
}

/// How hard the model thinks: Claude's `--effort` levels and Pi's thinking
/// levels, which add `off` and `minimal`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Effort {
    Off,
    Minimal,
    Low,
    Medium,
    High,
    XHigh,
    Max,
}

impl Effort {
    pub const ALL: [Effort; 7] =
        [Effort::Off, Effort::Minimal, Effort::Low, Effort::Medium, Effort::High, Effort::XHigh, Effort::Max];

    pub fn as_cli(self) -> &'static str {
        match self {
            Effort::Off => "off",
            Effort::Minimal => "minimal",
            Effort::Low => "low",
            Effort::Medium => "medium",
            Effort::High => "high",
            Effort::XHigh => "xhigh",
            Effort::Max => "max",
        }
    }

    pub fn from_cli(value: &str) -> Option<Self> {
        Self::ALL.into_iter().find(|effort| effort.as_cli() == value)
    }

    pub fn label(self) -> &'static str {
        match self {
            Effort::Off => "Off",
            Effort::Minimal => "Minimal",
            Effort::Low => "Low",
            Effort::Medium => "Medium",
            Effort::High => "High",
            Effort::XHigh => "Extra high",
            Effort::Max => "Max",
        }
    }
}

#[derive(Clone, Debug)]
pub struct LaunchConfig {
    pub provider: Provider,
    pub cwd: PathBuf,
    /// A Claude model alias, or `provider/model` for Pi. `None` is the default.
    pub model: Option<String>,
    pub permission_mode: PermissionMode,
    /// `None` leaves effort at the model's default.
    pub effort: Option<Effort>,
    /// Session to resume (a Claude session id or a Pi session file), used to
    /// restart a thread whose process stopped.
    pub resume: Option<String>,
}

/// Everything the background reader threads report back to the UI.
#[derive(Debug)]
pub enum ProcessEvent {
    Message(Value),
    Stderr(String),
    Exited(Option<i32>),
}

pub struct AgentProcess {
    child: Arc<Mutex<Child>>,
    stdin_tx: mpsc::Sender<String>,
}

impl AgentProcess {
    pub fn spawn(config: &LaunchConfig, events: smol::channel::Sender<ProcessEvent>) -> Result<Self> {
        let (binary, command) = match config.provider {
            Provider::Claude => claude_command(config)?,
            Provider::Pi => pi_command(config)?,
        };
        Self::spawn_command(&binary, command, events)
    }

    fn spawn_command(
        binary: &Path,
        mut command: Command,
        events: smol::channel::Sender<ProcessEvent>,
    ) -> Result<Self> {
        command.stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::piped());
        for key in NESTED_SESSION_ENV {
            command.env_remove(key);
        }
        // Take the agent down with us if the app dies without cleaning up.
        unsafe {
            command.pre_exec(|| {
                libc::prctl(libc::PR_SET_PDEATHSIG, libc::SIGTERM);
                Ok(())
            });
        }

        let mut child = command
            .spawn()
            .with_context(|| format!("failed to start {}", binary.display()))?;
        let stdout = child.stdout.take().context("missing stdout pipe")?;
        let stderr = child.stderr.take().context("missing stderr pipe")?;
        let mut stdin = child.stdin.take().context("missing stdin pipe")?;
        let child = Arc::new(Mutex::new(child));

        // Writes happen on their own thread so the UI never blocks on a full pipe.
        let (stdin_tx, stdin_rx) = mpsc::channel::<String>();
        thread::Builder::new()
            .name("agent-stdin".into())
            .spawn(move || {
                for line in stdin_rx {
                    if stdin.write_all(line.as_bytes()).is_err()
                        || stdin.write_all(b"\n").is_err()
                        || stdin.flush().is_err()
                    {
                        break;
                    }
                }
            })?;

        let stderr_events = events.clone();
        thread::Builder::new()
            .name("agent-stderr".into())
            .spawn(move || {
                for line in BufReader::new(stderr).lines().map_while(Result::ok) {
                    if !line.trim().is_empty() {
                        let _ = stderr_events.send_blocking(ProcessEvent::Stderr(line));
                    }
                }
            })?;

        let reader_child = child.clone();
        thread::Builder::new()
            .name("agent-stdout".into())
            .spawn(move || {
                // `lines()` splits on LF only, as Pi's JSONL framing requires.
                for line in BufReader::new(stdout).lines().map_while(Result::ok) {
                    let line = line.trim();
                    if line.is_empty() {
                        continue;
                    }
                    let event = match serde_json::from_str::<Value>(line) {
                        Ok(value) => ProcessEvent::Message(value),
                        Err(_) => ProcessEvent::Stderr(line.to_string()),
                    };
                    if events.send_blocking(event).is_err() {
                        return;
                    }
                }
                let code = reader_child
                    .lock()
                    .ok()
                    .and_then(|mut child| child.wait().ok())
                    .and_then(|status| status.code());
                let _ = events.send_blocking(ProcessEvent::Exited(code));
            })?;

        Ok(Self { child, stdin_tx })
    }

    pub fn send(&self, message: Value) {
        let _ = self.stdin_tx.send(message.to_string());
    }

    /// Sends a Claude user turn made of API content blocks (text, images).
    pub fn send_user_content(&self, content: Vec<Value>) {
        self.send(json!({
            "type": "user",
            "message": {"role": "user", "content": content},
            "parent_tool_use_id": null,
        }));
    }

    pub fn send_control(&self, request_id: &str, request: Value) {
        self.send(json!({"type": "control_request", "request_id": request_id, "request": request}));
    }

    pub fn respond_control(&self, request_id: &str, response: Value) {
        self.send(json!({
            "type": "control_response",
            "response": {"subtype": "success", "request_id": request_id, "response": response},
        }));
    }

    pub fn kill(&self) {
        if let Ok(mut child) = self.child.lock() {
            let _ = child.kill();
        }
    }
}

impl Drop for AgentProcess {
    fn drop(&mut self) {
        self.kill();
    }
}

fn claude_command(config: &LaunchConfig) -> Result<(PathBuf, Command)> {
    let binary = find_claude_binary()
        .ok_or_else(|| anyhow!("could not find the `claude` CLI on PATH (set VECTRUM_CODE_CLAUDE_BIN)"))?;
    let mut command = Command::new(&binary);
    command
        .args([
            "-p",
            "--input-format",
            "stream-json",
            "--output-format",
            "stream-json",
            "--verbose",
            "--include-partial-messages",
            "--permission-prompt-tool",
            "stdio",
            // Lets the user switch to "Full access" at runtime.
            "--allow-dangerously-skip-permissions",
            "--permission-mode",
            config.permission_mode.as_cli(),
        ])
        .current_dir(&config.cwd);
    if let Some(model) = &config.model {
        command.args(["--model", model]);
    }
    if let Some(effort) = config.effort {
        command.args(["--effort", effort.as_cli()]);
    }
    if let Some(session_id) = &config.resume {
        command.args(["--resume", session_id]);
    }
    Ok((binary, command))
}

fn pi_command(config: &LaunchConfig) -> Result<(PathBuf, Command)> {
    let binary = find_pi_binary()
        .ok_or_else(|| anyhow!("could not find the `pi` CLI on PATH (set VECTRUM_CODE_PI_BIN)"))?;
    let mut command = Command::new(&binary);
    command.args(["--mode", "rpc"]).current_dir(&config.cwd);
    if let Some((provider, model)) = config.model.as_deref().and_then(|model| model.split_once('/')) {
        command.args(["--provider", provider, "--model", model]);
    }
    if let Some(effort) = config.effort {
        command.args(["--thinking", effort.as_cli()]);
    }
    if let Some(session) = &config.resume {
        command.args(["--session", session]);
    }
    // Pi is a Node script; a desktop launcher's PATH may not reach the Node it
    // was installed with, which usually sits next to it.
    if let Some(dir) = binary.parent() {
        let path = std::env::var_os("PATH").unwrap_or_default();
        let paths = std::iter::once(dir.to_path_buf()).chain(std::env::split_paths(&path));
        if let Ok(joined) = std::env::join_paths(paths) {
            command.env("PATH", joined);
        }
    }
    Ok((binary, command))
}

pub fn find_claude_binary() -> Option<PathBuf> {
    if let Some(path) = std::env::var_os("VECTRUM_CODE_CLAUDE_BIN") {
        return Some(PathBuf::from(path));
    }
    if let Some(path) = crate::settings::get(|settings| settings.claude_bin.clone()) {
        return Some(expand_home(&path));
    }
    let on_path = std::env::var_os("PATH")
        .into_iter()
        .flat_map(|paths| std::env::split_paths(&paths).collect::<Vec<_>>())
        .map(|dir| dir.join("claude"));
    // Desktop launchers often start apps without the user's shell PATH.
    let home = std::env::var_os("HOME").map(PathBuf::from).unwrap_or_default();
    let fallbacks = [
        home.join(".local/bin/claude"),
        home.join(".claude/local/claude"),
        PathBuf::from("/usr/local/bin/claude"),
        PathBuf::from("/usr/bin/claude"),
    ];
    on_path.chain(fallbacks).find(|path| is_executable(path))
}

pub fn find_pi_binary() -> Option<PathBuf> {
    if let Some(path) = std::env::var_os("VECTRUM_CODE_PI_BIN") {
        return Some(PathBuf::from(path));
    }
    if let Some(path) = crate::settings::get(|settings| settings.pi_bin.clone()) {
        return Some(expand_home(&path));
    }
    let on_path = std::env::var_os("PATH")
        .into_iter()
        .flat_map(|paths| std::env::split_paths(&paths).collect::<Vec<_>>())
        .map(|dir| dir.join("pi"));
    let home = std::env::var_os("HOME").map(PathBuf::from).unwrap_or_default();
    let mut fallbacks = vec![
        home.join(".local/bin/pi"),
        home.join(".npm-global/bin/pi"),
        home.join(".bun/bin/pi"),
        home.join(".volta/bin/pi"),
        PathBuf::from("/usr/local/bin/pi"),
        PathBuf::from("/usr/bin/pi"),
    ];
    // nvm keeps one bin directory per Node version.
    if let Ok(versions) = std::fs::read_dir(home.join(".nvm/versions/node")) {
        fallbacks.extend(versions.flatten().map(|version| version.path().join("bin/pi")));
    }
    on_path.chain(fallbacks).find(|path| is_executable(path))
}

/// `~/bin/claude` → `/home/me/bin/claude`.
fn expand_home(path: &str) -> PathBuf {
    match (path.strip_prefix("~/"), std::env::var_os("HOME")) {
        (Some(rest), Some(home)) => PathBuf::from(home).join(rest),
        _ => PathBuf::from(path),
    }
}

pub fn is_executable(path: &Path) -> bool {
    use std::os::unix::fs::PermissionsExt;
    path.metadata()
        .map(|meta| meta.is_file() && meta.permissions().mode() & 0o111 != 0)
        .unwrap_or(false)
}