use super::ServiceEffectError;
use super::lys_http::{LysHttp, segment};
use serde::Deserialize;
use std::sync::Arc;
use std::sync::mpsc;
#[derive(Clone, Debug, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ContextReading {
pub agent: String,
pub session: String,
pub context_percent: u64,
pub at_ms: i64,
}
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 {
#[must_use]
pub const fn new(http: Arc<dyn LysHttp>, bearer: String) -> Self {
Self { http, bearer }
}
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))
}
}
}
pub trait ContextFeed: Send {
fn next_reading(&mut self) -> Option<(String, Vec<u8>)>;
}
enum Signal {
Ready,
Stop,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Drained {
pub reported: u64,
pub refused: u64,
}
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
}
#[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 {
#[must_use]
pub fn new() -> Self {
let (sender, receiver) = mpsc::channel();
Self { sender, receiver }
}
#[must_use]
pub fn notifier(&self) -> liminal::channel::InboxNotifier {
let sender = self.sender.clone();
Arc::new(move || {
if sender.send(Signal::Ready).is_err() {
tracing::trace!("context watch closed; a wake was dropped");
}
})
}
}
#[derive(Debug)]
pub struct ContextWatch {
stop: mpsc::Sender<Signal>,
thread: Option<std::thread::JoinHandle<()>>,
}
impl ContextWatch {
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),
})
}
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();
}
}