megacode_agent_sdk/

client.rs

//! A session (spec slice-ar R8.1, R8.5, R8.6): one process, one reader task,
//! one writer task, and the handle a program holds.

use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use serde_json::{Value, json};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{ChildStdin, ChildStdout};
use tokio::sync::{mpsc, oneshot, watch};

use crate::error::Error;
use crate::options::{ApprovalCallback, Options};
use crate::process::{self, Exit, OptionsFile, Signal};
use crate::tools::{self, CancelSignal, Tool, ToolContext};
use crate::wire::{
    self, ApprovalRequest, Compacted, Decision, Disposition, Event, Info, Line, REASON_LIMIT,
    Responder, Response, State, ToolCall, UserKind,
};

/// How long a start waits for `init` (R6.2).
const INIT_WAIT: Duration = Duration::from_secs(60);
/// A turn's unread events past this end the session (R8.5.5).
const UNREAD_LIMIT: usize = 10_000;
/// How long `close` waits for the exit after `quit` (R8.6.2).
const QUIT_WAIT: Duration = Duration::from_secs(5);
/// How long `close` waits for handler and callback tasks before aborting
/// them (R8.6.2).
const TASK_WAIT: Duration = Duration::from_secs(5);

/// What `compact` did (R8.1.2).
#[derive(Debug, Clone, PartialEq)]
pub enum Compaction {
    /// A turn is running; the compaction runs after it.
    Queued,
    /// The compaction ran: `Some` when one was written.
    Done(Option<Compacted>),
}

/// A queue of events that never blocks its writer (R8.5.3).
struct Queue {
    unread: AtomicUsize,
    overflowed: AtomicBool,
}

struct QueueTx {
    tx: mpsc::UnboundedSender<Result<Event, Error>>,
    shared: Arc<Queue>,
}

impl QueueTx {
    /// Queues `event`; `false` once the reader left more than
    /// [`UNREAD_LIMIT`] unread. A dropped reader discards it.
    fn push(&self, event: Event) -> bool {
        if self.tx.is_closed() {
            return true;
        }
        let unread = self.shared.unread.fetch_add(1, Ordering::SeqCst) + 1;
        let _ = self.tx.send(Ok(event));
        if unread > UNREAD_LIMIT {
            self.shared.overflowed.store(true, Ordering::SeqCst);
            return false;
        }
        true
    }

    fn fail(&self, error: Error) {
        let _ = self.tx.send(Err(error));
    }
}

fn queue() -> (QueueTx, Events) {
    let (tx, rx) = mpsc::unbounded_channel();
    let shared = Arc::new(Queue {
        unread: AtomicUsize::new(0),
        overflowed: AtomicBool::new(false),
    });
    (
        QueueTx {
            tx,
            shared: shared.clone(),
        },
        Events {
            rx,
            shared,
            done: false,
        },
    )
}

/// A stream of events: a turn's, or the session's between turns.
pub struct Events {
    rx: mpsc::UnboundedReceiver<Result<Event, Error>>,
    shared: Arc<Queue>,
    done: bool,
}

impl Events {
    /// The next event; `None` once the stream ended. An error ends it.
    pub async fn next(&mut self) -> Option<Result<Event, Error>> {
        if self.done {
            return None;
        }
        if self.shared.overflowed.load(Ordering::SeqCst) {
            self.done = true;
            return Some(Err(Error::SlowConsumer));
        }
        match self.rx.recv().await {
            Some(Ok(event)) => {
                self.shared.unread.fetch_sub(1, Ordering::SeqCst);
                Some(Ok(event))
            }
            Some(Err(e)) => {
                self.done = true;
                Some(Err(e))
            }
            None => {
                self.done = true;
                None
            }
        }
    }
}

/// One turn's events, `user` to `result` (from `query`, `init` first).
pub struct TurnStream {
    init: Option<Info>,
    events: Events,
    /// `query`'s session, closed once its turn ends.
    owner: Option<Client>,
    /// The session the turn runs in, alive while its events are read.
    _session: Option<Client>,
    closing: Option<tokio::task::JoinHandle<()>>,
}

impl TurnStream {
    /// The next event; `None` after the turn's `result`, once an abort took
    /// its queued prompt back (`Disposition::Dequeued`: the prompt never
    /// ran), and, from `query`, once the session was closed.
    pub async fn next(&mut self) -> Option<Result<Event, Error>> {
        if let Some(info) = self.init.take() {
            return Some(Ok(Event::Init(info)));
        }
        if let Some(closing) = self.closing.take() {
            let _ = closing.await;
            return None;
        }
        let next = self.events.next().await;
        let ended = matches!(next, None | Some(Ok(Event::Result(_))) | Some(Err(_)));
        if ended && let Some(client) = self.owner.take() {
            let closing = tokio::spawn(client.close());
            if next.is_none() {
                let _ = closing.await;
            } else {
                self.closing = Some(closing);
            }
        }
        next
    }
}

/// The turn the session is running for this program (R8.5.4).
struct Turn {
    prompt_id: String,
    /// The prompt's `user` line arrived: the lines after it are the turn's.
    /// Before it, whether the prompt was answered `started` or `queued`,
    /// the lines are another turn's (one a background agent's result
    /// opened) or between turns (a compaction), and the session's.
    opened: bool,
    /// The `id` of every `abort` that named its prompt, registered as the
    /// abort is queued: any one answered `dequeued` took the prompt back,
    /// so it never runs and the turn ends, whatever aborts follow it and
    /// whether or not its caller still waits.
    aborts: HashSet<String>,
    events: QueueTx,
}

/// A waiting command's key: the command and its `id`, as its `response`
/// echoes both, so an `approval` and a `client_tool_result` that share an
/// id never answer each other.
type Key = (String, String);

#[derive(Default)]
struct Session {
    /// Each waiting command's reply.
    pending: HashMap<Key, oneshot::Sender<Response>>,
    turn: Option<Turn>,
    /// `session_events()`'s sender, once taken.
    loose: Option<QueueTx>,
    /// `session_events()` was called.
    loose_given: bool,
    /// Why the session is over, once it is.
    ended: Option<Error>,
    /// Each running tool call's cancellation trigger, by call id; a call's
    /// task removes its own.
    calls: HashMap<String, watch::Sender<bool>>,
}

/// What the writer task takes.
enum Write {
    Line(String),
    /// Close stdin: nothing more is sent.
    Close,
}

/// What the reader, the writer and every handle share.
pub(crate) struct Shared {
    /// The writer task's queue: a line is queued whole, under the session
    /// lock, so lines keep their order and a dropped caller never leaves a
    /// part of one (R8.5.3).
    writer: mpsc::UnboundedSender<Write>,
    session: Mutex<Session>,
    ctl: mpsc::UnboundedSender<Signal>,
    exit: watch::Receiver<Option<Exit>>,
    next_id: AtomicU64,
    tools: HashMap<String, Tool>,
    on_approval: Option<ApprovalCallback>,
    tasks: Mutex<Tasks>,
}

/// The handler and callback tasks `close` waits for, then aborts: each
/// call's or request's own task, and the task its handler or callback runs
/// in (so a panic is caught there).
#[derive(Default)]
struct Tasks {
    own: Vec<tokio::task::JoinHandle<()>>,
    inner: Vec<tokio::task::AbortHandle>,
}

impl Shared {
    fn id(&self) -> String {
        format!("sdk-{}", self.next_id.fetch_add(1, Ordering::SeqCst))
    }

    fn ended(&self) -> Option<Error> {
        self.session.lock().expect("session lock").ended.clone()
    }

    /// Ends the session with `cause`: every stream and waiting command
    /// gets it. The first cause wins.
    fn end(&self, cause: Error) {
        let mut s = self.session.lock().expect("session lock");
        if s.ended.is_some() {
            return;
        }
        s.ended = Some(cause.clone());
        if let Some(turn) = s.turn.take() {
            turn.events.fail(cause.clone());
        }
        if let Some(loose) = s.loose.take() {
            loose.fail(cause.clone());
        }
        s.pending.clear();
        for fire in s.calls.values() {
            fire.send_replace(true);
        }
    }

    /// Fires every running call's cancellation signal.
    fn cancel_calls(&self) {
        let s = self.session.lock().expect("session lock");
        for fire in s.calls.values() {
            fire.send_replace(true);
        }
    }

    fn track(&self, task: tokio::task::JoinHandle<()>) {
        let mut t = self.tasks.lock().expect("tasks lock");
        t.own.retain(|h| !h.is_finished());
        t.own.push(task);
    }

    fn track_inner(&self, task: tokio::task::AbortHandle) {
        let mut t = self.tasks.lock().expect("tasks lock");
        t.inner.retain(|h| !h.is_finished());
        t.inner.push(task);
    }

    /// Waits up to [`TASK_WAIT`] for the tracked tasks, then aborts every
    /// one, those started meanwhile included. The tasks stay tracked while
    /// it waits, and a guard aborts them, so a dropped `close` leaves none
    /// running on.
    async fn finish_tasks(&self) {
        struct AbortAll<'a>(&'a Shared);
        impl Drop for AbortAll<'_> {
            fn drop(&mut self) {
                if let Ok(mut t) = self.0.tasks.lock() {
                    let t = std::mem::take(&mut *t);
                    for task in t.own {
                        task.abort();
                    }
                    for task in t.inner {
                        task.abort();
                    }
                }
            }
        }
        let _abort = AbortAll(self);
        let deadline = tokio::time::Instant::now() + TASK_WAIT;
        while tokio::time::Instant::now() < deadline && !self.tasks_finished() {
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
    }

    fn tasks_finished(&self) -> bool {
        let t = self.tasks.lock().expect("tasks lock");
        t.own.iter().all(|h| h.is_finished())
    }

    /// Writes `command` with `body` and waits for the `response` whose `id`
    /// is `id`.
    async fn command_with_id(
        &self,
        command: &str,
        mut body: Value,
        id: String,
    ) -> Result<Response, Error> {
        body["type"] = json!(command);
        let line = wire::encode(&body)?;
        self.send_line(command, id, line).await
    }

    /// Queues an encoded `command` line and waits for its `response`.
    async fn send_line(&self, command: &str, id: String, line: String) -> Result<Response, Error> {
        let key = (command.to_string(), id);
        let reply = {
            let mut s = self.session.lock().expect("session lock");
            if let Some(e) = &s.ended {
                return Err(e.clone());
            }
            if s.pending.contains_key(&key) {
                let id = &key.1;
                return Err(Error::Refused {
                    code: "bad_argument".into(),
                    message: format!("an answer to {id} is already waiting for its response"),
                });
            }
            if self.writer.send(Write::Line(line)).is_err() {
                return Err(Error::SessionEnded { diagnostics: None });
            }
            let (reply, rx) = oneshot::channel();
            s.pending.insert(key, reply);
            rx
        };
        self.answer(reply).await
    }

    /// Waits for a queued command's `response`: `Refused` when it is an
    /// error, the session's end when it never comes.
    async fn answer(&self, reply: oneshot::Receiver<Response>) -> Result<Response, Error> {
        match reply.await {
            Ok(r) if r.ok => Ok(r),
            Ok(r) => {
                let (code, message) = r.error.unwrap_or_default();
                Err(Error::Refused { code, message })
            }
            Err(_) => Err(self
                .ended()
                .unwrap_or(Error::SessionEnded { diagnostics: None })),
        }
    }

    /// Answers the approval request `request_id` (R8.3.3).
    async fn respond(&self, request_id: &str, decision: Decision) -> Result<(), Error> {
        let mut body = json!({ "id": request_id, "decision": decision.as_wire() });
        if let Decision::Deny {
            reason: Some(reason),
        } = &decision
        {
            body["reason"] = json!(wire::cut(reason, REASON_LIMIT));
        }
        // `approval`'s `id` names the request, so the SDK's own id cannot
        // ride on it: the answer is matched by the request's id, and a
        // second answer while one is in flight is refused before it is sent.
        self.command_with_id("approval", body, request_id.to_string())
            .await
            .map(|_| ())
    }

    /// The stop of R8.6.2: `quit` and the end of stdin are queued, then
    /// the exit is awaited, escalating to signals whether or not the writer
    /// got them through (a harness not reading stdin is stopped all the same).
    async fn shut_down(self: &Arc<Self>) {
        let quit = wire::encode(&json!({ "type": "quit", "id": self.id() })).expect("fits");
        let _ = self.writer.send(Write::Line(quit));
        let _ = self.writer.send(Write::Close);
        let mut exit = self.exit.clone();
        process::stop(&self.ctl, &mut exit, QUIT_WAIT).await;
    }
}

/// Owned by every `Client` clone; the last one dropped kills the process
/// (R8.6.3: a best-effort backstop that never blocks).
struct Handle {
    shared: Arc<Shared>,
}

impl Drop for Handle {
    fn drop(&mut self) {
        let _ = self.shared.ctl.send(Signal::Kill);
    }
}

/// A session: cheap to clone, every method on `&self`, usable while a
/// turn's stream is read (R8.1.2).
#[derive(Clone)]
pub struct Client {
    handle: Arc<Handle>,
    info: Arc<Info>,
}

impl Client {
    /// Starts a session and reads its `init`.
    pub async fn start(options: Options) -> Result<Client, Error> {
        Self::open(None, options).await
    }

    /// Continues the session `session_id` (R2.5): the options send only
    /// what a resume may give, and the same tools.
    pub async fn resume(session_id: &str, options: Options) -> Result<Client, Error> {
        Self::open(Some(session_id), options).await
    }

    async fn open(resume: Option<&str>, options: Options) -> Result<Client, Error> {
        Self::open_with(resume, options, INIT_WAIT).await
    }

    /// [`Client::open`] waiting `init_wait` for `init`.
    async fn open_with(
        resume: Option<&str>,
        options: Options,
        init_wait: Duration,
    ) -> Result<Client, Error> {
        let binary = crate::binary::locate(&options).await?;
        let body = serde_json::to_vec(&options.file_json(resume.is_some()))
            .map_err(|e| Error::Protocol(e.to_string()))?;
        let file = Arc::new(
            OptionsFile::write(&body)
                .map_err(|e| Error::Handshake(format!("cannot write the options file: {e}")))?,
        );
        let mut proc = process::spawn(process::Launch {
            binary: &binary,
            resume,
            mock: options.mock.as_deref(),
            cwd: options.cwd.as_deref(),
            env: &options.env,
            options: file.clone(),
        })
        .await?;
        let failure = match tokio::time::timeout(init_wait, read_init(&mut proc.stdout)).await {
            Err(_) => Error::Handshake(format!("no init within {} s", init_wait.as_secs())),
            Ok(Err(e)) => e,
            Ok(Ok(None)) => {
                let d =
                    process::diagnostics(&mut proc.exit, &proc.stderr, &mut proc.stderr_done).await;
                if d.exit_code == Some(2) {
                    // harness refused its flags or options, in one stderr
                    // line naming the field (R2.1).
                    Error::Refused {
                        code: "bad_argument".into(),
                        message: d.stderr_tail.last().cloned().unwrap_or_default(),
                    }
                } else {
                    Error::SessionEnded {
                        diagnostics: Some(d),
                    }
                }
            }
            Ok(Ok(Some((info, v)))) => match check(&info, v, options.tools.len()) {
                Err(why) => Error::Handshake(why),
                Ok(()) => {
                    file.remove();
                    return Ok(Self::run(info, proc, options));
                }
            },
        };
        // R8.6.4: stdin closed, the process stopped and reaped, the options
        // file removed, then the error.
        let process::Process {
            stdin,
            ctl,
            mut exit,
            ..
        } = proc;
        drop(stdin);
        process::stop(&ctl, &mut exit, QUIT_WAIT).await;
        file.remove();
        Err(failure)
    }

    /// The session after a good `init`: the reader task started.
    fn run(info: Info, proc: process::Process, options: Options) -> Client {
        let (writer, lines) = mpsc::unbounded_channel();
        let shared = Arc::new(Shared {
            writer,
            session: Mutex::new(Session::default()),
            ctl: proc.ctl,
            exit: proc.exit.clone(),
            next_id: AtomicU64::new(1),
            tools: options
                .tools
                .into_iter()
                .map(|t| (t.name.clone(), t))
                .collect(),
            on_approval: options.on_approval,
            tasks: Mutex::new(Tasks::default()),
        });
        tokio::spawn(write(proc.stdin, lines, Arc::downgrade(&shared)));
        tokio::spawn(read(
            shared.clone(),
            proc.stdout,
            proc.exit,
            proc.stderr,
            proc.stderr_done,
        ));
        Client {
            handle: Arc::new(Handle { shared }),
            info: Arc::new(info),
        }
    }

    fn shared(&self) -> &Arc<Shared> {
        &self.handle.shared
    }

    /// `init`'s data.
    pub fn info(&self) -> &Info {
        &self.info
    }

    /// The session's id.
    pub fn id(&self) -> &str {
        &self.info.session_id
    }

    /// Starts a turn with `prompt` and returns its events, `user` to
    /// `result`. One turn at a time: until the wire's `result` (whether or
    /// not the stream is read), or an [`Client::abort`] that took the
    /// prompt back, another `send` is `TurnInProgress`.
    pub async fn send(&self, prompt: &str) -> Result<TurnStream, Error> {
        let shared = self.shared();
        let id = shared.id();
        let line = wire::encode(&json!({ "type": "prompt", "id": id, "text": prompt }))?;
        let (events, reply) = {
            let mut s = shared.session.lock().expect("session lock");
            if let Some(e) = &s.ended {
                return Err(e.clone());
            }
            if s.turn.is_some() {
                return Err(Error::TurnInProgress);
            }
            let (tx, events) = queue();
            s.turn = Some(Turn {
                prompt_id: id.clone(),
                opened: false,
                aborts: HashSet::new(),
                events: tx,
            });
            let (reply, rx) = oneshot::channel();
            let key = ("prompt".to_string(), id.clone());
            s.pending.insert(key.clone(), reply);
            if shared.writer.send(Write::Line(line)).is_err() {
                s.turn = None;
                s.pending.remove(&key);
                return Err(Error::SessionEnded { diagnostics: None });
            }
            (events, rx)
        };
        shared.answer(reply).await?;
        Ok(TurnStream {
            init: None,
            events,
            owner: None,
            _session: Some(self.clone()),
            closing: None,
        })
    }

    /// Joins the running turn at its next landing point.
    pub async fn steer(&self, text: &str) -> Result<Disposition, Error> {
        self.disposition("steer", json!({ "text": text })).await
    }

    /// Interrupts the program's turn, naming its prompt: a running one is
    /// `Interrupting`; one still `queued` (behind a compaction) is taken
    /// back, `Dequeued`, never runs, and its stream ends with no `result`.
    /// With no turn of the program's, whatever runs is interrupted.
    pub async fn abort(&self) -> Result<Disposition, Error> {
        let shared = self.shared();
        let id = shared.id();
        // The turn read, the abort registered with it and the line queued
        // under one lock: the abort names the prompt that precedes it on
        // the wire.
        let reply = {
            let mut s = shared.session.lock().expect("session lock");
            if let Some(e) = &s.ended {
                return Err(e.clone());
            }
            let mut body = json!({ "type": "abort", "id": id });
            if let Some(turn) = s.turn.as_ref() {
                body["prompt_id"] = json!(turn.prompt_id);
            }
            let line = wire::encode(&body)?;
            if shared.writer.send(Write::Line(line)).is_err() {
                return Err(Error::SessionEnded { diagnostics: None });
            }
            if let Some(turn) = s.turn.as_mut() {
                turn.aborts.insert(id.clone());
            }
            let (reply, rx) = oneshot::channel();
            s.pending.insert(("abort".to_string(), id), reply);
            rx
        };
        let r = shared.answer(reply).await?;
        Ok(Disposition::from_wire(
            r.disposition.as_deref().unwrap_or(""),
        ))
    }

    /// Compacts the context: at once when idle, after the turn when one
    /// runs. `Done` carries the compaction its own `response` names
    /// (`data.compacted`); every `compacted` line goes to the running turn's
    /// stream or [`Client::session_events`].
    pub async fn compact(&self) -> Result<Compaction, Error> {
        let r = self.command("compact", json!({})).await?;
        Ok(match r.disposition.as_deref() {
            Some("queued") => Compaction::Queued,
            _ => Compaction::Done(
                r.data
                    .and_then(|mut d| d.get_mut("compacted").map(Value::take))
                    .and_then(|c| serde_json::from_value::<Compacted>(c).ok()),
            ),
        })
    }

    /// Puts the running or next task away, or ends away.
    pub async fn away(&self, on: bool) -> Result<Disposition, Error> {
        self.disposition("away", json!({ "on": on })).await
    }

    /// The session's state.
    pub async fn state(&self) -> Result<State, Error> {
        let r = self.command("get_state", json!({})).await?;
        let data = r
            .data
            .ok_or_else(|| Error::Protocol("a `get_state` response with no `data`".into()))?;
        serde_json::from_value(data)
            .map_err(|e| Error::Protocol(format!("a `get_state` response: {e}")))
    }

    /// Answers the approval request `request_id` (R8.3.3).
    pub async fn respond(&self, request_id: &str, decision: Decision) -> Result<(), Error> {
        self.shared().respond(request_id, decision).await
    }

    /// The events that arrive while no turn's stream is open; the first
    /// call takes the stream, later calls get `None`. Until it is taken,
    /// those events are dropped.
    pub fn session_events(&self) -> Option<Events> {
        let mut s = self.shared().session.lock().expect("session lock");
        if std::mem::replace(&mut s.loose_given, true) {
            return None;
        }
        let (tx, events) = queue();
        match &s.ended {
            Some(e) => tx.fail(e.clone()),
            None => s.loose = Some(tx),
        }
        Some(events)
    }

    /// Ends the session (R8.6.2): every running call's cancellation signal
    /// fires; `quit`, up to 5 s for the exit, SIGTERM, up to 5 s more,
    /// SIGKILL, the process reaped; meanwhile up to 5 s for handler and
    /// callback tasks, then they are aborted. Returns within about 15 s.
    pub async fn close(self) {
        /// Kills the process when `close` is dropped before the stop
        /// finished, whatever clones are still held.
        struct KillOnDrop<'a>(Option<&'a Shared>);
        impl Drop for KillOnDrop<'_> {
            fn drop(&mut self) {
                if let Some(shared) = self.0 {
                    let _ = shared.ctl.send(Signal::Kill);
                }
            }
        }
        let shared = self.shared().clone();
        shared.cancel_calls();
        let mut kill = KillOnDrop(Some(&shared));
        tokio::join!(shared.shut_down(), shared.finish_tasks());
        kill.0 = None;
        // Tasks started while the process stopped.
        shared.finish_tasks().await;
    }

    async fn disposition(&self, command: &str, body: Value) -> Result<Disposition, Error> {
        let r = self.command(command, body).await?;
        Ok(Disposition::from_wire(
            r.disposition.as_deref().unwrap_or(""),
        ))
    }

    async fn command(&self, command: &str, body: Value) -> Result<Response, Error> {
        let id = self.shared().id();
        let mut body = body;
        body["id"] = json!(id);
        self.shared().command_with_id(command, body, id).await
    }
}

impl ApprovalRequest {
    /// Answers this request (R8.3.3), as [`Client::respond`] does by id;
    /// `SessionEnded` once its session is gone.
    pub async fn respond(&self, decision: Decision) -> Result<(), Error> {
        let shared = self
            .responder
            .0
            .as_ref()
            .and_then(std::sync::Weak::upgrade)
            .ok_or(Error::SessionEnded { diagnostics: None })?;
        shared.respond(&self.id, decision).await
    }
}

/// Starts a session, sends `prompt`, and yields `init` and then that turn's
/// events up to its `result`; then the session is closed (R8.1.1).
pub async fn query(prompt: &str, options: Options) -> Result<TurnStream, Error> {
    let client = Client::start(options).await?;
    let mut stream = match client.send(prompt).await {
        Ok(stream) => stream,
        Err(e) => {
            client.close().await;
            return Err(e);
        }
    };
    stream.init = Some((*client.info).clone());
    stream._session = None;
    stream.owner = Some(client);
    Ok(stream)
}

/// The writer task: writes each queued line whole, in order, until `Close`,
/// the queue's end or a failed write; then stdin is closed. A failed write
/// ends the session at once (every waiting command settled, every running
/// call's signal fired) and kills the process, whether or not its stdout
/// ever ends (a descendant can hold it open). It holds the session weakly:
/// the queue ends when the last handle goes.
async fn write(
    mut stdin: ChildStdin,
    mut lines: mpsc::UnboundedReceiver<Write>,
    shared: std::sync::Weak<Shared>,
) {
    while let Some(Write::Line(line)) = lines.recv().await {
        if stdin.write_all(line.as_bytes()).await.is_err() || stdin.flush().await.is_err() {
            if let Some(shared) = shared.upgrade() {
                shared.end(Error::SessionEnded { diagnostics: None });
                let _ = shared.ctl.send(Signal::Kill);
            }
            break;
        }
    }
    let _ = stdin.shutdown().await;
}

/// Reads up to the first `init`, skipping lines of unknown type; `None` at
/// the end of stdout.
async fn read_init(stdout: &mut BufReader<ChildStdout>) -> Result<Option<(Info, u64)>, Error> {
    let mut buf = Vec::new();
    loop {
        buf.clear();
        if stdout.read_until(b'\n', &mut buf).await.unwrap_or(0) == 0 {
            return Ok(None);
        }
        let line = trim_newline(&buf);
        match wire::parse(line).map_err(|e| Error::Handshake(e.to_string()))? {
            Line::Init(info, v) => return Ok(Some((*info, v))),
            Line::Other => continue,
            _ => {
                return Err(Error::Handshake(format!(
                    "the first line was not init: {}",
                    String::from_utf8_lossy(&line[..line.len().min(200)])
                )));
            }
        }
    }
}

fn trim_newline(buf: &[u8]) -> &[u8] {
    let buf = buf.strip_suffix(b"\n").unwrap_or(buf);
    buf.strip_suffix(b"\r").unwrap_or(buf)
}

/// R6.2: protocol `v` 1, a harness of this SDK's major.minor, and
/// `client_tool_result` among the commands when the session has tools.
fn check(info: &Info, v: u64, tools: usize) -> Result<(), String> {
    if v != 1 {
        return Err(format!(
            "the session speaks protocol v{v}; this SDK speaks v1"
        ));
    }
    let ours = env!("CARGO_PKG_VERSION");
    let minor = |version: &str| version.splitn(3, '.').take(2).collect::<Vec<_>>().join(".");
    if minor(&info.harness_version) != minor(ours) {
        return Err(format!(
            "harness {} does not match this SDK ({ours}); it needs harness {}.x",
            info.harness_version,
            minor(ours)
        ));
    }
    if tools > 0 && !info.commands.iter().any(|c| c == "client_tool_result") {
        return Err(format!(
            "harness {} does not take client_tool_result, and the session has tools",
            info.harness_version
        ));
    }
    Ok(())
}

/// How long the reader keeps reading stdout once the process has exited:
/// its last lines are already in the pipe, while a descendant holding the
/// pipe open may never end it.
const DRAIN_AFTER_EXIT: Duration = Duration::from_millis(500);

/// The reader (R8.5.3): reads stdout until its end, or until the process
/// has exited and `DRAIN_AFTER_EXIT` has passed, never waiting on the
/// program; settles responses, routes events, ends the session on a
/// protocol error or a slow consumer.
async fn read(
    shared: Arc<Shared>,
    mut stdout: BufReader<ChildStdout>,
    mut exit: watch::Receiver<Option<Exit>>,
    stderr: Arc<Mutex<VecDeque<String>>>,
    mut stderr_done: watch::Receiver<bool>,
) {
    // A read cut by the exit or the drain's end leaves what it took in
    // `buf`, which is cleared only after a whole line.
    let mut buf = Vec::new();
    // Once the SDK ended the session, the rest of stdout is read and
    // dropped, so the process is never held writing while it stops.
    let mut stopped = false;
    let mut drain_until: Option<tokio::time::Instant> = None;
    loop {
        let read = match drain_until {
            Some(until) => {
                match tokio::time::timeout_at(until, stdout.read_until(b'\n', &mut buf)).await {
                    Ok(read) => read,
                    Err(_) => break,
                }
            }
            None => tokio::select! {
                read = stdout.read_until(b'\n', &mut buf) => read,
                _ = exit.wait_for(Option::is_some) => {
                    drain_until = Some(tokio::time::Instant::now() + DRAIN_AFTER_EXIT);
                    continue;
                }
            },
        };
        if read.unwrap_or(0) == 0 {
            break;
        }
        if stopped {
            buf.clear();
            continue;
        }
        let routed = wire::parse(trim_newline(&buf)).and_then(|line| route(&shared, line));
        buf.clear();
        if let Err(cause) = routed {
            stop_with(&shared, cause);
            stopped = true;
        }
    }
    let diagnostics = process::diagnostics(&mut exit, &stderr, &mut stderr_done).await;
    shared.end(Error::SessionEnded {
        diagnostics: Some(diagnostics),
    });
}

/// Ends the session for `cause` and stops the process (R8.6).
fn stop_with(shared: &Arc<Shared>, cause: Error) {
    shared.end(cause);
    let shared = shared.clone();
    tokio::spawn(async move { shared.shut_down().await });
}

/// One line's effect.
fn route(shared: &Arc<Shared>, line: Line) -> Result<(), Error> {
    let mut s = shared.session.lock().expect("session lock");
    match line {
        Line::ClientToolCall(call) => start_call(shared, &mut s, call),
        Line::ClientToolCancel(id) => {
            if let Some(fire) = s.calls.get(&id) {
                fire.send_replace(true);
            }
        }
        Line::Response(r) => {
            let Some(id) = r.id.clone() else {
                return Ok(());
            };
            // A refused prompt takes no turn, whether or not the caller of
            // `send` still waits for its answer.
            if r.command == "prompt" && !r.ok && s.turn.as_ref().is_some_and(|t| t.prompt_id == id)
            {
                s.turn = None;
            }
            // A prompt an abort took back never runs: its turn ends here.
            if r.command == "abort"
                && r.disposition.as_deref() == Some("dequeued")
                && s.turn.as_ref().is_some_and(|t| t.aborts.contains(&id))
            {
                s.turn = None;
            }
            if let Some(reply) = s.pending.remove(&(r.command.clone(), id)) {
                let _ = reply.send(r);
            }
        }
        Line::Event(event) => {
            let mut event = *event;
            if let Event::ApprovalRequest(request) = &mut event {
                request.responder = Responder(Some(Arc::downgrade(shared)));
                if let Some(callback) = &shared.on_approval {
                    start_approval(shared, callback.clone(), request.clone());
                    return Ok(());
                }
            }
            // The program's turn opens at its prompt's `user` line: the one
            // prompt in flight is the program's, so the first prompt-kind
            // `user` line is its own, and the next `result` ends it. A turn
            // a background agent's result opens has no such line.
            if let Some(turn) = s.turn.as_mut() {
                turn.opened =
                    turn.opened || matches!(&event, Event::User(u) if u.kind == UserKind::Message);
            }
            if let Some(turn) = s.turn.as_ref().filter(|t| t.opened) {
                let end = matches!(event, Event::Result(_));
                let kept = turn.events.push(event);
                if end {
                    s.turn = None;
                }
                if !kept {
                    return Err(Error::SlowConsumer);
                }
                return Ok(());
            }
            if let Some(loose) = &s.loose
                && !loose.push(event)
            {
                return Err(Error::SlowConsumer);
            }
        }
        // A second `init`, and lines of types this SDK does not read.
        Line::Init(..) | Line::Other => {}
    }
    Ok(())
}

/// Runs a `client_tool_call` (R8.4) in its own task: the handler in a task
/// of its own (a panic is the call's error), then the result, unless the
/// call's signal fired first (R8.4.4). A `bad_argument` refusal is followed
/// by an `is_error` result with its message (R8.4.5).
fn start_call(shared: &Arc<Shared>, s: &mut Session, call: ToolCall) {
    let (fire, signal) = tools::signal();
    if s.ended.is_some() {
        fire.send_replace(true);
    }
    s.calls.insert(call.id.clone(), fire);
    let tool = shared.tools.get(&call.name).cloned();
    let calling = shared.clone();
    shared.track(tokio::spawn(async move {
        let ToolCall {
            id,
            name,
            input,
            too_deep,
        } = call;
        let outcome = match tool {
            _ if too_deep => Err("the arguments are nested too deeply".to_string()),
            None => Err(format!("this program has no tool named {name:?}")),
            Some(tool) => {
                let ctx = ToolContext::new(id.clone(), signal.clone());
                let handler = tool.handler.clone();
                let running = tokio::spawn(async move { handler(input, ctx).await });
                calling.track_inner(running.abort_handle());
                match running.await {
                    Ok(outcome) => outcome,
                    Err(e) => Err(format!("the tool's handler {}", tools::failure_text(e))),
                }
            }
        };
        answer_call(&calling, &id, &outcome, &signal).await;
        calling
            .session
            .lock()
            .expect("session lock")
            .calls
            .remove(&id);
    }));
}

async fn answer_call(
    shared: &Shared,
    id: &str,
    outcome: &Result<tools::ToolOutput, String>,
    signal: &CancelSignal,
) {
    if signal.is_fired() {
        return;
    }
    const COMMAND: &str = "client_tool_result";
    let line = tools::result_line(id, outcome);
    if let Err(Error::Refused { code, message }) =
        shared.send_line(COMMAND, id.to_string(), line).await
        && code == "bad_argument"
        && !signal.is_fired()
    {
        let line = tools::error_line(id, &message);
        let _ = shared.send_line(COMMAND, id.to_string(), line).await;
    }
}

/// Runs the approval callback on `request` in its own task (R8.3.2) and
/// answers with its decision; a panic, or a decision not among the
/// request's choices, answers `Deny` with the failure as the reason.
fn start_approval(shared: &Arc<Shared>, callback: ApprovalCallback, request: ApprovalRequest) {
    let answering = shared.clone();
    shared.track(tokio::spawn(async move {
        let id = request.id.clone();
        let choices = request.choices.clone();
        let running = tokio::spawn(async move { callback(request).await });
        answering.track_inner(running.abort_handle());
        let decision = match running.await {
            Ok(d) if choices.iter().any(|c| c.as_str() == d.as_wire()) => d,
            Ok(d) => Decision::Deny {
                reason: Some(format!(
                    "the approval callback chose {}, which is not among the choices offered ({})",
                    d.as_wire(),
                    choices
                        .iter()
                        .map(|c| c.as_str())
                        .collect::<Vec<_>>()
                        .join(", ")
                )),
            },
            Err(e) => Decision::Deny {
                reason: Some(format!("the approval callback {}", tools::failure_text(e))),
            },
        };
        let _ = answering.respond(&id, decision).await;
    }));
}

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

    #[test]
    fn a_start_waits_60_s_for_init() {
        assert_eq!(INIT_WAIT, Duration::from_secs(60));
    }

    /// The test's own folder: on every path out, success or panic, the
    /// files it names are removed, then the (by then empty) folder.
    struct Scratch {
        dir: std::path::PathBuf,
        files: Vec<std::path::PathBuf>,
    }

    impl Scratch {
        fn new(tag: u32) -> Self {
            let dir = std::env::temp_dir().join(format!(
                "megacode-agent-sdk-test-{}-{tag}",
                std::process::id()
            ));
            std::fs::create_dir(&dir).unwrap();
            Self {
                dir,
                files: Vec::new(),
            }
        }

        /// A path in the folder, removed with it.
        fn file(&mut self, name: &str) -> std::path::PathBuf {
            let path = self.dir.join(name);
            self.files.push(path.clone());
            path
        }
    }

    impl Drop for Scratch {
        fn drop(&mut self) {
            for file in &self.files {
                let _ = std::fs::remove_file(file);
            }
            let _ = std::fs::remove_dir(&self.dir);
        }
    }

    #[tokio::test]
    async fn no_init_in_time_is_a_handshake_error_leaving_nothing_behind() {
        use std::os::unix::fs::PermissionsExt;
        let mut scratch = Scratch::new(line!());
        let script = scratch.file("harness-silent");
        let pid = scratch.file("pid");
        let args = scratch.file("args");
        std::fs::write(
            &script,
            format!(
                "#!/bin/sh\necho $$ > '{}'\nfor a in \"$@\"; do printf '%s\\n' \"$a\"; done > '{}'\n\
                 while read -r _; do :; done\n",
                pid.display(),
                args.display()
            ),
        )
        .unwrap();
        std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o755)).unwrap();

        let started = std::time::Instant::now();
        let options = Options::new().binary(&script);
        match Client::open_with(None, options, Duration::from_secs(1)).await {
            Err(Error::Handshake(why)) => assert!(why.contains("no init within 1 s"), "{why}"),
            Err(e) => panic!("{e:?}"),
            Ok(_) => panic!("started"),
        }
        let took = started.elapsed();
        assert!(took >= Duration::from_secs(1), "{took:?}");
        assert!(
            took < Duration::from_secs(5),
            "stdin closed ends the fake: {took:?}"
        );
        let pid: i32 = std::fs::read_to_string(&pid)
            .unwrap()
            .trim()
            .parse()
            .unwrap();
        // SAFETY: signal 0 checks for the process; nothing is sent.
        assert!(unsafe { libc::kill(pid, 0) } != 0, "the process was reaped");
        let sent = std::fs::read_to_string(&args).unwrap();
        let sent: Vec<&str> = sent.lines().collect();
        let file = std::path::PathBuf::from(
            sent[sent.iter().position(|a| *a == "--options").unwrap() + 1],
        );
        assert!(!file.exists() && !file.parent().unwrap().exists());
    }
}