codecraft/server/src/terminal.rs

//! Web terminal: a shell in a PTY, attached over a WebSocket (xterm.js).
//!
//! Runs in the user's sandbox like the bash tool. Terminals survive a lost
//! connection for a while (mobile apps get suspended) and replay their recent
//! output when a client re-attaches.

use crate::app::App;
use crate::store::User;
use anyhow::{Result, anyhow};
use axum::extract::ws::{Message as WsMessage, WebSocket};
use serde_json::{Value, json};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::broadcast;

/// Output kept for re-attaching clients.
const SCROLLBACK: usize = 256 * 1024;
/// Detached terminals outside a chat are stopped after this time; terminals that
/// belong to a chat stay open until they are closed, exit or the chat is deleted.
const DETACHED_TTL: Duration = Duration::from_secs(30 * 60);
const MAX_PER_USER: usize = 16;

pub struct Term {
    pub id: String,
    pub owner_id: String,
    /// Chat the terminal is a tab of.
    pub session_id: Option<String>,
    pub cwd: PathBuf,
    pub created_at: i64,
    writer: tokio::sync::Mutex<pty_process::OwnedWritePty>,
    child: Mutex<Option<tokio::process::Child>>,
    buffer: Mutex<Vec<u8>>,
    tx: broadcast::Sender<Vec<u8>>,
    attached: Mutex<usize>,
    detached_since: Mutex<Option<Instant>>,
    exited: Mutex<Option<i32>>,
}

impl Term {
    fn alive(&self) -> bool {
        self.exited.lock().unwrap().is_none()
    }

    fn kill(&self) {
        if let Some(mut c) = self.child.lock().unwrap().take() {
            if let Some(pid) = c.id() {
                // The shell's whole process group.
                unsafe {
                    libc::kill(-(pid as i32), libc::SIGHUP);
                }
            }
            let _ = c.start_kill();
        }
    }

    pub fn info(&self) -> Value {
        json!({
            "id": self.id,
            "cwd": self.cwd,
            "session_id": self.session_id,
            "created_at": self.created_at,
            "alive": self.alive(),
            "attached": *self.attached.lock().unwrap() > 0,
        })
    }
}

#[derive(Default)]
pub struct Terminals {
    map: Mutex<HashMap<String, Arc<Term>>>,
}

impl Terminals {
    pub fn list(&self, owner_id: &str) -> Vec<Value> {
        let mut v: Vec<_> = self.map.lock().unwrap().values().filter(|t| t.owner_id == owner_id).map(|t| (t.created_at, t.info())).collect();
        v.sort_by_key(|(c, _)| *c);
        v.into_iter().map(|(_, i)| i).collect()
    }

    /// Terminals of one chat (oldest first).
    pub fn for_session(&self, session_id: &str, owner_id: &str) -> Vec<Value> {
        let mut v: Vec<_> = self
            .map
            .lock()
            .unwrap()
            .values()
            .filter(|t| t.owner_id == owner_id && t.session_id.as_deref() == Some(session_id))
            .map(|t| (t.created_at, t.info()))
            .collect();
        v.sort_by_key(|(c, _)| *c);
        v.into_iter().map(|(_, i)| i).collect()
    }

    /// Close all terminals of a deleted chat.
    pub fn kill_session(&self, session_id: &str) {
        let gone: Vec<Arc<Term>> = {
            let mut map = self.map.lock().unwrap();
            let ids: Vec<String> = map.values().filter(|t| t.session_id.as_deref() == Some(session_id)).map(|t| t.id.clone()).collect();
            ids.iter().filter_map(|id| map.remove(id)).collect()
        };
        gone.iter().for_each(|t| t.kill());
    }

    pub fn get(&self, id: &str, owner_id: &str) -> Option<Arc<Term>> {
        self.map.lock().unwrap().get(id).filter(|t| t.owner_id == owner_id).cloned()
    }

    pub fn kill(&self, id: &str, owner_id: &str) -> bool {
        let t = {
            let mut map = self.map.lock().unwrap();
            match map.get(id) {
                Some(t) if t.owner_id == owner_id => map.remove(id),
                _ => None,
            }
        };
        if let Some(t) = t {
            t.kill();
            true
        } else {
            false
        }
    }

    /// Stop terminals that exited or stayed detached too long.
    pub fn reap(&self) {
        let mut map = self.map.lock().unwrap();
        map.retain(|_, t| {
            let stale = t.session_id.is_none() && t.detached_since.lock().unwrap().is_some_and(|d| d.elapsed() > DETACHED_TTL);
            let gone = !t.alive() && *t.attached.lock().unwrap() == 0;
            if stale || gone {
                t.kill();
                false
            } else {
                true
            }
        });
    }
}

/// Start a shell for `user` in `cwd` (already access-checked).
pub fn spawn(app: &Arc<App>, user: &User, cwd: &Path, session_id: Option<String>, cols: u16, rows: u16) -> Result<Arc<Term>> {
    if app.terminals.list(&user.id).iter().filter(|t| t["alive"] == true).count() >= MAX_PER_USER {
        anyhow::bail!("too many open terminals (max {MAX_PER_USER}) – close one first");
    }
    let (pty, pts) = pty_process::open()?;
    pty.resize(pty_process::Size::new(rows.max(5), cols.max(20)))?;
    // Interactive shell: pager and git prompts work as usual.
    let shell = ["env", "PAGER=less", "GIT_TERMINAL_PROMPT=1", "bash", "-l"];
    let cmd = match crate::sandbox::spec_for(app, Some(&user.id)) {
        Some(spec) => {
            if !crate::sandbox::available() {
                anyhow::bail!("the terminal is disabled for this account: the sandbox (bubblewrap) is not available");
            }
            pty_process::Command::new("bwrap").env_clear().args(crate::sandbox::bwrap_args(&spec, cwd, &[], false, shell[0], &shell[1..]))
        }
        None => {
            let c = pty_process::Command::new(shell[0]).args(&shell[1..]).current_dir(cwd).env("TERM", "xterm-256color").env("CODECRAFT", "1");
            match crate::gitcreds::global_config(app, user) {
                Some(cfg) => c.env("GIT_CONFIG_GLOBAL", cfg),
                None => c,
            }
        }
    };
    let child = cmd.kill_on_drop(true).spawn(pts).map_err(|e| anyhow!("could not start the shell: {e}"))?;
    let (mut reader, writer) = pty.into_split();
    let (tx, _) = broadcast::channel(1024);
    let term = Arc::new(Term {
        id: crate::types::new_id(),
        owner_id: user.id.clone(),
        session_id,
        cwd: cwd.to_path_buf(),
        created_at: crate::types::now_ms(),
        writer: tokio::sync::Mutex::new(writer),
        child: Mutex::new(Some(child)),
        buffer: Mutex::default(),
        tx,
        attached: Mutex::new(0),
        detached_since: Mutex::new(Some(Instant::now())),
        exited: Mutex::new(None),
    });
    app.terminals.map.lock().unwrap().insert(term.id.clone(), term.clone());

    // PTY output -> scrollback + attached clients.
    let t = term.clone();
    tokio::spawn(async move {
        let mut buf = vec![0u8; 16 * 1024];
        loop {
            match reader.read(&mut buf).await {
                Ok(0) | Err(_) => break,
                Ok(n) => {
                    let chunk = buf[..n].to_vec();
                    {
                        let mut sb = t.buffer.lock().unwrap();
                        sb.extend_from_slice(&chunk);
                        if sb.len() > SCROLLBACK {
                            let cut = sb.len() - SCROLLBACK;
                            sb.drain(..cut);
                        }
                    }
                    let _ = t.tx.send(chunk);
                }
            }
        }
        let code = t.child.lock().unwrap().as_mut().and_then(|c| c.try_wait().ok().flatten()).and_then(|s| s.code()).unwrap_or(0);
        *t.exited.lock().unwrap() = Some(code);
        let _ = t.tx.send(format!("\r\n\x1b[2m[process exited with code {code}]\x1b[0m\r\n").into_bytes());
    });
    Ok(term)
}

/// Attach a WebSocket: binary frames are keyboard input, text frames JSON
/// control messages ({"type":"resize","cols":..,"rows":..}).
pub async fn attach(term: Arc<Term>, mut socket: WebSocket) {
    let mut rx = term.tx.subscribe();
    *term.attached.lock().unwrap() += 1;
    *term.detached_since.lock().unwrap() = None;
    let replay = term.buffer.lock().unwrap().clone();
    let hello = json!({"type": "hello", "id": term.id, "cwd": term.cwd, "alive": term.alive()}).to_string();
    let ok = socket.send(WsMessage::Text(hello.into())).await.is_ok() && (replay.is_empty() || socket.send(WsMessage::Binary(replay.into())).await.is_ok());
    if ok {
        loop {
            tokio::select! {
                out = rx.recv() => match out {
                    Ok(bytes) => {
                        if socket.send(WsMessage::Binary(bytes.into())).await.is_err() {
                            break;
                        }
                    }
                    Err(broadcast::error::RecvError::Lagged(_)) => continue,
                    Err(_) => break,
                },
                msg = socket.recv() => match msg {
                    Some(Ok(WsMessage::Binary(b))) => {
                        if term.writer.lock().await.write_all(&b).await.is_err() {
                            break;
                        }
                    }
                    Some(Ok(WsMessage::Text(t))) => {
                        if let Ok(v) = serde_json::from_str::<Value>(&t) {
                            match v["type"].as_str() {
                                Some("resize") => {
                                    let cols = v["cols"].as_u64().unwrap_or(80).clamp(20, 500) as u16;
                                    let rows = v["rows"].as_u64().unwrap_or(24).clamp(5, 300) as u16;
                                    let _ = term.writer.lock().await.resize(pty_process::Size::new(rows, cols));
                                }
                                Some("input") => {
                                    let data = v["data"].as_str().unwrap_or("").as_bytes().to_vec();
                                    if term.writer.lock().await.write_all(&data).await.is_err() {
                                        break;
                                    }
                                }
                                _ => {}
                            }
                        }
                    }
                    Some(Ok(WsMessage::Close(_))) | None | Some(Err(_)) => break,
                    _ => {}
                },
            }
        }
    }
    let mut n = term.attached.lock().unwrap();
    *n = n.saturating_sub(1);
    if *n == 0 {
        *term.detached_since.lock().unwrap() = Some(Instant::now());
    }
}

pub fn spawn_reaper(app: Arc<App>) {
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(Duration::from_secs(60)).await;
            app.terminals.reap();
        }
    });
}