alktty/src/session.rs

//! Consumer half — `TtySession`, the typed client wrapper around the
//! `alk/tty` wire protocol (per alkcall's protocol-crate pattern).
//!
//! Two constructors:
//!
//! - [`TtySession::connect_direct`] — for a direct `alk/tty` ALPN
//!   connection. The consumer dials the transport (TLS, QUIC,
//!   WebSocket), negotiates the `alk/tty` ALPN, and hands the
//!   `Connection` to `connect_direct`. The session takes ownership of
//!   the connection's single `BiStream` (yield-once per connection,
//!   ADR-065), writes the negotiation frame, and exposes the typed
//!   methods.
//!
//! - [`TtySession::open_via_channels`] — for the `alk/channels`
//!   multiplexed path. The consumer holds a
//!   [`alkcall::channels::client::ChannelClient`], calls
//!   `open_via_channels(client, params)`, which invokes
//!   `channels/tty/sub` on channel 0, adopts the resulting channel,
//!   builds a `Connection` from the reassembled read half + mux write
//!   half, and runs the typed-methods flow directly in raw chunk mode
//!   (ADR-009: the open op's `params` — validated by the registry's
//!   input schema — *are* the negotiation; no second negotiation
//!   frame is written on the channel's data stream).
//!
//! The session handle exposes:
//! - [`TtySession::send_stdin`] / [`TtySession::close_stdin`] — write
//!   stdin chunks, close stdin (zero-length sentinel).
//! - [`TtySession::recv_stdout`] / [`TtySession::recv_stderr`] —
//!   streams of stdout/stderr chunks (`Stream<Item = Bytes>`).
//! - [`TtySession::resize`] — send a `Resize` control message.
//! - [`TtySession::signal`] — send a `Signal` control message.
//! - [`TtySession::wait`] — await the `Exit` control chunk (the
//!   process exit code).
//!
//! The session does NOT re-serialize the negotiation request itself
//! — the caller passes a `NegotiateRequest` (direct path) or a
//! `serde_json::Value` (channels path — the open op's `input` is the
//! negotiation). The direct path writes the frame and switches to
//! raw-chunk mode; the channels path starts in raw-chunk mode. See
//! ADR-052 for the two-carriage model.
//!
//! # WASM
//!
//! `TtySession` is wasm-clean (no `tokio::process`, no `std::thread`,
//! no `libc`). The consumer half is exactly the part a browser-side
//! or Python-wasm adapter would use — it runs the wire protocol
//! against a `Connection` the consumer dials (a WebSocket binary
//! stream, a WebTransport bidi stream, etc.). The producer half
//! (`TtyAdapter` + `register_openable`) runs on a real OS with a
//! backend that can spawn processes.

use std::collections::HashMap;
use std::pin::Pin;

use bytes::Bytes;
use futures::Stream;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::sync::{mpsc, Mutex};
use tokio::task::JoinHandle;
use tracing::{debug, warn};

use alkcall::channels::client::ChannelClient;
use alkcall::core::Connection;

use crate::control::ControlMessage;
use crate::negotiation::{NegotiateRequest, NegotiationError, NegotiationWriter};
use crate::wire::{Chunk, ChunkReader, ChunkWriter, RawError, STREAM_CTRL_IN, STREAM_CTRL_OUT};

/// Errors from the typed consumer client.
///
/// `#[non_exhaustive]` so new variants are additive (the same
/// two-way-door pattern as [`crate::backend::TtyError`], and the same
/// justification alkcall gives its consumer-facing `AdapterError`):
/// session drivers accrue failure modes (the channels-path fail-fast
/// variant was added pre-1.0), and an exhaustive match on this enum in
/// a downstream consumer would turn every addition into a breaking
/// change.
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum TtySessionError {
    /// The underlying transport I/O failed (not a clean EOF).
    #[error("io: {0}")]
    Io(#[from] std::io::Error),
    /// The raw-chunk codec errored (invalid stream type, chunk too
    /// large, transport I/O).
    #[error("wire: {0}")]
    Wire(#[from] RawError),
    /// The negotiation frame failed to serialize.
    #[error("negotiation serialize: {0}")]
    NegotiationSerialize(#[from] serde_json::Error),
    /// The negotiation frame failed to write (framing I/O error).
    #[error("negotiation write: {0}")]
    NegotiationWrite(#[from] NegotiationError),
    /// The channels open op failed (returned a `CallError` or
    /// `channel_id` missing from the response).
    #[error("channels open: {0}")]
    ChannelsOpen(String),
    /// The channels open op's `params` failed the consumer's local
    /// `NegotiateRequest` parse — the fail-fast check in
    /// [`TtySession::open_via_channels`] before a channel is allocated.
    /// (R5: the parse previously surfaced as `NegotiationSerialize`,
    /// a variant whose name and doc describe serializing the
    /// negotiation frame, not parsing open-op params.)
    #[error("invalid open params: {0}")]
    InvalidParams(String),
    /// The server sent a negotiation error frame (the first frame on
    /// the stream is a length-prefixed JSON `{"error":"..."}` rather
    /// than a raw chunk).
    #[error("negotiation rejected: {error}")]
    NegotiationRejected {
        error: String,
        fields: HashMap<String, String>,
    },
    /// The session ended (server closed the stream) before an `Exit`
    /// control chunk arrived. `wait()` returns this when the
    /// stdout/stderr pumps drain and no exit chunk was observed.
    #[error("session ended without exit chunk")]
    NoExitChunk,
    /// The `Exit` control chunk's JSON payload failed to parse.
    #[error("malformed exit chunk: {0}")]
    MalformedExitChunk(String),
}

/// A live `alk/tty` session — the typed consumer-side handle.
///
/// Constructed via [`TtySession::connect_direct`] (direct `alk/tty`
/// ALPN) or [`TtySession::open_via_channels`] (multiplexed over
/// `alk/channels`). The session owns the negotiation frame exchange
/// and exposes typed methods for stdin/stdout/stderr/control/exit.
///
/// The session drives a single read pump task (chunks →
/// stdout/stderr/exit channels) and holds the write half for stdin +
/// control messages. Dropping the session cancels the read pump and
/// closes the write half.
pub struct TtySession {
    /// The write half of the bidi stream, wrapped in a `ChunkWriter`.
    /// `send_stdin`, `resize`, `signal`, and `close_stdin` write
    /// through this. Behind a `Mutex` so the methods can take `&self`
    /// and the caller doesn't need `&mut self` to drive the session.
    writer: Mutex<ChunkWriter<Box<dyn AsyncWrite + Send + Unpin>>>,
    /// stdout chunks from the read pump. The caller drains this via
    /// `recv_stdout()`. `Option` so `recv_stdout()` can take it
    /// (calling twice returns an empty stream the second time).
    stdout_rx: Mutex<Option<mpsc::Receiver<Bytes>>>,
    /// stderr chunks from the read pump. `None` for PTY-mode backends
    /// (stdout/stderr merged into stdout by the kernel PTY) or after
    /// `recv_stderr()` has taken it.
    stderr_rx: Mutex<Option<mpsc::Receiver<Bytes>>>,
    /// The exit outcome, resolved by the read pump when it observes the
    /// `Exit` control chunk. `wait()` awaits this. `Option<ExitOutcome>`
    /// starts as `None`; the pump sends `Some(Exited(code))` on exit
    /// chunk, `Some(MalformedExit(msg))` on a malformed exit chunk, or
    /// `Some(NoExitChunk)` on stream close.
    exit_code: tokio::sync::watch::Receiver<Option<ExitOutcome>>,
    /// The read pump task handle. Dropping the session aborts it.
    _read_pump: JoinHandle<()>,
}

/// The cloneable outcome the read pump resolves into the exit watch
/// channel. `wait()` maps this to a [`TtySessionError`] (or the exit
/// code). Kept separate from `TtySessionError` because the watch channel
/// requires `Clone`, and `TtySessionError` carries non-`Clone` payloads
/// (`std::io::Error`, `serde_json::Error`).
#[derive(Debug, Clone)]
enum ExitOutcome {
    /// The `Exit` control chunk was observed with this code.
    Exited(i32),
    /// A `STREAM_CTRL_OUT` chunk failed to parse as a `ControlMessage`.
    MalformedExit(String),
    /// The stream closed before an `Exit` chunk was observed.
    NoExitChunk,
}

impl TtySession {
    /// Connect directly over a `alk/tty` ALPN connection.
    ///
    /// The consumer dials the transport (TLS, QUIC, WebSocket),
    /// negotiates `alk/tty`, and hands the `Connection` here. The
    /// session takes the connection's single `BiStream` (yield-once
    /// per single-stream connection, ADR-065), writes the negotiation
    /// frame, and starts the read pump.
    ///
    /// `negotiate` is the [`NegotiateRequest`] the session writes as
    /// the first frame. The caller builds it (the typed shape is
    /// easier to construct than a raw JSON `Value`); the session
    /// serializes it.
    pub async fn connect_direct(
        connection: Connection,
        negotiate: NegotiateRequest,
    ) -> Result<Self, TtySessionError> {
        let stream = connection.accept_bi().await.map_err(|e| {
            std::io::Error::new(std::io::ErrorKind::ConnectionReset, format!("{e}"))
        })?;
        Self::from_bidi_stream(stream, negotiate).await
    }

    /// Open a TTY session via `alk/channels` — the multiplexed path.
    ///
    /// The consumer holds a [`ChannelClient`], calls
    /// `open_via_channels(client, params)`, which invokes
    /// `channels/tty/sub` on channel 0, adopts the resulting channel,
    /// builds a `Connection` from the reassembled read half + mux
    /// write half, and runs the typed-methods flow directly in raw
    /// chunk mode (ADR-009: the open op's `params` — validated by the
    /// registry's input schema — *are* the negotiation; no second
    /// negotiation frame is written on the channel's data stream).
    ///
    /// `params` is the `NegotiateRequest` as a `serde_json::Value` —
    /// the channels open op takes a `Value`, not a typed struct. The
    /// session parses it as a `NegotiateRequest` before opening (so a
    /// malformed request fails fast, before a channel is allocated)
    /// and the producer parses the same value from the open op.
    ///
    /// Failures before the channel opens (ACL denial, unknown op,
    /// channel cap) surface as [`TtySessionError::ChannelsOpen`]; a
    /// params value that fails the local `NegotiateRequest` parse
    /// surfaces as [`TtySessionError::InvalidParams`]. Post-open
    /// failures (a `NegotiateRequest` parse failure of a
    /// schema-valid-but-unparseable params value, unknown backend,
    /// allocate failure, ownership denial) arrive as a negotiation
    /// error frame on the channel stream — the session surfaces those
    /// as [`TtySessionError::NegotiationRejected`] via the same `0x00`
    /// disambiguation read the direct path uses.
    pub async fn open_via_channels(
        client: &ChannelClient,
        params: serde_json::Value,
    ) -> Result<Self, TtySessionError> {
        let _: NegotiateRequest = serde_json::from_value(params.clone())
            .map_err(|e| TtySessionError::InvalidParams(e.to_string()))?;
        let (channel_id, send, recv) = client
            .open_channel(
                crate::channels::OP_TTY_OPEN,
                params,
                crate::channels::TTY_ALPN,
            )
            .await
            .map_err(TtySessionError::ChannelsOpen)?;
        debug!("tty: opened channel {channel_id} via channels");

        let remote_addr = client.manager().remote_addr();
        let source = alkcall::channels::source::channel_source(recv, send, remote_addr);
        let channel_conn =
            Connection::from_source(source, crate::channels::TTY_ALPN.as_bytes().to_vec());

        Self::from_bidi_stream_via(channel_conn).await
    }

    /// Shared inner: take a `BiStream`, write the negotiation frame,
    /// start the read pump. Used by both `connect_direct` and (after
    /// the channels open) `open_via_channels`.
    async fn from_bidi_stream(
        stream: alkcall::core::BiStream,
        negotiate: NegotiateRequest,
    ) -> Result<Self, TtySessionError> {
        let (read, write) = tokio::io::split(stream);
        Self::from_halves(read, write, negotiate).await
    }

    /// Like `from_bidi_stream` but takes the channel's `Connection`
    /// directly (the channels path already has the `Connection` from
    /// `Connection::from_source`). The negotiation already happened in
    /// the open op (ADR-009) — the stream starts in raw-chunk mode.
    async fn from_bidi_stream_via(channel_conn: Connection) -> Result<Self, TtySessionError> {
        let stream = channel_conn.accept_bi().await.map_err(|e| {
            std::io::Error::new(std::io::ErrorKind::ConnectionReset, format!("{e}"))
        })?;
        let (read, write) = tokio::io::split(stream);
        Self::from_halves_raw(read, write).await
    }

    /// Core inner: take a read half and a write half, write the
    /// negotiation frame, spawn the read pump, return the session.
    async fn from_halves<R, W>(
        read: R,
        write: W,
        negotiate: NegotiateRequest,
    ) -> Result<Self, TtySessionError>
    where
        R: AsyncRead + Send + Unpin + 'static,
        W: AsyncWrite + Send + Unpin + 'static,
    {
        let boxed_write: Box<dyn AsyncWrite + Send + Unpin> = Box::new(write);
        let mut neg_writer = NegotiationWriter::new(boxed_write);
        let body = serde_json::to_vec(&negotiate)?;
        neg_writer.write_frame(&body).await?;
        let writer = ChunkWriter::new(neg_writer.into_inner());

        let mut reader = ChunkReader::new(read);
        // Disambiguate the first response frame (ADR-052 §5): a
        // negotiation error frame's 4-byte length prefix starts with
        // `0x00`, while a raw chunk's first byte is a `stream_type` in
        // `{1, 2, 4}` (the server never sends `0` or `3`). If the server
        // rejected the negotiation, read the error frame and return
        // `NegotiationRejected`; otherwise hand the peeked reader to the
        // read pump.
        let mut first_byte_peeked = false;
        match reader.peek_stream_type().await {
            Ok(0x00) => {
                return Err(read_negotiation_error(reader.into_inner()).await);
            }
            Ok(_) => first_byte_peeked = true,
            Err(RawError::ConnectionClosed) => {
                // The server closed cleanly without a response. Fall
                // through to the read pump, which resolves `NoExitChunk`.
            }
            Err(e) => return Err(TtySessionError::Wire(e)),
        }

        Self::start_pump(writer, reader, first_byte_peeked)
    }

    /// Core inner for the channels path (ADR-009): the negotiation
    /// already happened in the open op — the stream is already in
    /// raw-chunk mode. The peek still applies: the producer sends a
    /// `0x00`-prefixed error frame on any post-open failure (a
    /// `NegotiateRequest` parse failure of the open op's `input`,
    /// unknown backend, allocate failure, ownership denial), and a raw
    /// chunk (`stream_type` in `{1, 2, 4}`) on success.
    async fn from_halves_raw<R, W>(read: R, write: W) -> Result<Self, TtySessionError>
    where
        R: AsyncRead + Send + Unpin + 'static,
        W: AsyncWrite + Send + Unpin + 'static,
    {
        let writer = ChunkWriter::new(Box::new(write) as Box<dyn AsyncWrite + Send + Unpin>);
        let mut reader = ChunkReader::new(read);
        let mut first_byte_peeked = false;
        match reader.peek_stream_type().await {
            Ok(0x00) => {
                return Err(read_negotiation_error(reader.into_inner()).await);
            }
            Ok(_) => first_byte_peeked = true,
            Err(RawError::ConnectionClosed) => {
                // The server closed cleanly without a response. Fall
                // through to the read pump, which resolves `NoExitChunk`.
            }
            Err(e) => return Err(TtySessionError::Wire(e)),
        }

        Self::start_pump(writer, reader, first_byte_peeked)
    }

    /// Wire up the exit watch + stdout/stderr channels and spawn the
    /// read pump. Shared by `from_halves` and `from_halves_raw`.
    fn start_pump<R>(
        writer: ChunkWriter<Box<dyn AsyncWrite + Send + Unpin>>,
        reader: ChunkReader<R>,
        first_byte_peeked: bool,
    ) -> Result<Self, TtySessionError>
    where
        R: AsyncRead + Send + Unpin + 'static,
    {
        let (stdout_tx, stdout_rx) = mpsc::channel::<Bytes>(64);
        let (stderr_tx, stderr_rx) = mpsc::channel::<Bytes>(64);
        let (exit_tx, exit_rx) = tokio::sync::watch::channel::<Option<ExitOutcome>>(None);

        let read_pump = tokio::spawn(read_pump(
            reader,
            first_byte_peeked,
            stdout_tx,
            stderr_tx,
            exit_tx,
        ));

        Ok(Self {
            writer: Mutex::new(writer),
            stdout_rx: Mutex::new(Some(stdout_rx)),
            stderr_rx: Mutex::new(Some(stderr_rx)),
            exit_code: exit_rx,
            _read_pump: read_pump,
        })
    }

    /// Send stdin bytes. Writes a stdin chunk (stream_type 0) with the
    /// given payload. An empty `bytes` writes a zero-length sentinel
    /// (client stdin EOF — see `tty-wire.md` §"Sentinels"); callers
    /// that want to signal EOF should use [`Self::close_stdin`] instead,
    /// which is explicit.
    pub async fn send_stdin(&self, bytes: Bytes) -> Result<(), TtySessionError> {
        let mut writer = self.writer.lock().await;
        let chunk = Chunk::stdin(bytes);
        writer.write_chunk(&chunk).await?;
        Ok(())
    }

    /// Close stdin — send a zero-length stdin chunk (the EOF sentinel)
    /// and flush. The server closes the backend's stdin but keeps
    /// pumping stdout + the exit chunk (see `tty-wire.md` §"Stdin
    /// Closure").
    pub async fn close_stdin(&self) -> Result<(), TtySessionError> {
        let mut writer = self.writer.lock().await;
        let chunk = Chunk::stdin(Bytes::new());
        writer.write_chunk(&chunk).await?;
        Ok(())
    }

    /// Send a `Resize` control message (client→server, `STREAM_CTRL_IN`).
    /// `pixel_width`/`pixel_height` default to 0 (most terminals don't
    /// report pixel dimensions).
    pub async fn resize(
        &self,
        cols: u16,
        rows: u16,
        pixel_width: u16,
        pixel_height: u16,
    ) -> Result<(), TtySessionError> {
        let mut writer = self.writer.lock().await;
        let msg = ControlMessage::Resize {
            cols,
            rows,
            pixel_width,
            pixel_height,
        };
        let json = msg.to_json()?;
        let chunk = Chunk::ctrl_in(json);
        writer.write_chunk(&chunk).await?;
        Ok(())
    }

    /// Send a `Signal` control message (client→server,
    /// `STREAM_CTRL_IN`). `name` is an uppercase string from the
    /// supported set (`HUP`, `INT`, `QUIT`, `TERM`, `KILL`, `USR1`,
    /// `USR2`, `TSTP`, `CONT` — see [`crate::control::signal_from_name`]).
    /// Unknown names are forwarded as-is; the backend decides whether
    /// to ignore or fall back to its default kill.
    pub async fn signal(&self, name: &str) -> Result<(), TtySessionError> {
        let mut writer = self.writer.lock().await;
        let msg = ControlMessage::Signal {
            name: name.to_string(),
        };
        let json = msg.to_json()?;
        let chunk = Chunk::ctrl_in(json);
        writer.write_chunk(&chunk).await?;
        Ok(())
    }

    /// Get the stdout stream. Returns a `Stream<Item = Bytes>` that
    /// yields stdout chunks as they arrive. The stream ends when the
    /// server's stdout reaches EOF (a zero-length stdout sentinel
    /// chunk, see `tty-wire.md` §"Sentinels").
    ///
    /// This consumes the stdout receiver — calling it twice returns
    /// an empty stream the second time (the receiver is behind a
    /// `Mutex<Option<...>>` and is taken).
    pub async fn recv_stdout(&self) -> Pin<Box<dyn Stream<Item = Bytes> + Send>> {
        let mut guard = self.stdout_rx.lock().await;
        if let Some(rx) = guard.take() {
            return Box::pin(futures::stream::unfold(rx, |mut rx| async move {
                rx.recv().await.map(|bytes| (bytes, rx))
            }));
        }
        // Already taken — return an empty stream.
        Box::pin(futures::stream::empty())
    }

    /// Get the stderr stream. `None` for PTY-mode backends
    /// (stdout/stderr merged into stdout by the kernel PTY), or if
    /// already taken. The stream ends when the server's stderr reaches
    /// EOF.
    pub async fn recv_stderr(&self) -> Option<Pin<Box<dyn Stream<Item = Bytes> + Send>>> {
        let mut guard = self.stderr_rx.lock().await;
        let rx = guard.take()?;
        Some(Box::pin(futures::stream::unfold(rx, |mut rx| async move {
            rx.recv().await.map(|bytes| (bytes, rx))
        })))
    }

    /// Await the `Exit` control chunk and return the process exit
    /// code. The exit chunk is the last control chunk before stream
    /// close (ADR-055). `code` is `i32` matching
    /// `std::process::ExitStatus::code()`; negative values are
    /// signal-terminated, `-1` is the adapter's best-effort "backend
    /// could not determine the exit code" sentinel.
    ///
    /// Returns [`TtySessionError::NoExitChunk`] if the session ends
    /// (server closed the stream) before an exit chunk is observed.
    pub async fn wait(&self) -> Result<i32, TtySessionError> {
        let mut rx = self.exit_code.clone();
        // If the pump already resolved before we started waiting, the
        // watch's current value is `Some(_)` — return it.
        {
            let borrow = rx.borrow();
            if let Some(outcome) = borrow.as_ref() {
                return outcome_to_result(outcome);
            }
        }
        // Wait for the read pump to send a value.
        rx.changed()
            .await
            .map_err(|_| TtySessionError::NoExitChunk)?;
        let borrow = rx.borrow();
        match borrow.as_ref() {
            Some(outcome) => outcome_to_result(outcome),
            None => Err(TtySessionError::NoExitChunk),
        }
    }
}

impl Drop for TtySession {
    fn drop(&mut self) {
        self._read_pump.abort();
    }
}

/// Map a resolved [`ExitOutcome`] to the `wait()` result.
fn outcome_to_result(outcome: &ExitOutcome) -> Result<i32, TtySessionError> {
    match outcome {
        ExitOutcome::Exited(code) => Ok(*code),
        ExitOutcome::MalformedExit(msg) => Err(TtySessionError::MalformedExitChunk(msg.clone())),
        ExitOutcome::NoExitChunk => Err(TtySessionError::NoExitChunk),
    }
}

/// Read a negotiation error frame (ADR-052 §5) from the raw transport
/// and map it to [`TtySessionError::NegotiationRejected`]. The caller
/// has already peeked the first byte (`0x00`); this reads the remaining
/// 3 length bytes, the body, and parses the `{"error": "...", ...}`
/// JSON. Any non-`error` fields are collected into the `fields` map.
async fn read_negotiation_error<R>(mut read: R) -> TtySessionError
where
    R: AsyncRead + Unpin,
{
    use tokio::io::AsyncReadExt;

    let mut len_rest = [0u8; 3];
    if let Err(e) = read.read_exact(&mut len_rest).await {
        return TtySessionError::Wire(RawError::Io(e));
    }
    let length = u32::from_be_bytes([0x00, len_rest[0], len_rest[1], len_rest[2]]) as usize;
    let mut body = vec![0u8; length];
    if let Err(e) = read.read_exact(&mut body).await {
        return TtySessionError::Wire(RawError::Io(e));
    }

    let value: serde_json::Value = match serde_json::from_slice(&body) {
        Ok(v) => v,
        Err(_) => {
            return TtySessionError::NegotiationRejected {
                error: String::from_utf8_lossy(&body).into_owned(),
                fields: HashMap::new(),
            };
        }
    };

    let mut fields = HashMap::new();
    let mut error = String::new();
    if let Some(obj) = value.as_object() {
        for (k, v) in obj {
            if k == "error" {
                if let Some(s) = v.as_str() {
                    error = s.to_string();
                }
            } else if let Some(s) = v.as_str() {
                fields.insert(k.clone(), s.to_string());
            }
        }
    }
    TtySessionError::NegotiationRejected { error, fields }
}

/// The read pump: reads chunks off the bidi stream's read half and
/// routes them to the stdout/stderr/exit channels. The pump owns the
/// `ChunkReader`. When the stream closes (clean EOF or transport
/// error), the pump drains the stdout/stderr channels (drops the
/// senders) and resolves the exit watch with `NoExitChunk` if no
/// `Exit` chunk was observed, or with the parsed exit code if one was.
///
/// The pump distinguishes the four stream types:
/// - `STREAM_STDOUT` (1) → stdout channel
/// - `STREAM_STDERR` (2) → stderr channel
/// - `STREAM_CTRL_OUT` (4) → control message; parses as
///   `ControlMessage` and, if it's `Exit`, resolves the exit watch
/// - `STREAM_STDIN` (0) / `STREAM_CTRL_IN` (3) — client→server only;
///   the server shouldn't send these, the pump ignores them (with a
///   debug log)
async fn read_pump<R>(
    mut reader: ChunkReader<R>,
    mut first_byte_peeked: bool,
    stdout_tx: mpsc::Sender<Bytes>,
    stderr_tx: mpsc::Sender<Bytes>,
    exit_tx: tokio::sync::watch::Sender<Option<ExitOutcome>>,
) where
    R: AsyncRead + Send + Unpin + 'static,
{
    let mut exit_resolved = false;
    loop {
        let read = if first_byte_peeked {
            first_byte_peeked = false;
            reader.read_chunk_after_peek().await
        } else {
            reader.read_chunk().await
        };
        match read {
            Ok(chunk) => match chunk.stream_type {
                crate::wire::STREAM_STDOUT => {
                    if stdout_tx.send(chunk.bytes).await.is_err() {
                        debug!("tty: stdout receiver dropped, ending read pump");
                        break;
                    }
                }
                crate::wire::STREAM_STDERR => {
                    if stderr_tx.send(chunk.bytes).await.is_err() {
                        debug!("tty: stderr receiver dropped, ending read pump");
                        break;
                    }
                }
                STREAM_CTRL_OUT => match ControlMessage::from_slice(&chunk.bytes) {
                    Ok(ControlMessage::Exit { code }) => {
                        let _ = exit_tx.send(Some(ExitOutcome::Exited(code)));
                        exit_resolved = true;
                        debug!("tty: exit chunk received, code={code}");
                        break;
                    }
                    Ok(other) => {
                        debug!("tty: ignoring non-exit control on STREAM_CTRL_OUT: {other:?}");
                    }
                    Err(e) => {
                        let _ = exit_tx.send(Some(ExitOutcome::MalformedExit(e.to_string())));
                        exit_resolved = true;
                        break;
                    }
                },
                STREAM_CTRL_IN | crate::wire::STREAM_STDIN => {
                    debug!(
                        "tty: ignoring client→server stream_type {} from server",
                        chunk.stream_type
                    );
                }
                other => {
                    debug!("tty: ignoring unknown stream_type {other}");
                }
            },
            Err(RawError::ConnectionClosed) => {
                debug!("tty: read pump: stream closed");
                break;
            }
            Err(e) => {
                warn!("tty: read pump: chunk read error: {e}");
                break;
            }
        }
    }
    // Drain the channels (drop the senders so the receivers observe
    // EOF). If no exit chunk was observed, resolve the watch with
    // `NoExitChunk`.
    drop(stdout_tx);
    drop(stderr_tx);
    if !exit_resolved {
        let _ = exit_tx.send(Some(ExitOutcome::NoExitChunk));
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::backend::{MockBackend, TtyBackend};
    use crate::negotiation::NegotiateRequest;
    use alkcall::core::auth::Identity;
    use alkcall::core::types::Connection;
    use futures::stream::StreamExt;
    use std::collections::HashMap;
    use std::sync::Arc;
    use tokio::io::duplex;

    /// Build a `NegotiateRequest` for tests — minimal valid shape.
    fn test_negotiate(backend: &str) -> NegotiateRequest {
        NegotiateRequest {
            carriage: "raw".to_string(),
            backend: backend.to_string(),
            tty: None,
            cmd: vec!["true".to_string()],
            cwd: None,
            env: HashMap::new(),
            backend_params: serde_json::Map::new(),
        }
    }

    /// Build a `TtySession` wired to a `drive_session` over a duplex
    /// pair, with a `MockBackend` registered as `"mock"`. The session
    /// is the client side; the server side runs `drive_session` on
    /// the other half of the duplex. Returns the session and the
    /// server-side task handle.
    async fn wire_session_and_server(
        backend: Arc<dyn TtyBackend>,
        identity: Option<Identity>,
    ) -> (TtySession, tokio::task::JoinHandle<()>) {
        let mut backends: HashMap<String, Arc<dyn TtyBackend>> = HashMap::new();
        backends.insert("mock".to_string(), backend);
        let backends = Arc::new(backends);

        let (client, server) = duplex(64 * 1024);
        let (server_read, server_write) = tokio::io::split(server);
        let server_task = tokio::spawn(async move {
            crate::adapter::drive_session(server_write, server_read, backends, None, identity)
                .await;
        });

        let client_conn = Connection::from_bidi(client, b"alk/tty".to_vec(), None);
        let session = TtySession::connect_direct(client_conn, test_negotiate("mock"))
            .await
            .expect("connect_direct");
        (session, server_task)
    }

    #[tokio::test]
    async fn connect_direct_writes_negotiation_frame() {
        // The session writes the negotiation frame on construction.
        // If the server side reads it and dispatches to the backend,
        // the session is wired. A `MockBackend` resolves to exit 0
        // immediately; the session's `wait()` should observe it.
        let backend = Arc::new(MockBackend::with_exit_code(0));
        let identity = Some(Identity {
            id: "alice".to_string(),
            scopes: vec![crate::adapter::TTY_OPEN_SCOPE.to_string()],
            resources: HashMap::new(),
        });
        let (session, _server) = wire_session_and_server(backend, identity).await;
        let code = tokio::time::timeout(std::time::Duration::from_secs(5), session.wait())
            .await
            .expect("wait didn't time out")
            .expect("wait returns exit code");
        assert_eq!(code, 0);
    }

    #[tokio::test]
    async fn send_stdin_round_trips_to_backend() {
        // Use a backend that captures stdin — but `MockBackend` doesn't
        // expose the stdin channel to the test. This test just verifies
        // `send_stdin` doesn't error; the adapter tests cover the
        // stdin-to-backend pump.
        let backend = Arc::new(MockBackend::with_exit_code(0));
        let identity = Some(Identity {
            id: "alice".to_string(),
            scopes: vec![crate::adapter::TTY_OPEN_SCOPE.to_string()],
            resources: HashMap::new(),
        });
        let (session, _server) = wire_session_and_server(backend, identity).await;
        session
            .send_stdin(Bytes::from_static(b"hello"))
            .await
            .expect("send_stdin");
        session.close_stdin().await.expect("close_stdin");
        let _ = tokio::time::timeout(std::time::Duration::from_secs(5), session.wait()).await;
    }

    #[tokio::test]
    async fn resize_and_signal_dont_error() {
        // The session writes control chunks; whether the backend
        // receives them is the adapter's concern (covered by the
        // adapter tests). This test verifies the typed methods
        // serialize and write without error.
        let backend = Arc::new(MockBackend::with_exit_code(0));
        let identity = Some(Identity {
            id: "alice".to_string(),
            scopes: vec![crate::adapter::TTY_OPEN_SCOPE.to_string()],
            resources: HashMap::new(),
        });
        let (session, _server) = wire_session_and_server(backend, identity).await;
        session.resize(80, 24, 0, 0).await.expect("resize");
        session.signal("INT").await.expect("signal");
        let _ = tokio::time::timeout(std::time::Duration::from_secs(5), session.wait()).await;
    }

    #[tokio::test]
    async fn recv_stdout_yields_backend_stdout() {
        // `MockBackend` doesn't pump stdout (it resolves exit
        // immediately), so the stdout stream should be empty. This
        // test verifies the stream API works and ends cleanly.
        let backend = Arc::new(MockBackend::with_exit_code(0));
        let identity = Some(Identity {
            id: "alice".to_string(),
            scopes: vec![crate::adapter::TTY_OPEN_SCOPE.to_string()],
            resources: HashMap::new(),
        });
        let (session, _server) = wire_session_and_server(backend, identity).await;
        let stdout = session.recv_stdout().await;
        let collected: Vec<Bytes> = stdout.collect().await;
        // The backend's stdout stream ends immediately (MockBackend
        // drops its stdout sender on exit), so the stream should be
        // empty or near-empty.
        assert!(
            collected.is_empty() || collected.iter().all(|b| b.is_empty()),
            "mock backend produces no stdout, got {collected:?}"
        );
    }

    #[tokio::test]
    async fn wait_returns_no_exit_chunk_when_server_drops_without_exit() {
        // If the server side drops the stream before sending an exit
        // chunk, `wait()` should return `NoExitChunk`. We simulate
        // this by wiring the session to a duplex where the "server"
        // reads the negotiation frame (so the client's write succeeds)
        // then drops without sending anything back.
        let (client, mut server) = duplex(64);
        let server_handle = tokio::spawn(async move {
            use tokio::io::AsyncReadExt;
            // Read the 4-byte length prefix + body so the client's
            // negotiation write succeeds (duplex buffers are small;
            // a partial write would block and the test would hang).
            let mut len_buf = [0u8; 4];
            let _ = server.read_exact(&mut len_buf).await;
            let len = u32::from_be_bytes(len_buf) as usize;
            let mut body = vec![0u8; len];
            let _ = server.read_exact(&mut body).await;
            // Drop `server` — the client's read pump hits EOF.
        });
        let client_conn = Connection::from_bidi(client, b"alk/tty".to_vec(), None);
        let session = TtySession::connect_direct(client_conn, test_negotiate("mock"))
            .await
            .expect("connect_direct");
        let result = tokio::time::timeout(std::time::Duration::from_secs(5), session.wait())
            .await
            .expect("wait didn't time out");
        assert!(matches!(result, Err(TtySessionError::NoExitChunk)));
        let _ = server_handle.await;
    }

    /// `connect_direct` with a `Connection` whose stream is broken
    /// (server half dropped) should return an error from the
    /// negotiation frame write (`BrokenPipe`).
    #[tokio::test]
    async fn connect_direct_errors_when_stream_is_broken() {
        let (client, server) = duplex(64);
        drop(server);
        let client_conn = Connection::from_bidi(client, b"alk/tty".to_vec(), None);
        let result = TtySession::connect_direct(client_conn, test_negotiate("mock")).await;
        assert!(
            result.is_err(),
            "construction should fail when the negotiation write hits a broken pipe"
        );
    }

    /// A backend that emits fixed stdout/stderr chunks before resolving
    /// exit, so the consumer's read-pump routing can be tested with
    /// real data (the `MockBackend` emits nothing).
    struct EmittingBackend {
        stdout: Vec<Bytes>,
        stderr: Vec<Bytes>,
        exit_code: i32,
    }

    #[async_trait::async_trait]
    impl TtyBackend for EmittingBackend {
        async fn allocate(
            &self,
            _params: &crate::backend::TtyParams,
        ) -> Result<crate::backend::TtyHandle, crate::backend::TtyError> {
            use crate::backend::{TtyControlHandle, TtyHandle};
            use tokio_stream::wrappers::ReceiverStream;

            let (stdout_tx, stdout_rx) = mpsc::channel::<Bytes>(8);
            let (stderr_tx, stderr_rx) = mpsc::channel::<Bytes>(8);
            let (_stdin_tx, _stdin_rx) = mpsc::channel::<Bytes>(8);
            let (exit_tx, exit_rx) =
                tokio::sync::oneshot::channel::<Result<i32, crate::backend::TtyError>>();

            let stdout = self.stdout.clone();
            let stderr = self.stderr.clone();
            let code = self.exit_code;
            tokio::spawn(async move {
                for b in stdout {
                    let _ = stdout_tx.send(b).await;
                }
                for b in stderr {
                    let _ = stderr_tx.send(b).await;
                }
                let _ = exit_tx.send(Ok(code));
            });

            let stdout: Pin<Box<dyn Stream<Item = Bytes> + Send>> =
                Box::pin(ReceiverStream::new(stdout_rx));
            let stderr: Option<Pin<Box<dyn Stream<Item = Bytes> + Send>>> =
                Some(Box::pin(ReceiverStream::new(stderr_rx)));
            let stdin: Box<dyn AsyncWrite + Send + Unpin> = Box::new(tokio::io::sink());
            let control = Some(TtyControlHandle::new(Arc::new(
                crate::backend::MockControl::default(),
            )));
            let exit_code: crate::backend::BoxFuture<Result<i32, crate::backend::TtyError>> =
                Box::pin(async move {
                    exit_rx
                        .await
                        .map_err(|_| crate::backend::TtyError::WaitFailed {
                            message: "exit sender dropped".to_string(),
                        })
                        .and_then(|r| r)
                });

            Ok(TtyHandle {
                stdin,
                stdout,
                stderr,
                exit_code,
                control,
            })
        }
    }

    /// The consumer's read pump routes stdout and stderr chunks to the
    /// correct channels (L2). `MockBackend` emits nothing, so this uses
    /// an `EmittingBackend` that produces real stdout/stderr data.
    #[tokio::test]
    async fn recv_stdout_and_stderr_route_backend_data() {
        let backend = Arc::new(EmittingBackend {
            stdout: vec![Bytes::from_static(b"out1"), Bytes::from_static(b"out2")],
            stderr: vec![Bytes::from_static(b"err1")],
            exit_code: 0,
        });
        let identity = Some(Identity {
            id: "alice".to_string(),
            scopes: vec![crate::adapter::TTY_OPEN_SCOPE.to_string()],
            resources: HashMap::new(),
        });
        let (session, _server) = wire_session_and_server(backend, identity).await;

        let stdout = session.recv_stdout().await;
        let collected: Vec<Bytes> = stdout.collect().await;
        // The adapter emits a zero-length stdout sentinel after the
        // backend stream ends; filter it out to assert the data chunks.
        let data: Vec<Bytes> = collected.into_iter().filter(|b| !b.is_empty()).collect();
        assert_eq!(
            data,
            vec![Bytes::from_static(b"out1"), Bytes::from_static(b"out2")],
            "stdout chunks should route to the stdout stream"
        );

        let stderr = session.recv_stderr().await.expect("stderr present");
        let collected: Vec<Bytes> = stderr.collect().await;
        assert_eq!(
            collected,
            vec![Bytes::from_static(b"err1")],
            "stderr chunks should route to the stderr stream"
        );

        let code = tokio::time::timeout(std::time::Duration::from_secs(5), session.wait())
            .await
            .expect("wait didn't time out")
            .expect("wait returns exit code");
        assert_eq!(code, 0);
    }

    /// `wait()` surfaces a malformed exit chunk as
    /// `MalformedExitChunk`, not `NoExitChunk` (M2). The server sends a
    /// `STREAM_CTRL_OUT` chunk whose JSON fails to parse as a
    /// `ControlMessage`.
    #[tokio::test]
    async fn wait_returns_malformed_exit_chunk() {
        let (client, mut server) = duplex(64 * 1024);
        let server_handle = tokio::spawn(async move {
            use tokio::io::{AsyncReadExt, AsyncWriteExt};
            let mut len_buf = [0u8; 4];
            let _ = server.read_exact(&mut len_buf).await;
            let len = u32::from_be_bytes(len_buf) as usize;
            let mut body = vec![0u8; len];
            let _ = server.read_exact(&mut body).await;

            // Write a ctrl_out chunk with a malformed exit payload.
            let payload = br#"{"type":"exit","code":"not-a-number"}"#;
            let mut header = [0u8; 5];
            header[0] = crate::wire::STREAM_CTRL_OUT;
            header[1..].copy_from_slice(&(payload.len() as u32).to_be_bytes());
            let _ = server.write_all(&header).await;
            let _ = server.write_all(payload).await;
            let _ = server.flush().await;
        });
        let client_conn = Connection::from_bidi(client, b"alk/tty".to_vec(), None);
        let session = TtySession::connect_direct(client_conn, test_negotiate("mock"))
            .await
            .expect("connect_direct");
        let result = tokio::time::timeout(std::time::Duration::from_secs(5), session.wait())
            .await
            .expect("wait didn't time out");
        assert!(
            matches!(result, Err(TtySessionError::MalformedExitChunk(_))),
            "expected MalformedExitChunk, got {result:?}"
        );
        let _ = server_handle.await;
    }

    /// `connect_direct` returns `NegotiationRejected` when the server
    /// rejects the negotiation with an error frame (M1). The server
    /// reads the negotiation frame and writes back a length-prefixed
    /// `{"error":"unknown_backend","backend":"nope"}` frame.
    #[tokio::test]
    async fn connect_direct_returns_negotiation_rejected() {
        let (client, mut server) = duplex(64 * 1024);
        let server_handle = tokio::spawn(async move {
            use tokio::io::{AsyncReadExt, AsyncWriteExt};
            let mut len_buf = [0u8; 4];
            let _ = server.read_exact(&mut len_buf).await;
            let len = u32::from_be_bytes(len_buf) as usize;
            let mut body = vec![0u8; len];
            let _ = server.read_exact(&mut body).await;

            let err_body = br#"{"error":"unknown_backend","backend":"nope"}"#;
            let _ = server
                .write_all(&(err_body.len() as u32).to_be_bytes())
                .await;
            let _ = server.write_all(err_body).await;
            let _ = server.flush().await;
        });
        let client_conn = Connection::from_bidi(client, b"alk/tty".to_vec(), None);
        let result = TtySession::connect_direct(client_conn, test_negotiate("mock")).await;
        match result {
            Err(TtySessionError::NegotiationRejected { error, fields }) => {
                assert_eq!(error, "unknown_backend");
                assert_eq!(fields.get("backend").map(String::as_str), Some("nope"));
            }
            Ok(_) => panic!("expected NegotiationRejected, got Ok(session)"),
            Err(other) => panic!("expected NegotiationRejected, got {other:?}"),
        }
        let _ = server_handle.await;
    }

    // --- channels consumer path (L3, over the shared `testing` harness) ----

    use crate::testing::{tty_identity, wire_client_and_server};

    /// The negotiate params as the open op's `input` (same JSON shape
    /// the direct path serializes from `NegotiateRequest`).
    fn test_open_params(backend: &str) -> serde_json::Value {
        serde_json::json!({
            "carriage": "raw",
            "backend": backend,
            "cmd": ["true"],
        })
    }

    fn mock_backends(code: i32) -> Arc<HashMap<String, Arc<dyn TtyBackend>>> {
        let mut backends: HashMap<String, Arc<dyn TtyBackend>> = HashMap::new();
        backends.insert(
            "mock".to_string(),
            Arc::new(MockBackend::with_exit_code(code)),
        );
        Arc::new(backends)
    }

    /// End-to-end (L3): `TtySession::open_via_channels` opens the
    /// channel through the real `register_openable` producer path and
    /// resolves the session. The negotiation travels in the open op's
    /// `input` (ADR-009); the channel stream starts in raw-chunk mode.
    #[tokio::test]
    async fn open_via_channels_end_to_end_negotiates_and_waits() {
        let client =
            wire_client_and_server(mock_backends(0), None, Some(tty_identity("alice"))).await;

        let session = tokio::time::timeout(
            std::time::Duration::from_secs(10),
            TtySession::open_via_channels(&client, test_open_params("mock")),
        )
        .await
        .expect("open_via_channels timed out")
        .expect("session opens");

        let code = tokio::time::timeout(std::time::Duration::from_secs(10), session.wait())
            .await
            .expect("wait didn't time out")
            .expect("wait returns exit code");
        assert_eq!(code, 0);
    }

    /// The same end-to-end path with an unknown backend: the producer
    /// sends a `0x00`-prefixed error frame on the channel stream, the
    /// consumer's post-open peek disambiguates it, and the session
    /// surfaces `NegotiationRejected` (M1 on the channels path).
    #[tokio::test]
    async fn open_via_channels_surfaces_negotiation_rejected() {
        let client =
            wire_client_and_server(mock_backends(0), None, Some(tty_identity("alice"))).await;

        let result = tokio::time::timeout(
            std::time::Duration::from_secs(10),
            TtySession::open_via_channels(&client, test_open_params("nope")),
        )
        .await
        .expect("open_via_channels timed out");
        match result {
            Err(TtySessionError::NegotiationRejected { error, fields }) => {
                assert_eq!(error, "unknown_backend");
                assert_eq!(fields.get("backend").map(String::as_str), Some("nope"));
            }
            Ok(_) => panic!("expected NegotiationRejected, got Ok(session)"),
            Err(other) => panic!("expected NegotiationRejected, got {other:?}"),
        }
    }

    /// The open op's `input` is schema-validated by the registry
    /// (alkcall 0.4): a params value missing the required `backend`
    /// field is rejected before any handler runs, so the failure is a
    /// `ChannelsOpen` error (the open op fails), not a session error.
    #[tokio::test]
    async fn open_via_channels_fails_fast_on_schema_invalid_params() {
        let client =
            wire_client_and_server(mock_backends(0), None, Some(tty_identity("alice"))).await;

        let result = tokio::time::timeout(
            std::time::Duration::from_secs(10),
            TtySession::open_via_channels(
                &client,
                serde_json::json!({ "carriage": "raw", "cmd": ["true"] }),
            ),
        )
        .await
        .expect("open_via_channels timed out");
        assert!(
            matches!(result, Err(TtySessionError::InvalidParams(_))),
            "schema-invalid params fail at the local NegotiateRequest parse (fail-fast, pre-open)"
        );
    }

    /// `open_via_channels` with params that fail the local
    /// `NegotiateRequest` parse (not just the schema): fails fast,
    /// before a channel is allocated.
    #[tokio::test]
    async fn open_via_channels_fails_fast_on_unparseable_params() {
        let client =
            wire_client_and_server(mock_backends(0), None, Some(tty_identity("alice"))).await;

        let result = tokio::time::timeout(
            std::time::Duration::from_secs(10),
            TtySession::open_via_channels(
                &client,
                serde_json::json!({ "carriage": 42, "backend": "mock", "cmd": ["true"] }),
            ),
        )
        .await
        .expect("open_via_channels timed out");
        assert!(
            matches!(result, Err(TtySessionError::InvalidParams(_))),
            "unparseable params must fail before the open op"
        );
    }

    /// End-to-end with an emitting backend (L2's backend): stdout and
    /// stderr route through the channels data plane to the consumer's
    /// typed streams — the full producer+consumer channels path with
    /// real data.
    #[tokio::test]
    async fn open_via_channels_routes_backend_stdout_and_stderr() {
        let mut backends: HashMap<String, Arc<dyn TtyBackend>> = HashMap::new();
        backends.insert(
            "mock".to_string(),
            Arc::new(EmittingBackend {
                stdout: vec![Bytes::from_static(b"ch-out")],
                stderr: vec![Bytes::from_static(b"ch-err")],
                exit_code: 3,
            }),
        );
        let client =
            wire_client_and_server(Arc::new(backends), None, Some(tty_identity("alice"))).await;

        let session = tokio::time::timeout(
            std::time::Duration::from_secs(10),
            TtySession::open_via_channels(&client, test_open_params("mock")),
        )
        .await
        .expect("open_via_channels timed out")
        .expect("session opens");

        let stdout = session.recv_stdout().await;
        let collected: Vec<Bytes> = stdout.collect().await;
        let data: Vec<Bytes> = collected.into_iter().filter(|b| !b.is_empty()).collect();
        assert_eq!(
            data,
            vec![Bytes::from_static(b"ch-out")],
            "stdout should route through the channels data plane"
        );

        let stderr = session.recv_stderr().await.expect("stderr present");
        let collected: Vec<Bytes> = stderr.collect().await;
        assert_eq!(
            collected,
            vec![Bytes::from_static(b"ch-err")],
            "stderr should route through the channels data plane"
        );

        let code = tokio::time::timeout(std::time::Duration::from_secs(10), session.wait())
            .await
            .expect("wait didn't time out")
            .expect("wait returns exit code");
        assert_eq!(code, 3);
    }
}