liminal_server/services/

context_watch.rs

//! The context watch: a watcher over the `context-watch` subject that
//! reports each context reading to Lys, which decides what a crossed budget
//! does (ACCESS-003 R1).
//!
//! A runner publishes its session's context reading on `context-watch` (a
//! hot `publish` decided from its pass like any other); a granted `call` on
//! the service carries a reading the same way. Each reading is reported to
//! Lys's `POST /agents/{agent}/usage` as its `context_percent`, under a
//! stable event id (the message's own id, or the call's operation id), so a
//! reading reported again charges nothing. Lys keeps each crossing and asks
//! the session's runner to compact, give notice or stop; liminal decides
//! none of it and keeps no threshold.
//!
//! The watcher is told, never polls: the subscription's inbox notifier wakes
//! one watcher thread, which drains every ready reading and sleeps on its
//! channel again. Stopping the watch (or dropping it) ends the thread.

use super::ServiceEffectError;
use super::lys_http::{LysHttp, segment};
use serde::Deserialize;
use std::sync::Arc;
use std::sync::mpsc;

/// One context reading, as a runner publishes it.
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ContextReading {
    /// The agent (seat) whose session it is, as Lys names it.
    pub agent: String,
    /// The session the reading is of.
    pub session: String,
    /// The session's context, in percent of its window.
    pub context_percent: u64,
    /// When it was measured, in milliseconds since the Unix epoch.
    pub at_ms: i64,
}

/// Reports context readings to Lys over [`LysHttp`].
pub struct ContextReporter {
    http: Arc<dyn LysHttp>,
    bearer: String,
}

impl std::fmt::Debug for ContextReporter {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("ContextReporter")
            .field("http", &self.http)
            .finish_non_exhaustive()
    }
}

impl ContextReporter {
    /// Reports through `http` with `bearer`, this liminal's own Lys pass,
    /// never shown by `Debug`.
    #[must_use]
    pub const fn new(http: Arc<dyn LysHttp>, bearer: String) -> Self {
        Self { http, bearer }
    }

    /// Reports one reading's bytes to Lys under `event`.
    ///
    /// # Errors
    /// The bytes are not a reading, the event id is empty, or Lys refused or
    /// could not be asked, each by name.
    pub fn report(&self, payload: &[u8], event: &str) -> Result<(), ServiceEffectError> {
        if event.trim().is_empty() {
            return Err(ServiceEffectError(
                "a context reading needs an event id".to_owned(),
            ));
        }
        let reading: ContextReading = serde_json::from_slice(payload)
            .map_err(|error| ServiceEffectError(format!("not a context reading: {error}")))?;
        let route = format!("agents/{}/usage", segment(&reading.agent));
        let body = serde_json::json!({
            "event": event,
            "at_ms": reading.at_ms,
            "session": reading.session,
            "context_percent": reading.context_percent,
        });
        let answer = self.http.post(&self.bearer, &route, &body)?;
        if answer.succeeded() {
            Ok(())
        } else {
            Err(answer.refused(&route))
        }
    }
}

/// The ready readings on the watched subject, each with its message id.
pub trait ContextFeed: Send {
    /// The next ready reading without blocking, or `None` when none is
    /// ready.
    fn next_reading(&mut self) -> Option<(String, Vec<u8>)>;
}

enum Signal {
    Ready,
    Stop,
}

/// What a drain came to.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Drained {
    /// Readings Lys accepted.
    pub reported: u64,
    /// Readings refused or not reported, each logged by name.
    pub refused: u64,
}

/// Drains every ready reading from `feed` into `reporter`.
pub fn drain(feed: &mut dyn ContextFeed, reporter: &ContextReporter) -> Drained {
    let mut drained = Drained::default();
    while let Some((event, payload)) = feed.next_reading() {
        match reporter.report(&payload, &event) {
            Ok(()) => drained.reported = drained.reported.saturating_add(1),
            Err(error) => {
                drained.refused = drained.refused.saturating_add(1);
                tracing::warn!(event = %event, error = %error, "context reading not reported to Lys");
            }
        }
    }
    drained
}

/// The wake the watched subscription's inbox fires on every admitted
/// envelope, and the receiving end the watcher sleeps on.
#[derive(Debug)]
pub struct ContextWake {
    sender: mpsc::Sender<Signal>,
    receiver: mpsc::Receiver<Signal>,
}

impl std::fmt::Debug for Signal {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(match self {
            Self::Ready => "Ready",
            Self::Stop => "Stop",
        })
    }
}

impl Default for ContextWake {
    fn default() -> Self {
        Self::new()
    }
}

impl ContextWake {
    /// A fresh wake.
    #[must_use]
    pub fn new() -> Self {
        let (sender, receiver) = mpsc::channel();
        Self { sender, receiver }
    }

    /// The notifier to install on the subscription's inbox before it is
    /// registered, so no reading arrives without a wake.
    #[must_use]
    pub fn notifier(&self) -> liminal::channel::InboxNotifier {
        let sender = self.sender.clone();
        Arc::new(move || {
            // A closed watcher has nothing left to wake.
            if sender.send(Signal::Ready).is_err() {
                tracing::trace!("context watch closed; a wake was dropped");
            }
        })
    }
}

/// The running watcher; stopping or dropping it ends the thread.
#[derive(Debug)]
pub struct ContextWatch {
    stop: mpsc::Sender<Signal>,
    thread: Option<std::thread::JoinHandle<()>>,
}

impl ContextWatch {
    /// Starts the watcher over `feed`, woken by `wake`, reporting through
    /// `reporter`. It drains once at start, then once per wake.
    ///
    /// # Errors
    /// The watcher thread could not be spawned.
    pub fn start<F: ContextFeed + 'static>(
        mut feed: F,
        wake: ContextWake,
        reporter: Arc<ContextReporter>,
    ) -> Result<Self, ServiceEffectError> {
        let ContextWake { sender, receiver } = wake;
        let thread = std::thread::Builder::new()
            .name("liminal-context-watch".to_owned())
            .spawn(move || {
                drain(&mut feed, &reporter);
                while matches!(receiver.recv(), Ok(Signal::Ready)) {
                    drain(&mut feed, &reporter);
                }
            })
            .map_err(|error| {
                ServiceEffectError(format!("the context watch did not start: {error}"))
            })?;
        Ok(Self {
            stop: sender,
            thread: Some(thread),
        })
    }

    /// Stops the watcher and waits for its thread.
    pub fn stop(&mut self) {
        if self.stop.send(Signal::Stop).is_err() {
            tracing::debug!("context watch thread already stopped");
        }
        if let Some(thread) = self.thread.take()
            && thread.join().is_err()
        {
            tracing::warn!("the context watch thread panicked");
        }
    }
}

impl Drop for ContextWatch {
    fn drop(&mut self) {
        self.stop();
    }
}