codecraft/server/src/shells.rs

//! Background shells: long-running commands started by agents (dev servers,
//! watchers). Output is buffered; agents read it incrementally.

use crate::app::App;
use crate::tools::{Tool, ToolCtx, ToolOutput, arg_str, truncate};
use anyhow::{Result, bail};
use async_trait::async_trait;
use serde_json::{Value, json};
use std::collections::HashMap;
use std::process::Stdio;
use std::sync::{Arc, Mutex};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::Command;
use tokio_util::sync::CancellationToken;

const MAX_BUFFER: usize = 1024 * 1024;
const MAX_PER_SESSION: usize = 8;

pub struct Shell {
    pub id: String,
    pub session_id: String,
    pub command: String,
    pub started_at: i64,
    /// Everything printed so far (the oldest part is dropped beyond 1 MB).
    output: Mutex<String>,
    /// Bytes dropped from the front of `output`.
    dropped: Mutex<usize>,
    /// Absolute position up to which the agent has read.
    read_pos: Mutex<usize>,
    exit: Mutex<Option<i32>>,
    stop: CancellationToken,
}

impl Shell {
    pub fn running(&self) -> bool {
        self.exit.lock().unwrap().is_none()
    }

    pub fn json(&self, tail: usize) -> Value {
        let out = self.output.lock().unwrap();
        let start = out.len().saturating_sub(tail);
        let start = (start..=out.len()).find(|i| out.is_char_boundary(*i)).unwrap_or(out.len());
        json!({
            "id": self.id, "command": self.command, "started_at": self.started_at,
            "running": self.running(), "exit_code": *self.exit.lock().unwrap(), "output": &out[start..],
        })
    }

    /// Output not yet seen by the agent.
    fn take_new(&self) -> String {
        let out = self.output.lock().unwrap();
        let dropped = *self.dropped.lock().unwrap();
        let mut pos = self.read_pos.lock().unwrap();
        let from = pos.saturating_sub(dropped).min(out.len());
        let from = (from..=out.len()).find(|i| out.is_char_boundary(*i)).unwrap_or(out.len());
        let text = out[from..].to_string();
        *pos = dropped + out.len();
        text
    }

    pub fn kill(&self) {
        self.stop.cancel();
    }
}

pub type Shells = Mutex<HashMap<String, Arc<Shell>>>;

pub fn spawn(app: &Arc<App>, session_id: &str, command: &str, mut cmd: Command) -> Result<String> {
    let running = app.shells.lock().unwrap().values().filter(|s| s.session_id == session_id && s.running()).count();
    if running >= MAX_PER_SESSION {
        bail!("too many background shells in this session ({running}) – stop some with kill_shell");
    }
    cmd.stdin(Stdio::null()).stdout(Stdio::piped()).stderr(Stdio::null()).kill_on_drop(true);
    cmd.process_group(0);
    let mut child = cmd.spawn()?;
    let pid = child.id();
    let shell = Arc::new(Shell {
        id: format!("sh{}", &crate::types::new_id()[20..]),
        session_id: session_id.to_string(),
        command: command.to_string(),
        started_at: crate::types::now_ms(),
        output: Mutex::default(),
        dropped: Mutex::default(),
        read_pos: Mutex::default(),
        exit: Mutex::default(),
        stop: CancellationToken::new(),
    });
    app.shells.lock().unwrap().insert(shell.id.clone(), shell.clone());
    let stdout = child.stdout.take().unwrap();
    let s = shell.clone();
    tokio::spawn(async move {
        let mut lines = BufReader::new(stdout).split(b'\n');
        loop {
            tokio::select! {
                l = lines.next_segment() => match l {
                    Ok(Some(l)) => {
                        let mut out = s.output.lock().unwrap();
                        out.push_str(&String::from_utf8_lossy(&l));
                        out.push('\n');
                        if out.len() > MAX_BUFFER {
                            let mut cut = out.len() - MAX_BUFFER / 2;
                            while !out.is_char_boundary(cut) { cut += 1; }
                            out.drain(..cut);
                            *s.dropped.lock().unwrap() += cut;
                        }
                    }
                    _ => break,
                },
                _ = s.stop.cancelled() => {
                    if let Some(pid) = pid {
                        // SAFETY: plain syscall; negative pid targets the process group.
                        unsafe { libc::kill(-(pid as i32), libc::SIGTERM); }
                        tokio::time::sleep(std::time::Duration::from_secs(2)).await;
                        unsafe { libc::kill(-(pid as i32), libc::SIGKILL); }
                    }
                    break;
                }
            }
        }
        let code = child.wait().await.ok().and_then(|st| st.code()).unwrap_or(-1);
        *s.exit.lock().unwrap() = Some(code);
    });
    Ok(shell.id.clone())
}

/// Stop all shells of a session (on delete).
pub fn kill_session(app: &App, session_id: &str) {
    for s in app.shells.lock().unwrap().values().filter(|s| s.session_id == session_id) {
        s.kill();
    }
}

fn find(ctx: &ToolCtx, id: &str) -> Result<Arc<Shell>> {
    let shells = ctx.app.shells.lock().unwrap();
    match shells.get(id) {
        // Shells belong to the session tree that started them.
        Some(s) if ctx.app.root_session_id(&s.session_id) == ctx.app.root_session_id(&ctx.session_id) => Ok(s.clone()),
        _ => bail!("unknown shell id '{id}'"),
    }
}

pub struct BashOutput;

#[async_trait]
impl Tool for BashOutput {
    fn name(&self) -> &str {
        "bash_output"
    }
    fn description(&self) -> String {
        "Read new output of a background shell (started with bash run_in_background) since the last read, and whether it is still running.".into()
    }
    fn schema(&self) -> Value {
        json!({"type": "object", "properties": {"shell_id": {"type": "string"}, "filter": {"type": "string", "description": "Only lines matching this regex"}}, "required": ["shell_id"]})
    }
    fn read_only(&self) -> bool {
        true
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let shell = find(ctx, arg_str(&input, "shell_id")?)?;
        let mut text = shell.take_new();
        if let Some(f) = input["filter"].as_str() {
            let re = regex::Regex::new(f)?;
            text = text.lines().filter(|l| re.is_match(l)).collect::<Vec<_>>().join("\n");
        }
        let status = match *shell.exit.lock().unwrap() {
            None => "running".to_string(),
            Some(c) => format!("exited with code {c}"),
        };
        Ok(ToolOutput::text(format!("[{status}]\n{}", if text.is_empty() { "(no new output)".into() } else { truncate(&text, 30_000) })))
    }
}

pub struct KillShell;

#[async_trait]
impl Tool for KillShell {
    fn name(&self) -> &str {
        "kill_shell"
    }
    fn description(&self) -> String {
        "Stop a background shell.".into()
    }
    fn schema(&self) -> Value {
        json!({"type": "object", "properties": {"shell_id": {"type": "string"}}, "required": ["shell_id"]})
    }
    fn read_only(&self) -> bool {
        true
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let shell = find(ctx, arg_str(&input, "shell_id")?)?;
        shell.kill();
        Ok(ToolOutput::text(format!("Stopping shell {}.", shell.id)))
    }
}