liminal-cli 0.3.30

//! One bounded writer for answers and unsolicited logging lines.

use super::framing;
use super::notifications::{self, Batch};
use super::subscriptions::{Service, Subscriptions};
use serde_json::{Value, json};
use std::io::{self, Write};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

struct State<O> {
    output: O,
    error: Option<String>,
    /// The files this writer handed answers over in, removed when it drops.
    handed: Vec<std::path::PathBuf>,
}

/// Where an answer over the seat's inline bound is handed over whole
/// (SRV-052 R2): `bytes` is `[limits] answer_bytes`, `dir` the seat's own.
pub(super) struct Answers {
    pub bytes: usize,
    /// `None` when the seat has no home of its own (no `HOME`).
    pub dir: Option<std::path::PathBuf>,
}

pub(super) struct Writer<O> {
    state: Mutex<State<O>>,
    limit: usize,
    answers: Answers,
    warning_only: Arc<AtomicBool>,
    subscriptions: Arc<Subscriptions>,
    listens: Arc<super::listen::Listens>,
}

impl<O: Write> Writer<O> {
    pub const fn new(
        output: O,
        limit: usize,
        answers: Answers,
        warning_only: Arc<AtomicBool>,
        subscriptions: Arc<Subscriptions>,
        listens: Arc<super::listen::Listens>,
    ) -> Self {
        Self {
            state: Mutex::new(State {
                output,
                error: None,
                handed: Vec::new(),
            }),
            limit,
            answers,
            warning_only,
            subscriptions,
            listens,
        }
    }

    pub fn channel(&self, message: Value, subject: &super::listen::Subject) -> io::Result<()> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| io::Error::other("the host writer lock is poisoned"))?;
        if let Some(words) = &state.error {
            return Err(io::Error::other(words.clone()));
        }
        if !subject.active.load(Ordering::SeqCst) {
            return Ok(());
        }
        let result = self.put(&mut state, message);
        if let Err(error) = &result {
            state.error = Some(error.to_string());
        }
        drop(state);
        self.finish(result)
    }

    pub fn check(&self) -> io::Result<()> {
        let state = self
            .state
            .lock()
            .map_err(|_| io::Error::other("the host writer lock is poisoned"))?;
        state
            .error
            .as_ref()
            .map_or(Ok(()), |words| Err(io::Error::other(words.clone())))
    }

    pub fn write(
        &self,
        message: Value,
        service: Option<&Service>,
        terminal: bool,
    ) -> io::Result<()> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| io::Error::other("the host writer lock is poisoned"))?;
        if let Some(words) = &state.error {
            return Err(io::Error::other(words.clone()));
        }
        if let Some(service) = service {
            if !terminal && !service.active.load(Ordering::SeqCst) {
                return Ok(());
            }
            if self.warning_only.load(Ordering::SeqCst) && message["params"]["level"] == "info" {
                return Ok(());
            }
        }
        let result = if is_answer(&message) {
            self.put_answer(&mut state, &message)
        } else {
            self.put(&mut state, message)
        };
        if let Err(error) = &result {
            state.error = Some(error.to_string());
        }
        drop(state);
        self.finish(result)
    }

    pub fn record(&self, service: &Service, subject: &str, batch: &Batch) -> io::Result<()> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| io::Error::other("the host writer lock is poisoned"))?;
        if let Some(words) = &state.error {
            return Err(io::Error::other(words.clone()));
        }
        if !service.active.load(Ordering::SeqCst) {
            return Ok(());
        }
        let result: io::Result<()> = (|| {
            if !self.warning_only.load(Ordering::SeqCst) {
                for message in &batch.messages {
                    self.put(&mut state, notifications::message(service, message))?;
                }
            }
            if batch.dropped > 0 {
                self.put(
                    &mut state,
                    notifications::warning(
                        service,
                        subject,
                        "dropped",
                        &format!("{} messages on {subject} were not kept", batch.dropped),
                    ),
                )?;
            }
            Ok(())
        })();
        if let Err(error) = &result {
            state.error = Some(error.to_string());
        }
        drop(state);
        self.finish(result)
    }

    fn put(&self, state: &mut State<O>, mut message: Value) -> io::Result<()> {
        fit(&mut message, self.limit)?;
        framing::write_message(&mut state.output, &message)
    }

    /// An answer within the inline bound is written whole; a larger one is
    /// handed over whole as a file, or, when it cannot be, refuses only its
    /// own call by its size (SRV-052 R1, R2). Neither ends the seat.
    fn put_answer(&self, state: &mut State<O>, message: &Value) -> io::Result<()> {
        let bytes = size(message)?;
        if bytes <= self.answers.bytes {
            return framing::write_message(&mut state.output, message);
        }
        let over = format!(
            "the answer is {bytes} bytes, over this seat's inline bound of {}",
            self.answers.bytes
        );
        let id = message["id"].clone();
        let answer = match self.answers.dir.as_deref().map_or_else(
            || {
                Err(io::Error::new(
                    io::ErrorKind::NotFound,
                    "HOME is not set, so this seat has no directory of its own to hand it over in",
                ))
            },
            |dir| hand_over(dir, message),
        ) {
            Ok(handed) => {
                state.handed.push(handed.path.clone());
                handed.answer(&id, &over)
            }
            Err(error) => json!({"jsonrpc": "2.0", "id": id, "result": {
                "content": [{"type": "text", "text": format!(
                    "{over}; ask for less (flow answer \"last\", keep false) or raise [limits] answer_bytes; \
                     it could not be handed over as a file: {error}"
                )}],
                "isError": true,
            }}),
        };
        framing::write_message(&mut state.output, &answer)
    }

    fn finish(&self, result: io::Result<()>) -> io::Result<()> {
        if let Err(error) = result {
            let mut failures = vec![error.to_string()];
            if let Err(words) = self.subscriptions.close_all() {
                failures.push(words);
            }
            if let Err(words) = self.listens.close_all() {
                failures.push(words);
            }
            return Err(io::Error::other(failures.join("; ")));
        }
        Ok(())
    }
}

impl<O> Drop for Writer<O> {
    fn drop(&mut self) {
        if let Ok(state) = self.state.get_mut() {
            for path in state.handed.drain(..) {
                // A file already gone needs no removing; nothing else is touched.
                let _removed = std::fs::remove_file(path);
            }
            if let Some(dir) = &self.answers.dir {
                // Only an empty directory is removed; anything else stays.
                let _removed = std::fs::remove_dir(dir);
            }
        }
    }
}

/// A JSON-RPC response: an `id` and no `method`.
fn is_answer(message: &Value) -> bool {
    message.get("id").is_some() && message.get("method").is_none()
}

/// The characters of a handed-over answer's text kept in its preview.
const PREVIEW_CHARS: usize = 4096;

/// One answer written whole to a file the seat owns.
struct Handed {
    path: std::path::PathBuf,
    bytes: usize,
    sha256: String,
    preview: String,
}

impl Handed {
    fn answer(&self, id: &Value, over: &str) -> Value {
        json!({"jsonrpc": "2.0", "id": id, "result": {
            "content": [{"type": "text", "text": format!(
                "{over}; it is handed over whole at {} ({} bytes, sha256 {}): read the file for the answer",
                self.path.display(), self.bytes, self.sha256
            )}],
            "structuredContent": {
                "handed_over": {
                    "path": self.path.display().to_string(),
                    "bytes": self.bytes,
                    "sha256": self.sha256,
                    "media_type": "application/json",
                },
                "preview": self.preview,
            },
            "isError": false,
        }})
    }
}

/// Writes the answer's result (or error) to `dir`, mode 0700, as a file
/// named by its sha256, mode 0600.
fn hand_over(dir: &std::path::Path, message: &Value) -> io::Result<Handed> {
    use sha2::Digest as _;
    use std::os::unix::fs::{DirBuilderExt as _, OpenOptionsExt as _};
    let body = message
        .get("result")
        .or_else(|| message.get("error"))
        .unwrap_or(message);
    let bytes = serde_json::to_vec(body)?;
    let sha256 = hex::encode(sha2::Sha256::digest(&bytes));
    std::fs::DirBuilder::new()
        .recursive(true)
        .mode(0o700)
        .create(dir)?;
    let path = dir.join(format!("{sha256}.json"));
    let mut file = std::fs::OpenOptions::new()
        .write(true)
        .create(true)
        .truncate(true)
        .mode(0o600)
        .open(&path)?;
    file.write_all(&bytes)?;
    file.sync_all()?;
    let text = body
        .pointer("/content/0/text")
        .and_then(Value::as_str)
        .map_or_else(
            || String::from_utf8_lossy(&bytes).into_owned(),
            str::to_owned,
        );
    Ok(Handed {
        path,
        bytes: bytes.len(),
        sha256,
        preview: text.chars().take(PREVIEW_CHARS).collect(),
    })
}

fn size(message: &Value) -> io::Result<usize> {
    Ok(serde_json::to_vec(message)?.len().saturating_add(1))
}

fn fit(message: &mut Value, limit: usize) -> io::Result<()> {
    if size(message)? <= limit {
        return Ok(());
    }
    let channel = message["method"] == "notifications/claude/channel";
    let pointer = if channel {
        "/params/content"
    } else {
        "/params/data/words"
    };
    let Some(words) = message
        .pointer(pointer)
        .and_then(Value::as_str)
        .map(str::to_owned)
    else {
        return Err(io::Error::other(format!(
            "the host answer exceeds frame_bytes {limit}; its fields were not cut"
        )));
    };
    if channel {
        message["params"]["meta"]["truncated"] = Value::String("true".to_owned());
    }
    *message
        .pointer_mut(pointer)
        .ok_or_else(|| io::Error::other("the notification text field disappeared"))? =
        Value::String("…".to_owned());
    if size(message)? > limit {
        return Err(io::Error::other(format!(
            "the notification fields exceed frame_bytes {limit}; its fields were not cut"
        )));
    }
    let boundaries: Vec<_> = words
        .char_indices()
        .map(|(index, _)| index)
        .chain(std::iter::once(words.len()))
        .collect();
    let mut low = 0;
    let mut high = boundaries.len() - 1;
    while low < high {
        let mid = low + (high - low).div_ceil(2);
        *message
            .pointer_mut(pointer)
            .ok_or_else(|| io::Error::other("the notification text field disappeared"))? =
            Value::String(format!("{}…", &words[..boundaries[mid]]));
        if size(message)? <= limit {
            low = mid;
        } else {
            high = mid - 1;
        }
    }
    *message
        .pointer_mut(pointer)
        .ok_or_else(|| io::Error::other("the notification text field disappeared"))? =
        Value::String(format!("{}…", &words[..boundaries[low]]));
    Ok(())
}

#[cfg(test)]
#[path = "writer_tests.rs"]
mod tests;