clawmates/crates/cm-runtime/src/terminals.rs

//! Interactive terminal containers — one themed `agent-terminal` (zsh +
//! oh-my-zsh + powerlevel10k) container per agent, provisioned lazily on first
//! WebSocket connect and reused across reconnects. Unlike the agent tool
//! sandboxes, these are interactive PTYs (`exec -it zsh`), so a small idle
//! sweeper reaps containers that no live session has touched for a while.
//!
//! The agent→container mapping lives in Postgres (`agent_containers`) rather
//! than process memory, so ANY server replica reuses the same container instead
//! of spawning a duplicate — and registry-tracked terminals survive a redeploy
//! (their tmux sessions resume), since boot reconciliation removes only true
//! orphans (engine containers with no registry row).

use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use cm_domain::{AgentId, WorkspaceId};
use cm_sandbox::{DriveMount, PtySession, SandboxDriver, SandboxHandle, SandboxKind, SandboxSpec};
use sqlx::PgPool;

use crate::NodeDriverProvider;

/// The `agent_containers.kind` discriminator for terminal containers.
const KIND: &str = "terminal";

/// The per-agent node-local drive volume mounted at `~/drives` for a node-placed
/// agent — auto-created by Docker on first mount (no subpath). All of the agent's
/// containers on that node mount the SAME volume, so they share files. Both the
/// terminal and the sandbox use this exact mount (see `SandboxManager`).
pub(crate) fn node_local_drive_mount(agent_id: AgentId) -> DriveMount {
    DriveMount {
        volume: format!("clawmates_agent_{}", agent_id.as_uuid().simple()),
        subpath: String::new(),
        target: "/home/agent/drives".to_string(),
        read_only: false,
    }
}

/// Where the agent's Files drives live, so the Terminal can mount them.
#[derive(Clone)]
pub struct DriveConfig {
    /// The engine's named volume holding all file-drive blobs (e.g. `clawmates_filedata`).
    pub volume: String,
    /// The server's view of the blob root, to create per-agent subdirs before
    /// the (subpath) mount — Docker errors on a missing subpath.
    pub data_dir: PathBuf,
}

pub struct TerminalManager {
    driver: Arc<dyn SandboxDriver>,
    pool: PgPool,
    /// This manager's LOCAL placement id ("local"); remote placements come from
    /// `node_provider` keyed by a fleet node's id.
    node_id: String,
    image: String,
    egress: bool,
    /// Drive mounts; None disables the ~/drives mapping.
    drives: Option<DriveConfig>,
    /// Resolves drivers for remote fleet nodes (None ⇒ local-only deployment), so
    /// an agent placed on a node runs its terminal there too.
    node_provider: Option<Arc<dyn NodeDriverProvider>>,
}

impl std::fmt::Debug for TerminalManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TerminalManager")
            .field("image", &self.image)
            .field("node_id", &self.node_id)
            .finish_non_exhaustive()
    }
}

impl TerminalManager {
    pub fn new(
        driver: Arc<dyn SandboxDriver>,
        pool: PgPool,
        node_id: &str,
        image: &str,
        egress: bool,
        drives: Option<DriveConfig>,
    ) -> TerminalManager {
        TerminalManager {
            driver,
            pool,
            node_id: node_id.to_owned(),
            image: image.to_owned(),
            egress,
            drives,
            node_provider: None,
        }
    }

    /// Wire a fleet-node driver provider so an agent placed on a connected node
    /// runs its terminal container there (sharing the agent's node-local drives).
    /// Default (unset) = local-only, unchanged.
    pub fn with_node_provider(mut self, provider: Arc<dyn NodeDriverProvider>) -> TerminalManager {
        self.node_provider = Some(provider);
        self
    }

    /// The driver that owns containers on `node_id`: the local driver for "local"
    /// (or when the node isn't connected), else the node's remote driver.
    fn driver_for(&self, node_id: &str) -> Arc<dyn SandboxDriver> {
        if node_id != self.node_id {
            if let Some(d) = self.node_provider.as_ref().and_then(|p| p.driver(node_id)) {
                return d;
            }
        }
        self.driver.clone()
    }

    /// Where a NEW terminal container for this agent should run: the agent's
    /// workspace placement node if it's connected (and not draining), else "local"
    /// — mirroring `SandboxManager` so the terminal lands beside the sandbox.
    async fn placement_node(&self, agent_id: AgentId) -> String {
        match cm_db::repo::workspace_placement::for_agent(&self.pool, agent_id).await {
            Ok(Some(node)) if node != self.node_id => {
                if let Ok(nid) = node.parse::<uuid::Uuid>() {
                    if matches!(
                        cm_db::repo::nodes::status_of(&self.pool, cm_domain::NodeId::from(nid)).await,
                        Ok(Some(ref s)) if s == "draining"
                    ) {
                        return self.node_id.clone();
                    }
                }
                if self
                    .node_provider
                    .as_ref()
                    .and_then(|p| p.driver(&node))
                    .is_some()
                {
                    return node;
                }
                self.node_id.clone()
            }
            _ => self.node_id.clone(),
        }
    }

    /// Where the agent's terminal is / would run, WITHOUT provisioning: the node
    /// of an existing container, else the computed placement. For the ticket's
    /// `node` hint (the browser uses the WebRTC-vs-WS transport accordingly).
    pub async fn placement_node_for(&self, agent_id: AgentId) -> String {
        if let Ok(Some(row)) = cm_db::repo::agent_containers::get(&self.pool, agent_id, KIND).await
        {
            return row.node_id;
        }
        self.placement_node(agent_id).await
    }

    /// Ensure the agent's terminal container exists (on its placement node) and
    /// return `(node_id, container_name)` — what the node-bridge route needs to
    /// `docker exec` into it. `node_id` is "local" or a fleet-node UUID.
    pub async fn placement_for(
        &self,
        workspace_id: WorkspaceId,
        agent_id: AgentId,
    ) -> Result<(String, String), String> {
        let (handle, node_id) = self.ensure(workspace_id, agent_id).await?;
        Ok((node_id, handle.name))
    }

    /// The three Files drives mounted read-write at `~/drives/*`, each a per-agent
    /// subpath of the shared volume (subpath = isolation). Creates the subdirs
    /// first so the mount doesn't fail on an empty drive.
    async fn drive_mounts(&self, workspace_id: WorkspaceId, agent_id: AgentId) -> Vec<DriveMount> {
        let Some(cfg) = &self.drives else {
            return Vec::new();
        };
        let ws = workspace_id.to_string();
        let agent = agent_id.to_string();
        // (drive, scope, mount target) per `files::blob_key`: documents/received +
        // the Obsidian vault are agent-scoped; the shared ClawDrive is team-wide.
        // The vault lands at ~/obsidian (the agent's second-brain markdown vault).
        let drives: [(&str, String, &str); 4] = [
            ("documents", agent.clone(), "/home/agent/drives/documents"),
            ("received", agent.clone(), "/home/agent/drives/received"),
            ("shared", "shared".to_string(), "/home/agent/drives/shared"),
            ("vault", agent.clone(), "/home/agent/obsidian"),
        ];
        let mut out = Vec::with_capacity(drives.len());
        for (drive, scope, target) in drives {
            let subpath = format!("{ws}/{drive}/{scope}");
            let _ = tokio::fs::create_dir_all(cfg.data_dir.join(&subpath)).await;
            out.push(DriveMount {
                volume: cfg.volume.clone(),
                subpath,
                target: target.to_string(),
                read_only: false,
            });
        }
        out
    }

    async fn provision_one(
        &self,
        workspace_id: WorkspaceId,
        agent_id: AgentId,
        node: &str,
    ) -> Result<SandboxHandle, String> {
        let short = uuid::Uuid::now_v7().simple().to_string();
        // Local: per-agent subpaths of the shared gateway volume (today's path).
        // Remote: a single per-agent node-local volume at ~/drives, shared with
        // the agent's sandbox on that node.
        let mounts = if node == self.node_id {
            self.drive_mounts(workspace_id, agent_id).await
        } else {
            vec![node_local_drive_mount(agent_id)]
        };
        let spec = SandboxSpec {
            name: format!("tc-term-{}", &short[short.len() - 12..]),
            image: self.image.clone(),
            // A dev shell wants more headroom than a tool sandbox.
            memory_bytes: 1024 * 1024 * 1024,
            nano_cpus: 2_000_000_000,
            pids_limit: 512,
            egress: self.egress,
            kind: SandboxKind::Terminal,
            mounts,
        };
        self.driver_for(node)
            .provision(&spec)
            .await
            .map_err(|e| format!("terminal provision failed: {e}"))
    }

    /// The agent's terminal container, provisioned on first use; a dead one is
    /// replaced transparently. The mapping is read from / written to the shared
    /// registry, so a different replica finds the same container.
    async fn ensure(
        &self,
        workspace_id: WorkspaceId,
        agent_id: AgentId,
    ) -> Result<(SandboxHandle, String), String> {
        if let Ok(Some(row)) = cm_db::repo::agent_containers::get(&self.pool, agent_id, KIND).await
        {
            let driver = self.driver_for(&row.node_id);
            let handle = SandboxHandle {
                id: row.container_id.clone(),
                name: row.name.clone(),
            };
            if driver.health(&handle).await.unwrap_or(false) {
                return Ok((handle, row.node_id));
            }
            // Recorded but dead: clean both the container and the stale row.
            let _ = driver.destroy(&handle).await;
            let _ = cm_db::repo::agent_containers::delete(&self.pool, agent_id, KIND).await;
        }
        let node = self.placement_node(agent_id).await;
        let handle = self.provision_one(workspace_id, agent_id, &node).await?;
        cm_db::repo::agent_containers::upsert(
            &self.pool,
            agent_id,
            KIND,
            &node,
            &handle.id,
            &handle.name,
        )
        .await
        .map_err(|e| format!("registry upsert failed: {e}"))?;
        Ok((handle, node))
    }

    /// Open an interactive login zsh in the agent's terminal container. `env`
    /// adds session vars (e.g. a `CLAWMATES_USER` MOTD greeting). Counts a live
    /// session (paired with [`detach`]).
    pub async fn attach(
        &self,
        workspace_id: WorkspaceId,
        agent_id: AgentId,
        cols: u16,
        rows: u16,
        env: &[String],
        tmux_session: &str,
    ) -> Result<PtySession, String> {
        let (handle, node_id) = self.ensure(workspace_id, agent_id).await?;
        // tmux attach-or-create: the named session + its processes survive a WS
        // disconnect and resume on reconnect. Different session names are
        // independent tabs sharing the agent's one container (+ its ~/drives).
        // Local attach uses the local driver (unchanged); a remote terminal is
        // served via the node WebRTC/WS bridge, not here.
        let session = self
            .driver_for(&node_id)
            .attach_pty(
                &handle,
                // -A attach-or-create; -D detaches any stale client on reattach so
                // the resumed session redraws cleanly at the new client's size.
                &["tmux", "new-session", "-A", "-D", "-s", tmux_session],
                cols,
                rows,
                env,
            )
            .await
            .map_err(|e| format!("terminal attach failed: {e}"))?;
        let _ = cm_db::repo::agent_containers::add_session(&self.pool, agent_id, KIND, 1).await;
        Ok(session)
    }

    pub async fn resize(&self, exec_id: &str, cols: u16, rows: u16) -> Result<(), String> {
        self.driver
            .resize_pty(exec_id, cols, rows)
            .await
            .map_err(|e| format!("terminal resize failed: {e}"))
    }

    /// A session ended; the container stays warm for reconnects but becomes a
    /// candidate for the idle sweeper once no session remains.
    pub async fn detach(&self, agent_id: AgentId) {
        let _ = cm_db::repo::agent_containers::add_session(&self.pool, agent_id, KIND, -1).await;
    }

    /// Tear down a single agent's terminal container (on agent deletion / idle).
    /// Returns whether one existed.
    pub async fn release_agent(&self, agent_id: AgentId) -> bool {
        match cm_db::repo::agent_containers::get(&self.pool, agent_id, KIND).await {
            Ok(Some(row)) => {
                let handle = SandboxHandle {
                    id: row.container_id,
                    name: row.name,
                };
                if let Err(e) = self.driver.destroy(&handle).await {
                    eprintln!("terminal release: failed to remove {}: {e}", handle.id);
                }
                let _ = cm_db::repo::agent_containers::delete(&self.pool, agent_id, KIND).await;
                true
            }
            _ => false,
        }
    }

    /// Graceful shutdown: terminals are NOT destroyed — they persist (tmux
    /// sessions survive the redeploy) and become idle-reapable. We only zero the
    /// session counts since the WS connections are dropping.
    pub async fn shutdown(&self) {
        let _ = cm_db::repo::agent_containers::reset_sessions(&self.pool, KIND).await;
    }

    /// Reap idle terminals (no live session, untouched past `idle_ttl`).
    async fn reap_idle(&self, idle_ttl: Duration) -> usize {
        let rows =
            match cm_db::repo::agent_containers::idle(&self.pool, KIND, idle_ttl.as_secs() as i64)
                .await
            {
                Ok(r) => r,
                Err(e) => {
                    eprintln!("terminal reaper: idle query failed: {e}");
                    return 0;
                }
            };
        let mut reaped = 0;
        for r in rows {
            let handle = SandboxHandle {
                id: r.container_id.clone(),
                name: r.name.clone(),
            };
            let _ = self.driver.destroy(&handle).await;
            let _ = cm_db::repo::agent_containers::delete(&self.pool, r.agent_id, KIND).await;
            reaped += 1;
        }
        reaped
    }

    /// Reap orphan containers the engine holds that no registry row owns (and,
    /// when `min_age` > 0, that are older than it). Registry-tracked terminals
    /// are left alone.
    async fn reap_orphans(&self, min_age: Duration) -> usize {
        let managed = match self
            .driver
            .list_managed(SandboxKind::Terminal.label())
            .await
        {
            Ok(m) => m,
            Err(e) => {
                eprintln!("terminal reaper: list failed: {e}");
                return 0;
            }
        };
        let tracked: std::collections::HashSet<String> =
            cm_db::repo::agent_containers::all(&self.pool, KIND)
                .await
                .unwrap_or_default()
                .into_iter()
                .map(|r| r.container_id)
                .collect();
        let now_unix = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_secs() as i64)
            .unwrap_or(0);
        let min = min_age.as_secs() as i64;
        let mut reaped = 0;
        for m in managed {
            if tracked.contains(&m.id) {
                continue;
            }
            if min > 0 && now_unix - m.created_unix < min {
                continue;
            }
            let handle = SandboxHandle {
                id: m.id.clone(),
                name: m.id.clone(),
            };
            match self.driver.destroy(&handle).await {
                Ok(()) => reaped += 1,
                Err(e) => eprintln!("terminal reaper: failed to remove {}: {e}", m.id),
            }
        }

        // An agent placed on a fleet node runs its terminal there too, so sweep
        // each connected node as well — otherwise a node-placed terminal whose
        // registry row is gone can never be found again. TTL-only for the same
        // reason as the sandbox reaper: the boot pass uses ZERO and must not
        // touch containers on a shared node.
        if min > 0 {
            for node_id in self
                .node_provider
                .as_ref()
                .map(|p| p.node_ids())
                .unwrap_or_default()
            {
                if node_id == self.node_id {
                    continue;
                }
                let Some(driver) = self.node_provider.as_ref().and_then(|p| p.driver(&node_id))
                else {
                    continue;
                };
                let remote = match driver.list_managed(SandboxKind::Terminal.label()).await {
                    Ok(m) => m,
                    Err(e) => {
                        eprintln!("terminal reaper: list on node {node_id} failed: {e}");
                        continue;
                    }
                };
                for m in remote {
                    if tracked.contains(&m.id) || now_unix - m.created_unix < min {
                        continue;
                    }
                    let handle = SandboxHandle {
                        id: m.id.clone(),
                        name: m.id.clone(),
                    };
                    match driver.destroy(&handle).await {
                        Ok(()) => reaped += 1,
                        Err(e) => eprintln!(
                            "terminal reaper: failed to remove {} on node {node_id}: {e}",
                            m.id
                        ),
                    }
                }
            }
        }
        reaped
    }

    /// Boot reconciliation: remove only true orphans (engine containers with no
    /// registry row). Registry-tracked terminals are preserved across a redeploy.
    pub async fn reconcile_orphans(&self) -> usize {
        self.reap_orphans(Duration::ZERO).await
    }

    /// Background reaper: every `interval`, reap idle (`idle_ttl`) + orphan terminals.
    pub fn spawn_reaper(self: Arc<Self>, interval: Duration, idle_ttl: Duration) {
        tokio::spawn(async move {
            let mut tick = tokio::time::interval(interval);
            loop {
                tick.tick().await;
                let n = self.reap_idle(idle_ttl).await + self.reap_orphans(idle_ttl).await;
                if n > 0 {
                    eprintln!("terminal reaper: removed {n} terminal container(s)");
                }
            }
        });
    }
}