IdeaSDK/crates/infrastructure/src/pty/mod.rs

//! [`PortablePtyAdapter`] — local [`PtyPort`] implementation over the
//! `portable-pty` crate (ARCHITECTURE §5, L3).
//!
//! # Design
//!
//! `portable-pty` is a blocking, thread-oriented API: a master PTY gives a
//! `Box<dyn Read>` reader and a `Box<dyn Write>` writer, and the child process is
//! waited on a thread. We bridge that to the domain [`PtyPort`] as follows:
//!
//! - [`PtyHandle`] only carries a [`SessionId`]; the *real* OS handles (master
//!   PTY, writer, child) live in this adapter's registry keyed by that id. The
//!   domain never sees an OS handle (ARCHITECTURE §4).
//! - On [`spawn`](PtyPort::spawn) we open a PTY pair, spawn the command in the
//!   slave, then start **one reader thread** that pumps bytes from the master
//!   into a shared [`Broadcast`] hub. The hub does two things with every chunk:
//!   it appends to a bounded **scrollback ring buffer** (~100 KB, most recent
//!   bytes) and it fans the chunk out to every *currently subscribed* receiver.
//! - [`subscribe_output`](PtyPort::subscribe_output) registers a fresh
//!   subscriber and returns its receiver wrapped as the domain's blocking
//!   [`OutputStream`] iterator. It is **re-subscribable**: after a view tears
//!   down (navigation / layout change) a new view can re-attach to the *same*
//!   live PTY by subscribing again and repainting the scrollback first — no
//!   re-spawn. [`scrollback`](PtyPort::scrollback) returns that retained buffer.
//! - [`write`](PtyPort::write) / [`resize`](PtyPort::resize) act on the stored
//!   writer / master. [`kill`](PtyPort::kill) terminates the child, joins the
//!   reader thread, and returns the [`ExitStatus`].
//!
//! # Cross-platform note (spike, ARCHITECTURE §13.1)
//!
//! `portable-pty` abstracts ConPTY on Windows, but exit-code/signal semantics
//! differ: a Unix process killed by a signal reports `code: None` here, while
//! ConPTY surfaces a numeric code. Resize uses the same `PtySize` everywhere.
//! Points to validate on Windows: child `kill()` actually tears down ConPTY, and
//! the reader thread observes EOF promptly on exit. The code below avoids any
//! Unix-only assumption (no raw fds, no signals) so it should port as-is.

use std::collections::HashMap;
use std::collections::VecDeque;
use std::io::{Read, Write};
use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;

use async_trait::async_trait;
use portable_pty::{Child, CommandBuilder, MasterPty, NativePtySystem, PtySystem, SlavePty};

use domain::ports::{ExitStatus, OutputStream, PtyError, PtyHandle, PtyPort, SpawnSpec};
use domain::sandbox::{SandboxEnforcer, SandboxPlan};
use domain::terminal::PtySize;
use domain::SessionId;

/// Size of each read buffer pumped from the master PTY.
const READ_BUF: usize = 8 * 1024;

/// Maximum number of bytes retained in a session's scrollback ring buffer
/// (~100 KB, "the most recent output"). When the buffer would exceed this, the
/// oldest bytes are dropped so a re-attaching view repaints recent history.
const SCROLLBACK_CAP: usize = 100 * 1024;

/// The shared output hub of one PTY: a bounded scrollback ring buffer plus the
/// **single** currently-subscribed receiver. The reader thread feeds both; a
/// view subscribes on (re-)attach, and each new subscription supersedes the
/// previous one ([`Broadcast::subscribe`]) so a PTY is never fanned to two live
/// consumers at once.
///
/// A subscriber is a [`Sender`]; a send failing (receiver dropped because the
/// view detached) prunes it on the next chunk. The `Vec` is retained (rather
/// than an `Option`) only to keep the prune-on-send logic uniform.
#[derive(Default)]
struct Broadcast {
    /// Bounded ring buffer of the most recent output bytes.
    scrollback: VecDeque<u8>,
    /// Live subscribers; pruned lazily when their receiver is gone.
    subscribers: Vec<Sender<Vec<u8>>>,
    /// Set once the PTY hit EOF (process exited) — no more output will ever come.
    eof: bool,
}

impl Broadcast {
    /// Appends a chunk to the scrollback (trimming to [`SCROLLBACK_CAP`]) and
    /// fans it out to every live subscriber, dropping any that have gone away.
    fn push(&mut self, chunk: &[u8]) {
        self.scrollback.extend(chunk.iter().copied());
        let overflow = self.scrollback.len().saturating_sub(SCROLLBACK_CAP);
        if overflow > 0 {
            self.scrollback.drain(0..overflow);
        }
        self.subscribers
            .retain(|tx| tx.send(chunk.to_vec()).is_ok());
    }

    /// Registers a new subscriber, returning its receiver.
    ///
    /// **Single live consumer.** A PTY session is shown by exactly one view at a
    /// time; a re-attach supersedes the previous view. So a new subscription
    /// **drops every prior subscriber**, which ends the previous attach's pump
    /// thread (its receiver closes) instead of leaving it alive to fan the *same*
    /// bytes a second time — the root cause of on-resize output duplication.
    ///
    /// If the PTY already hit EOF, the returned stream is immediately closed
    /// (its sender is dropped) so a late re-attach to a finished session doesn't
    /// block forever waiting for output that will never come.
    fn subscribe(&mut self) -> Receiver<Vec<u8>> {
        let (tx, rx) = mpsc::channel();
        self.subscribers.clear();
        if !self.eof {
            self.subscribers.push(tx);
        }
        rx
    }

    /// Returns the retained scrollback as a contiguous byte vector.
    fn snapshot(&self) -> Vec<u8> {
        self.scrollback.iter().copied().collect()
    }

    /// Drops every subscriber's sender so their output streams end (EOF). Called
    /// by the reader thread when the PTY hits EOF (process exit). The scrollback
    /// is preserved so a late re-attach can still repaint the final output.
    fn close_subscribers(&mut self) {
        self.eof = true;
        self.subscribers.clear();
    }
}

/// A live PTY owned by the adapter.
struct LivePty {
    /// Master side — used for resize.
    master: Box<dyn MasterPty + Send>,
    /// Writer into the PTY (child stdin).
    writer: Box<dyn Write + Send>,
    /// The spawned child process.
    child: Box<dyn Child + Send + Sync>,
    /// Shared scrollback + subscriber hub, fed by the reader thread.
    output: Arc<Mutex<Broadcast>>,
    /// Handle of the reader thread, joined on kill.
    reader: Option<JoinHandle<()>>,
}

/// Local PTY adapter backed by `portable-pty`'s native PTY system.
///
/// Thread-safe: the registry of live PTYs is behind a [`Mutex`]; the adapter is
/// cloneable-as-`Arc` and injected as `Arc<dyn PtyPort>` at the composition root.
#[derive(Default)]
pub struct PortablePtyAdapter {
    sessions: Mutex<HashMap<SessionId, LivePty>>,
    /// Optional OS-sandbox enforcer (lot LP4-1). `None` ⇒ no sandboxing at all
    /// (the default, zero behaviour change). When set **and** a [`SpawnSpec`]
    /// carries a [`SandboxPlan`], the plan is enforced on the child at spawn
    /// (see [`spawn_command_sandboxed`]).
    sandbox_enforcer: Option<Arc<dyn SandboxEnforcer>>,
}

impl PortablePtyAdapter {
    /// Creates an empty adapter (no OS sandbox).
    #[must_use]
    pub fn new() -> Self {
        Self {
            sessions: Mutex::new(HashMap::new()),
            sandbox_enforcer: None,
        }
    }

    /// Additive builder: wires an OS-sandbox [`SandboxEnforcer`] (lot LP4-1). With
    /// it set, any [`SpawnSpec`] whose `sandbox` is `Some` has that plan enforced on
    /// the spawned child. Without it (the default), no spawn is ever sandboxed —
    /// `new()`'s behaviour is unchanged.
    #[must_use]
    pub fn with_sandbox_enforcer(mut self, enforcer: Arc<dyn SandboxEnforcer>) -> Self {
        self.sandbox_enforcer = Some(enforcer);
        self
    }
}

/// Maps the domain [`PtySize`] to the `portable-pty` one.
fn to_pty_size(size: PtySize) -> portable_pty::PtySize {
    portable_pty::PtySize {
        rows: size.rows,
        cols: size.cols,
        pixel_width: 0,
        pixel_height: 0,
    }
}

/// Spawns the child **under an OS sandbox** (lot LP4-1).
///
/// ## Why a dedicated thread instead of a `pre_exec` hook
///
/// `portable-pty` 0.9 exposes **no** user-injectable `pre_exec` closure: its
/// `CommandBuilder` has no such method, and `SlavePty::spawn_command` installs its
/// *own* internal `pre_exec` (for `setsid` / controlling-tty) before exec, with no
/// extension point. So we cannot run `enforcer.enforce(plan)` in the child's
/// `pre_exec` directly.
///
/// Instead we lean on a guaranteed kernel property: **a Landlock domain is
/// inherited across `fork` and preserved across `execve`**. We therefore restrict a
/// **dedicated throwaway thread** with `enforce`, then call `spawn_command` *from
/// that thread*. `std::process::Command` (which `spawn_command` drives) forks from
/// the calling thread — and because `portable-pty` always sets a `pre_exec`, std is
/// forced down the `fork`+`exec` path (never `posix_spawn`) — so the child inherits
/// the restricted thread's domain. The throwaway thread then dies, taking its
/// (irreversible) restriction with it, leaving IdeA's other threads untouched.
///
/// On `enforce` returning `Err` (fail-closed, e.g. a `Deny` posture on a kernel
/// without Landlock) the spawn is failed and **no child runs**. `Ok(Unsupported /
/// Degraded)` proceeds (best-effort).
///
/// Everything is **moved** into the thread (no borrow), so no `Sync` bound is
/// required of the slave; the resulting `Child` is `Send` and returned to the
/// caller.
fn spawn_command_sandboxed(
    slave: Box<dyn SlavePty + Send>,
    cmd: CommandBuilder,
    enforcer: Arc<dyn SandboxEnforcer>,
    plan: SandboxPlan,
) -> Result<Box<dyn Child + Send + Sync>, PtyError> {
    let join = std::thread::spawn(move || -> Result<Box<dyn Child + Send + Sync>, PtyError> {
        // Restrict THIS thread; the forked child inherits the domain.
        if let Err(e) = enforcer.enforce(&plan) {
            return Err(PtyError::Spawn(format!("sandbox enforcement failed: {e}")));
        }
        let child = slave
            .spawn_command(cmd)
            .map_err(|e| PtyError::Spawn(e.to_string()))?;
        // The slave is held by the child; drop our copy so EOF propagates on exit.
        drop(slave);
        Ok(child)
    });
    join.join()
        .map_err(|_| PtyError::Spawn("sandbox spawn thread panicked".to_owned()))?
}

/// Builds the `portable-pty` command from a domain [`SpawnSpec`].
fn to_command(spec: &SpawnSpec) -> CommandBuilder {
    let mut cmd = CommandBuilder::new(&spec.command);
    cmd.args(&spec.args);
    cmd.cwd(spec.cwd.as_str());
    for (k, v) in &spec.env {
        cmd.env(k, v);
    }
    cmd
}

#[async_trait]
impl PtyPort for PortablePtyAdapter {
    async fn spawn(&self, spec: SpawnSpec, size: PtySize) -> Result<PtyHandle, PtyError> {
        let pty_system = NativePtySystem::default();
        let pair = pty_system
            .openpty(to_pty_size(size))
            .map_err(|e| PtyError::Spawn(e.to_string()))?;

        let cmd = to_command(&spec);
        // Move the slave out of the pair (the master stays usable) so it can be
        // either spawned directly or handed to the sandboxed-spawn thread.
        let slave = pair.slave;
        let child = match (spec.sandbox.clone(), self.sandbox_enforcer.clone()) {
            // Sandboxed spawn: a plan AND an enforcer are present.
            (Some(plan), Some(enforcer)) => spawn_command_sandboxed(slave, cmd, enforcer, plan)?,
            // No plan or no enforcer ⇒ plain spawn (the default path, unchanged).
            _ => {
                let child = slave
                    .spawn_command(cmd)
                    .map_err(|e| PtyError::Spawn(e.to_string()))?;
                // The slave is held by the child; drop our copy so EOF propagates.
                drop(slave);
                child
            }
        };

        let writer = pair
            .master
            .take_writer()
            .map_err(|e| PtyError::Io(e.to_string()))?;
        let mut reader = pair
            .master
            .try_clone_reader()
            .map_err(|e| PtyError::Io(e.to_string()))?;

        // One reader thread per PTY pumps bytes into the broadcast hub until EOF.
        // The hub retains a scrollback ring buffer AND fans bytes out to every
        // current subscriber, so views can detach/re-attach without re-spawning.
        let output: Arc<Mutex<Broadcast>> = Arc::new(Mutex::new(Broadcast::default()));
        let output_for_reader = Arc::clone(&output);
        let reader_handle = std::thread::spawn(move || {
            let mut buf = [0u8; READ_BUF];
            loop {
                match reader.read(&mut buf) {
                    Ok(0) => break,
                    Ok(n) => {
                        if let Ok(mut hub) = output_for_reader.lock() {
                            hub.push(&buf[..n]);
                        }
                    }
                    Err(_) => break,
                }
            }
            // EOF (process exited): end every attached stream by dropping its
            // sender, while preserving the scrollback for any late re-attach.
            if let Ok(mut hub) = output_for_reader.lock() {
                hub.close_subscribers();
            }
        });

        // The PTY layer owns the handle identity: it mints a fresh session id and
        // the caller (the `OpenTerminal` use case) adopts it as the
        // `TerminalSession.id`. This keeps the OS handle out of the domain while
        // giving everyone a single, agreed-upon id (ARCHITECTURE §4).
        let handle = PtyHandle {
            session_id: SessionId::new_random(),
        };
        let live = LivePty {
            master: pair.master,
            writer,
            child,
            output,
            reader: Some(reader_handle),
        };
        self.sessions
            .lock()
            .map_err(|_| PtyError::Io("pty registry poisoned".to_owned()))?
            .insert(handle.session_id, live);
        Ok(handle)
    }

    fn write(&self, handle: &PtyHandle, data: &[u8]) -> Result<(), PtyError> {
        let mut map = self
            .sessions
            .lock()
            .map_err(|_| PtyError::Io("pty registry poisoned".to_owned()))?;
        let live = map.get_mut(&handle.session_id).ok_or(PtyError::NotFound)?;
        live.writer
            .write_all(data)
            .map_err(|e| PtyError::Io(e.to_string()))?;
        live.writer.flush().map_err(|e| PtyError::Io(e.to_string()))
    }

    fn resize(&self, handle: &PtyHandle, size: PtySize) -> Result<(), PtyError> {
        let map = self
            .sessions
            .lock()
            .map_err(|_| PtyError::Io("pty registry poisoned".to_owned()))?;
        let live = map.get(&handle.session_id).ok_or(PtyError::NotFound)?;
        live.master
            .resize(to_pty_size(size))
            .map_err(|e| PtyError::Io(e.to_string()))
    }

    fn subscribe_output(&self, handle: &PtyHandle) -> Result<OutputStream, PtyError> {
        let map = self
            .sessions
            .lock()
            .map_err(|_| PtyError::Io("pty registry poisoned".to_owned()))?;
        let live = map.get(&handle.session_id).ok_or(PtyError::NotFound)?;
        let rx = live
            .output
            .lock()
            .map_err(|_| PtyError::Io("pty output hub poisoned".to_owned()))?
            .subscribe();
        Ok(Box::new(rx.into_iter()))
    }

    fn scrollback(&self, handle: &PtyHandle) -> Result<Vec<u8>, PtyError> {
        let map = self
            .sessions
            .lock()
            .map_err(|_| PtyError::Io("pty registry poisoned".to_owned()))?;
        let live = map.get(&handle.session_id).ok_or(PtyError::NotFound)?;
        let snapshot = live
            .output
            .lock()
            .map_err(|_| PtyError::Io("pty output hub poisoned".to_owned()))?
            .snapshot();
        Ok(snapshot)
    }

    async fn kill(&self, handle: &PtyHandle) -> Result<ExitStatus, PtyError> {
        // Remove from the registry so the writer/master drop and the child is
        // fully owned here while we tear it down.
        let mut live = {
            let mut map = self
                .sessions
                .lock()
                .map_err(|_| PtyError::Io("pty registry poisoned".to_owned()))?;
            map.remove(&handle.session_id).ok_or(PtyError::NotFound)?
        };

        // Ask the child to terminate, then wait for its real status.
        let _ = live.child.kill();
        let status = live.child.wait().map_err(|e| PtyError::Io(e.to_string()))?;

        // Dropping master/writer closes the PTY; the reader thread then sees EOF.
        // Dropping the broadcast hub drops every subscriber's sender, so any
        // still-attached view's output stream ends cleanly too.
        if let Some(reader) = live.reader.take() {
            let _ = reader.join();
        }

        Ok(ExitStatus {
            code: exit_code(&status),
        })
    }
}

/// Extracts a portable exit code. `portable-pty`'s `ExitStatus` exposes
/// `exit_code(): u32`; `0` is success. We surface it as `i32`.
fn exit_code(status: &portable_pty::ExitStatus) -> Option<i32> {
    Some(status.exit_code() as i32)
}

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

    /// Regression: a re-attach (new subscribe) must end the previous subscriber's
    /// stream so only ONE consumer is fed. Two live subscribers would each pump
    /// the same bytes to the bridge → the on-resize output duplication.
    #[test]
    fn subscribe_supersedes_the_previous_subscriber() {
        let mut hub = Broadcast::default();
        let first = hub.subscribe();
        let second = hub.subscribe();

        hub.push(b"hello");

        // The superseded receiver's stream is closed (sender dropped) and gets
        // nothing; only the current subscriber receives the chunk.
        assert!(first.recv().is_err(), "old subscriber must be disconnected");
        assert_eq!(second.recv().unwrap(), b"hello".to_vec());
    }

    #[test]
    fn subscribe_after_eof_yields_a_closed_stream() {
        let mut hub = Broadcast::default();
        hub.close_subscribers();
        let rx = hub.subscribe();
        assert!(rx.recv().is_err(), "no output will ever come after EOF");
    }

    #[test]
    fn snapshot_retains_scrollback_across_resubscribe() {
        let mut hub = Broadcast::default();
        let _first = hub.subscribe();
        hub.push(b"abc");
        // A re-attach drops the old subscriber but the scrollback is preserved.
        let _second = hub.subscribe();
        assert_eq!(hub.snapshot(), b"abc".to_vec());
    }
}