liminal-cli 0.3.4

//! One turn is active; draft creation and all pipe writes run off the reader.

use super::config::Harness;
use super::exchange::{Delivery, Exchange};
use super::handshake::{Session, Version};
use super::permission::{Decision, Request, decide};
use super::queue::Prompts;
use super::wire::{Frame, Output};
use agent_client_protocol_schema::{v1, v2};
use serde_json::{Value, json};
use sha2::Digest as _;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::task::JoinSet;

#[path = "pump_config.rs"]
mod config;
#[path = "pump_sessions.rs"]
mod sessions;
#[path = "pump_turn.rs"]
mod turn;

struct Pending {
    id: Value,
    request: Request,
    draft: Option<String>,
    digest: Option<String>,
    expected_digest: String,
    cancelled: bool,
    notified: bool,
}

pub(super) enum Publish {
    Stream(Value),
    Refuse(String),
}

pub(super) struct Pump {
    pub session: Session,
    output: Output,
    exchange: Arc<Exchange>,
    prompts: Prompts,
    next_id: u64,
    turn: Option<u64>,
    configuring: Option<u64>,
    opening: Option<(u64, sessions::Opening)>,
    acknowledged: bool,
    idle: bool,
    cancelled: bool,
    pending: BTreeMap<String, Pending>,
    early: BTreeMap<String, Value>,
    notices: BTreeSet<String>,
    pub drafts: JoinSet<(String, anyhow::Result<Value>)>,
    publications: mpsc::Sender<Publish>,
    permission_bound: usize,
    content_bytes: usize,
    /// How many conversations a listing names and how many title characters each keeps.
    listing: (usize, usize),
}

impl Pump {
    pub fn new(
        session: Session,
        output: Output,
        exchange: Arc<Exchange>,
        config: &Harness,
        publications: mpsc::Sender<Publish>,
    ) -> Self {
        Self {
            session,
            output,
            exchange,
            prompts: Prompts::new(config.prompt_queue),
            next_id: 3,
            turn: None,
            configuring: None,
            opening: None,
            acknowledged: false,
            idle: false,
            cancelled: false,
            pending: BTreeMap::new(),
            early: BTreeMap::new(),
            notices: BTreeSet::new(),
            drafts: JoinSet::new(),
            publications,
            permission_bound: config.pending_permissions,
            content_bytes: config.content_bytes,
            listing: (config.listed_sessions, config.title_chars),
        }
    }

    pub fn frame(&mut self, frame: Frame) -> anyhow::Result<()> {
        match frame {
            // A new conversation's first updates may precede its answer; it is adopted on that answer.
            Frame::Notification { method, params }
                if method == v1::CLIENT_METHOD_NAMES.session_update && self.foreign(&params) => {}
            Frame::Notification { method, params }
                if method == v1::CLIENT_METHOD_NAMES.session_update =>
            {
                if self.session.version == Version::Two {
                    let note: v2::UpdateSessionNotification =
                        serde_json::from_value(params.clone())
                            .map_err(|_| anyhow::anyhow!("ACP 2 update params refused"))?;
                    anyhow::ensure!(
                        note.session_id.to_string() == self.session.id,
                        "state update names a foreign session"
                    );
                    if matches!(
                        note.update,
                        v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle(_))
                    ) && self.turn.is_some()
                    {
                        self.idle = true;
                    }
                }
                if let Some(item) = super::stream::project(
                    self.session.version,
                    params,
                    &self.session.id,
                    self.content_bytes,
                )? {
                    if item["type"] == "config_options" {
                        self.session.config = item["options"].clone();
                    }
                    let item = if self.replaying() {
                        super::stream::replayed(item, self.content_bytes)?
                    } else {
                        item
                    };
                    self.publications
                        .try_send(Publish::Stream(item))
                        .map_err(|_| anyhow::anyhow!("stream_queue reached or publisher closed"))?;
                }
                if self.idle && self.acknowledged {
                    self.finish_turn()?;
                }
            }
            Frame::Notification { method, .. } if method.starts_with('_') => {}
            Frame::Notification { method, .. } => {
                anyhow::ensure!(
                    method == v1::PROTOCOL_LEVEL_METHOD_NAMES.cancel_request,
                    "ACP notification is not advertised"
                );
            }
            Frame::Response { id, result }
                if self
                    .opening
                    .as_ref()
                    .is_some_and(|(call, _)| id == json!(call)) =>
            {
                self.opened(result)?;
            }
            Frame::Response { id, result }
                if self.configuring.is_some_and(|call| id == json!(call)) =>
            {
                self.configured(result)?;
            }
            Frame::Response { id, result } => self.turn_response(&id, result)?,
            Frame::Request { id, method, params }
                if method == v1::CLIENT_METHOD_NAMES.session_request_permission =>
            {
                self.permission(id, params)?;
            }
            Frame::Request { id, .. } => self.output.refuse(
                id,
                -32601,
                "method is not advertised: no filesystem, terminal or credential callbacks",
            )?,
        }
        Ok(())
    }

    pub fn delivery(&mut self, delivery: Delivery) -> anyhow::Result<()> {
        match delivery {
            Delivery::Inbox(value) => {
                if self.consume_notice(&value) {
                    return Ok(());
                }
                if let Some(change) = value.get("config") {
                    return self.configure(change);
                }
                if let Some(change) = value.get("session") {
                    return self.session_control(change);
                }
                if let Some(text) = self.prompt_text(&value)? {
                    if self.turn.is_none() {
                        self.start_prompt(text)?;
                    } else {
                        match self.prompts.push(text) {
                            Ok(()) => self
                                .publications
                                .try_send(Publish::Stream(
                                    json!({"type":"prompt_queued","count":self.prompts.len()}),
                                ))
                                .map_err(|_| {
                                    anyhow::anyhow!("stream_queue reached or publisher closed")
                                })?,
                            Err(error) => self
                                .publications
                                .try_send(Publish::Refuse(error.to_string()))
                                .map_err(|_| {
                                    anyhow::anyhow!(
                                        "stream_queue reached while refusing prompt_queue overflow"
                                    )
                                })?,
                        }
                    }
                }
            }
            Delivery::Cancel => self.cancel()?,
            Delivery::State(value) => self.decision(value)?,
        }
        Ok(())
    }

    fn start_prompt(&mut self, text: String) -> anyhow::Result<()> {
        let id = self.next_id;
        self.next_id = id
            .checked_add(1)
            .ok_or_else(|| anyhow::anyhow!("ACP request id exhausted"))?;
        self.turn = Some(id);
        self.acknowledged = false;
        self.idle = false;
        self.cancelled = false;
        // The person's own words open the turn on the stream, so a transcript has both sides.
        self.publications
            .try_send(Publish::Stream(json!({"type":"request","text":text})))
            .map_err(|_| anyhow::anyhow!("stream_queue reached or publisher closed"))?;
        self.output.request(
            id,
            v1::AGENT_METHOD_NAMES.session_prompt,
            super::handshake::prompt(&self.session, text)?,
        )
    }

    fn next_prompt(&mut self) -> anyhow::Result<()> {
        if let Some(text) = self.prompts.pop() {
            self.start_prompt(text)?;
        }
        Ok(())
    }

    fn finish_turn(&mut self) -> anyhow::Result<()> {
        self.turn = None;
        self.idle = false;
        self.acknowledged = false;
        self.cancel_pending()?;
        self.next_prompt()
    }

    pub fn cancel(&mut self) -> anyhow::Result<()> {
        if self.turn.is_some() {
            self.cancelled = true;
            self.cancel_pending()?;
            self.output.notification(
                v1::AGENT_METHOD_NAMES.session_cancel,
                json!({"sessionId":self.session.id}),
            )?;
        }
        Ok(())
    }

    fn cancel_pending(&mut self) -> anyhow::Result<()> {
        for pending in self
            .pending
            .values_mut()
            .filter(|pending| !pending.cancelled)
        {
            pending.cancelled = true;
            self.output.response(
                pending.id.clone(),
                decide(&pending.request, Decision::Cancelled),
            )?;
        }
        self.pending.retain(|_, pending| pending.draft.is_none());
        self.early.clear();
        Ok(())
    }

    fn permission(&mut self, id: Value, params: Value) -> anyhow::Result<()> {
        let Ok(mut request) = Request::decode(
            self.session.version,
            params,
            &self.session.id,
            self.content_bytes,
        ) else {
            return self.output.refuse(
                id,
                -32602,
                "permission params refused: session, title, options or content_bytes",
            );
        };
        if self.cancelled || self.turn.is_none() {
            return self
                .output
                .response(id, decide(&request, Decision::Cancelled));
        }
        let key = serde_json::to_string(&id)?;
        anyhow::ensure!(
            !self.pending.contains_key(&key),
            "ACP permission request id is already held"
        );
        if self.pending.len() >= self.permission_bound {
            return self.output.refuse(
                id,
                -32001,
                &format!(
                    "pending_permissions {} reached: {} held requests",
                    self.permission_bound,
                    self.pending.len()
                ),
            );
        }
        let exchange = Arc::clone(&self.exchange);
        let job_key = key.clone();
        let operation = super::permission::operation()?;
        // The task submits the draft bytes; the waiting host retains only a digest and selection ids.
        let words = serde_json::to_string(&request.words)?;
        let expected_digest = hex::encode(sha2::Sha256::digest(words.as_bytes()));
        let act = request.act.clone();
        let target = request.target.clone();
        let arguments = request.arguments.clone();
        request.discard_payload();
        self.drafts.spawn_blocking(move || {
            // An ordinary permission names no target and no arguments: the server's contract.
            let body = json!({"agent":exchange.name,"operation":operation,"act":act,"target":target,"arguments":arguments,"words":words});
            (job_key, super::permission::held(exchange.request("agents.permission", &body)))
        });
        self.pending.insert(
            key,
            Pending {
                id,
                request,
                draft: None,
                digest: None,
                expected_digest,
                cancelled: false,
                notified: false,
            },
        );
        Ok(())
    }

    pub fn drafted(&mut self, key: &str, answer: &Value) -> anyhow::Result<()> {
        let draft = answer["draft"]
            .as_str()
            .filter(|value| !value.is_empty() && value.len() <= self.content_bytes)
            .ok_or_else(|| anyhow::anyhow!("permission hold has no bounded draft id"))?
            .to_owned();
        let pending = self
            .pending
            .get_mut(key)
            .ok_or_else(|| anyhow::anyhow!("permission hold answered an unknown request"))?;
        let words = answer["words"]
            .as_str()
            .ok_or_else(|| anyhow::anyhow!("permission hold has no approver words"))?;
        let digest = hex::encode(sha2::Sha256::digest(words.as_bytes()));
        anyhow::ensure!(
            digest == pending.expected_digest,
            "permission hold changed the approver words"
        );
        if pending.cancelled {
            self.pending.remove(key);
            return Ok(());
        }
        pending.draft = Some(draft.clone());
        pending.digest = Some(digest);
        if let Some(decision) = self.early.remove(&draft) {
            self.decision(decision)?;
        }
        Ok(())
    }

    fn decision(&mut self, value: Value) -> anyhow::Result<()> {
        let value = value
            .get("last_outcome")
            .filter(|outcome| outcome.is_object())
            .cloned()
            .unwrap_or(value);
        let Some(outcome) = value.get("outcome").and_then(Value::as_str) else {
            return Ok(());
        };
        let decision = match outcome {
            "executed" => Decision::Executed,
            "refused" => Decision::Refused,
            _ => return Ok(()),
        };
        let draft = value["draft"]
            .as_str()
            .filter(|draft| !draft.is_empty())
            .ok_or_else(|| anyhow::anyhow!("draft decision has no id"))?;
        let key = self
            .pending
            .iter()
            .find(|(_, pending)| pending.draft.as_deref() == Some(draft))
            .map(|(key, _)| key.clone());
        let Some(key) = key else {
            if self
                .pending
                .values()
                .any(|pending| pending.draft.is_none() && !pending.cancelled)
            {
                anyhow::ensure!(
                    self.early.len() < self.permission_bound,
                    "pending_permissions bound reached by early decisions"
                );
                self.early.insert(draft.to_owned(), value);
            }
            return Ok(());
        };
        let pending = self
            .pending
            .remove(&key)
            .ok_or_else(|| anyhow::anyhow!("draft request disappeared"))?;
        if value.get("digest").is_some() {
            anyhow::ensure!(
                value["digest"].as_str() == pending.digest.as_deref(),
                "draft decision digest differs from the held act"
            );
        }
        // This subject is server-only; its executed outcome follows the store's atomic receipt.
        if !pending.notified {
            anyhow::ensure!(
                self.notices.len() < self.permission_bound,
                "pending_permissions bound reached by outcome notices"
            );
            self.notices.insert(draft.to_owned());
        }
        self.output
            .response(pending.id, decide(&pending.request, decision))
    }

    fn consume_notice(&mut self, value: &Value) -> bool {
        if value.get("text").is_some()
            || !matches!(value["outcome"].as_str(), Some("executed" | "refused"))
        {
            return false;
        }
        let Some(draft) = value["draft"].as_str() else {
            return false;
        };
        if self.notices.remove(draft) {
            return true;
        }
        if let Some(pending) = self
            .pending
            .values_mut()
            .find(|pending| pending.draft.as_deref() == Some(draft))
        {
            // The inbox cannot decide a permission; its duplicate notice carries no new prompt.
            pending.notified = true;
            return true;
        }
        false
    }
}