johnandersen777.bsky.social/asyncai

Switch repository

Loading…

//! Push-driven watches over MCP resources (ADR 0005).
//!
//! One tokio task per watch waits on its subscription. Each update triggers
//! one resource read, events are extracted (event-log `seq` or state diff),
//! filtered by optional Python in the watch's own read-only REPL scope, and
//! debounced into a `WatchNotice`. Nothing polls: the task sleeps only for a
//! debounce window after an event, or for backoff after a failure.

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use tokio::time::Instant;

use super::{CodeModeRun, Failure, Shared};
use crate::resources::{
    default_notice, item_key, truncate, EventExtractor, ResourceSource, ResourceWatch, Shape,
    WatchSignal,
};

/// Watch limits (ADR 0005).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WatchLimits {
    pub max_watches: usize,
    /// Quiet time after the last event before a notice is sent.
    pub debounce: Duration,
    /// Longest a burst is held from its first event.
    pub max_wait: Duration,
    pub max_lines: usize,
    pub max_events: usize,
    pub resubscribe_attempts: u32,
    pub resubscribe_backoff: Duration,
}

impl Default for WatchLimits {
    fn default() -> Self {
        Self {
            max_watches: 5,
            debounce: Duration::from_secs(3),
            max_wait: Duration::from_secs(15),
            max_lines: 20,
            max_events: 100,
            resubscribe_attempts: 5,
            resubscribe_backoff: Duration::from_secs(1),
        }
    }
}

/// Reply token that publishes nothing.
pub const NO_REPLY: &str = "NO_REPLY";

/// Instruction appended to every watch notice (and recorded in sessions).
pub const WATCH_INSTRUCTION: &str = "This is a watch notification, not a user message. \
Tell the channel only if the change matters to what was asked; otherwise reply exactly NO_REPLY \
as your message text (not inside run_code).";

/// Debounced changes from one watch, for one agent turn.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WatchNotice {
    pub name: String,
    pub uri: String,
    pub lines: Vec<String>,
    pub dropped: usize,
}

impl WatchNotice {
    /// Message text for the model.
    pub fn render(&self) -> String {
        let mut body = self.lines.join("\n");
        if self.dropped > 0 {
            body.push_str(&format!("\n+{} more", self.dropped));
        }
        format!(
            "<watch-event name=\"{}\" uri=\"{}\">\n{}\n</watch-event>\n{WATCH_INSTRUCTION}",
            self.name, self.uri, body
        )
    }

    /// Source label for the orchestrator.
    pub fn source(&self) -> String {
        format!("watch:{}", self.name)
    }
}

struct WatchEntry {
    uri: String,
    on_event: Option<String>,
    shape: Shape,
    stats: Arc<WatchStats>,
    task: JoinHandle<()>,
}

#[derive(Default)]
struct WatchStats {
    events: AtomicU64,
    notices: AtomicU64,
}

/// Owns every watch of one agent. Watches outlive runs and REPL scopes.
pub struct WatchManager {
    resources: Arc<dyn ResourceSource>,
    filter_shared: Arc<Shared>,
    notices: mpsc::UnboundedSender<WatchNotice>,
    watches: Mutex<HashMap<String, WatchEntry>>,
    limits: WatchLimits,
    verbose: bool,
    /// Watch specs saved here on every start/stop, restored at startup.
    state_path: Option<PathBuf>,
    /// Saved specs that failed to restore for a transient reason (server
    /// down, subscribe failed). Kept in the state file until started again
    /// or stopped; invalid specs are dropped instead.
    unrestored: Mutex<Vec<WatchSpec>>,
}

/// A saved watch.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct WatchSpec {
    pub name: String,
    pub uri: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub on_event: Option<String>,
}

impl Drop for WatchManager {
    fn drop(&mut self) {
        if let Ok(watches) = self.watches.lock() {
            for entry in watches.values() {
                entry.task.abort();
            }
        }
    }
}

impl WatchManager {
    pub(crate) fn new(
        resources: Arc<dyn ResourceSource>,
        filter_shared: Arc<Shared>,
        notices: mpsc::UnboundedSender<WatchNotice>,
        limits: WatchLimits,
        verbose: bool,
        state_path: Option<PathBuf>,
    ) -> Self {
        Self {
            resources,
            filter_shared,
            notices,
            watches: Mutex::new(HashMap::new()),
            limits,
            verbose,
            state_path,
            unrestored: Mutex::new(Vec::new()),
        }
    }

    /// Starts the watches saved in the state file and sends one notice
    /// naming them. Returns how many were restored.
    pub async fn restore(&self) -> usize {
        let Some(path) = &self.state_path else {
            return 0;
        };
        let specs: Vec<WatchSpec> = match std::fs::read_to_string(path) {
            Ok(text) => match serde_json::from_str(&text) {
                Ok(specs) => specs,
                Err(error) => {
                    crate::log::emit(
                        "error",
                        "watch.restore_failed",
                        json!({"path": path, "error": error.to_string()}),
                    );
                    return 0;
                }
            },
            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return 0,
            Err(error) => {
                crate::log::emit(
                    "error",
                    "watch.restore_failed",
                    json!({"path": path, "error": error.to_string()}),
                );
                return 0;
            }
        };
        let mut lines = Vec::new();
        for spec in specs {
            match self
                .start(&spec.name, &spec.uri, spec.on_event.clone())
                .await
            {
                Ok(_) => lines.push(format!("restored {} {}", spec.name, spec.uri)),
                Err(error) if is_invalid_spec(&error) => lines.push(format!(
                    "dropped {} (invalid, removed from the state file): {error}",
                    spec.name
                )),
                Err(error) => {
                    lines.push(format!(
                        "could not restore {} yet (kept in the state file; watch_start it again, or watch_stop to forget it): {error}",
                        spec.name
                    ));
                    self.unrestored.lock().expect("unrestored lock").push(spec);
                }
            }
        }
        let restored = lines
            .iter()
            .filter(|line| line.starts_with("restored"))
            .count();
        // The file keeps running and transiently failed specs; invalid
        // specs are dropped. The notice tells the agent either way.
        self.save();
        crate::log::emit(
            "info",
            "watch.restored",
            json!({"path": path, "restored": restored, "lines": lines}),
        );
        if !lines.is_empty() {
            lines.push("changes made while the agent was down may be missing".to_string());
            let _ = self.notices.send(WatchNotice {
                name: "restore".to_string(),
                uri: path.display().to_string(),
                lines,
                dropped: 0,
            });
        }
        restored
    }

    /// Writes the current watch specs (temp file + rename).
    fn save(&self) {
        let Some(path) = &self.state_path else {
            return;
        };
        let mut specs: Vec<WatchSpec> = self
            .watches
            .lock()
            .expect("watch lock")
            .iter()
            .map(|(name, entry)| WatchSpec {
                name: name.clone(),
                uri: entry.uri.clone(),
                on_event: entry.on_event.clone(),
            })
            .collect();
        for spec in self.unrestored.lock().expect("unrestored lock").iter() {
            if !specs.iter().any(|known| known.name == spec.name) {
                specs.push(spec.clone());
            }
        }
        specs.sort_by(|a, b| a.name.cmp(&b.name));
        let temp = path.with_extension("tmp");
        let result = serde_json::to_vec_pretty(&specs)
            .map_err(std::io::Error::other)
            .and_then(|bytes| std::fs::write(&temp, bytes))
            .and_then(|_| std::fs::rename(&temp, path));
        if let Err(error) = result {
            crate::log::emit(
                "error",
                "watch.save_failed",
                json!({"path": path, "error": error.to_string()}),
            );
        }
    }

    pub async fn start(
        &self,
        name: &str,
        uri: &str,
        on_event: Option<String>,
    ) -> Result<Value, String> {
        {
            let watches = self.watches.lock().expect("watch lock");
            if watches.contains_key(name) {
                return Err(format!(
                    "watch {name:?} already exists; watch_stop it first"
                ));
            }
            if watches.len() >= self.limits.max_watches {
                return Err(format!(
                    "at most {} watches; watch_stop one first",
                    self.limits.max_watches
                ));
            }
        }

        let on_event = on_event.filter(|code| !code.trim().is_empty());
        let has_filter = on_event.is_some();
        let filter = match on_event.clone() {
            Some(code) if !code.trim().is_empty() => {
                let run = CodeModeRun::new(self.filter_shared.clone())
                    .map_err(|error| error.to_string())?;
                // Dry run on a sample event: syntax errors surface now, in the
                // model's turn, instead of on the first real event.
                let sample = json!({
                    "watch": name, "uri": uri, "type": "VALIDATE", "key": "sample",
                    "object": {"metadata": {"name": "sample", "namespace": "default"}},
                });
                if let (Err(Failure::Python(traceback)), _) =
                    run.execute(filter_program(&sample, &code), false).await
                {
                    // Missing keys are expected on the sample; syntax errors
                    // and undefined names are bugs on every event.
                    let last = traceback.lines().last().unwrap_or("error");
                    if traceback.contains("SyntaxError") {
                        return Err(format!("on_event does not parse: {last}"));
                    }
                    if last.contains("NameError") {
                        return Err(format!(
                            "on_event uses an undefined name: {last}. It runs in its own scope: only `event`, \
                             helpers (dig, items, names, count_by, to_json), json/re/asyncio, and read-only \
                             tools exist there; functions from run_code do not."
                        ));
                    }
                }
                Some((run, code))
            }
            _ => None,
        };

        let subscription = self
            .resources
            .watch(uri)
            .await
            .map_err(|error| error.to_string())?;
        let baseline = self
            .resources
            .read_resource(uri)
            .await
            .map_err(|error| error.to_string())?;
        let extractor = EventExtractor::baseline(&baseline, self.limits.max_events);
        let shape = extractor.shape();

        let stats = Arc::new(WatchStats::default());
        let task = tokio::spawn(run_watch(WatchTask {
            name: name.to_string(),
            uri: uri.to_string(),
            resources: self.resources.clone(),
            subscription,
            extractor,
            filter,
            notices: self.notices.clone(),
            limits: self.limits.clone(),
            stats: stats.clone(),
            verbose: self.verbose,
            last_lines: HashMap::new(),
        }));

        {
            let mut watches = self.watches.lock().expect("watch lock");
            if watches.contains_key(name) {
                task.abort();
                return Err(format!(
                    "watch {name:?} already exists; watch_stop it first"
                ));
            }
            watches.insert(
                name.to_string(),
                WatchEntry {
                    uri: uri.to_string(),
                    on_event,
                    shape,
                    stats,
                    task,
                },
            );
        }
        self.forget_unrestored(name);
        self.save();
        crate::log::emit(
            "info",
            "watch.started",
            json!({"watch": name, "uri": uri, "shape": shape, "filter": has_filter}),
        );
        Ok(json!({"name": name, "uri": uri, "shape": shape}))
    }

    fn forget_unrestored(&self, name: &str) -> bool {
        let mut unrestored = self.unrestored.lock().expect("unrestored lock");
        let before = unrestored.len();
        unrestored.retain(|spec| spec.name != name);
        unrestored.len() != before
    }

    pub fn stop(&self, name: &str) -> bool {
        if self.forget_unrestored(name) {
            self.save();
            return true;
        }
        let removed = self.watches.lock().expect("watch lock").remove(name);
        match removed {
            Some(entry) => {
                entry.task.abort();
                self.save();
                crate::log::emit(
                    "info",
                    "watch.stopped",
                    json!({"watch": name, "uri": entry.uri}),
                );
                true
            }
            None => false,
        }
    }

    pub fn list(&self) -> Value {
        let watches = self.watches.lock().expect("watch lock");
        let mut list: Vec<Value> = watches
            .iter()
            .map(|(name, entry)| {
                json!({
                    "name": name,
                    "uri": entry.uri,
                    "shape": entry.shape,
                    "events": entry.stats.events.load(Ordering::Relaxed),
                    "notices": entry.stats.notices.load(Ordering::Relaxed),
                    "running": !entry.task.is_finished(),
                })
            })
            .collect();
        for spec in self.unrestored.lock().expect("unrestored lock").iter() {
            list.push(json!({
                "name": spec.name,
                "uri": spec.uri,
                "running": false,
                "restore_failed": true,
            }));
        }
        list.sort_by(|a, b| a["name"].as_str().cmp(&b["name"].as_str()));
        Value::Array(list)
    }
}

struct WatchTask {
    name: String,
    uri: String,
    resources: Arc<dyn ResourceSource>,
    subscription: ResourceWatch,
    extractor: EventExtractor,
    filter: Option<(CodeModeRun, String)>,
    notices: mpsc::UnboundedSender<WatchNotice>,
    limits: WatchLimits,
    stats: Arc<WatchStats>,
    verbose: bool,
    /// Last notice line per item key; an identical next line is dropped.
    last_lines: HashMap<String, String>,
}

/// A saved spec that can never start (bad filter, over a limit), as
/// opposed to a transient failure (server down, subscribe failed).
fn is_invalid_spec(error: &str) -> bool {
    error.starts_with("on_event ")
        || error.starts_with("at most ")
        || error.contains("already exists")
}

/// Identity of the item an event is about, for line dedupe.
fn event_key(event: &Value) -> String {
    if let Some(key) = event.get("key").and_then(Value::as_str) {
        return key.to_string();
    }
    match event
        .get("object")
        .or_else(|| event.get("item"))
        .or_else(|| event.get("new"))
    {
        Some(object) => item_key(object),
        None => event.to_string(),
    }
}

/// Collapses repeated lines, keeping first-seen order: `line ×3`.
fn collapse(lines: Vec<String>) -> Vec<String> {
    let mut counted: Vec<(String, usize)> = Vec::new();
    for line in lines {
        match counted.iter_mut().find(|(known, _)| *known == line) {
            Some((_, count)) => *count += 1,
            None => counted.push((line, 1)),
        }
    }
    counted
        .into_iter()
        .map(|(line, count)| {
            if count > 1 {
                format!("{line} ×{count}")
            } else {
                line
            }
        })
        .collect()
}

impl WatchTask {
    fn log(&self, level: &str, event: &str, mut fields: Value) {
        if level == "info" && !self.verbose {
            return;
        }
        fields["watch"] = json!(self.name);
        fields["uri"] = json!(self.uri);
        crate::log::emit(level, event, fields);
    }

    /// Reads once and returns notice lines for the new events.
    async fn on_update(&mut self) -> Vec<String> {
        let content = match self.resources.read_resource(&self.uri).await {
            Ok(content) => content,
            Err(error) => return vec![format!("watch read failed: {error}")],
        };
        let events = self.extractor.extract(&content);
        let mut lines = Vec::new();
        for mut event in events {
            self.stats.events.fetch_add(1, Ordering::Relaxed);
            if let Some(object) = event.as_object_mut() {
                object.insert("watch".to_string(), json!(self.name));
                object.insert("uri".to_string(), json!(self.uri));
            }
            let mut line = self.filter_event(&event).await;
            let mut repeated = false;
            if let Some(text) = &line {
                let key = event_key(&event);
                repeated = self.last_lines.get(&key) == Some(text);
                if repeated {
                    line = None;
                } else {
                    self.last_lines.insert(key, text.clone());
                }
            }
            self.log(
                "info",
                "watch.event",
                json!({"type": event.get("type"), "kept": line.is_some(), "repeated": repeated}),
            );
            lines.extend(line);
        }
        lines
    }

    /// `None` drops the event.
    async fn filter_event(&self, event: &Value) -> Option<String> {
        let Some((run, code)) = &self.filter else {
            return Some(default_notice(event));
        };
        let (outcome, _) = run.execute(filter_program(event, code), false).await;
        match outcome {
            Ok(completion) => {
                let prints = completion.prints.trim();
                let value = completion.value.map(|value| match value {
                    Value::String(text) => text,
                    other => other.to_string(),
                });
                let line = match (prints.is_empty(), value) {
                    (true, None) => return None,
                    (true, Some(value)) => value,
                    (false, None) => prints.to_string(),
                    (false, Some(value)) => format!("{prints} {value}"),
                };
                Some(truncate(&line, 500))
            }
            Err(Failure::Python(traceback)) | Err(Failure::Host(traceback)) => {
                let last = traceback.lines().last().unwrap_or("error").to_string();
                Some(format!("on_event failed: {}", truncate(&last, 300)))
            }
        }
    }

    fn flush(&self, lines: &mut Vec<String>) {
        if lines.is_empty() {
            return;
        }
        let mut collapsed = collapse(std::mem::take(lines));
        let dropped = collapsed.len().saturating_sub(self.limits.max_lines);
        collapsed.truncate(self.limits.max_lines);
        let notice = WatchNotice {
            name: self.name.clone(),
            uri: self.uri.clone(),
            lines: collapsed,
            dropped,
        };
        self.stats.notices.fetch_add(1, Ordering::Relaxed);
        crate::log::emit(
            "info",
            "watch.notice",
            json!({"watch": self.name, "uri": self.uri, "lines": notice.lines, "dropped": dropped}),
        );
        let _ = self.notices.send(notice);
    }

    /// Resubscribes with exponential backoff; reconciles changes missed
    /// while disconnected. `false` when every attempt failed.
    async fn resubscribe(&mut self, reason: &str, lines: &mut Vec<String>) -> bool {
        let mut backoff = self.limits.resubscribe_backoff;
        for attempt in 1..=self.limits.resubscribe_attempts {
            self.log(
                "warn",
                "watch.resubscribing",
                json!({"reason": reason, "attempt": attempt, "backoff_ms": backoff.as_millis() as u64}),
            );
            tokio::time::sleep(backoff).await;
            backoff = backoff.saturating_mul(2);
            if let Ok(subscription) = self.resources.watch(&self.uri).await {
                self.subscription = subscription;
                lines.extend(self.on_update().await);
                return true;
            }
        }
        false
    }
}

async fn run_watch(mut task: WatchTask) {
    let mut lines: Vec<String> = Vec::new();
    let mut flush_at: Option<Instant> = None;
    let mut burst_started: Option<Instant> = None;
    task.log("info", "watch.subscribed", json!({}));
    loop {
        tokio::select! {
            signal = task.subscription.signals.recv() => {
                match signal {
                    Some(WatchSignal::Updated) => {
                        lines.extend(task.on_update().await);
                    }
                    Some(WatchSignal::Ended(reason)) => {
                        if !task.resubscribe(&reason, &mut lines).await {
                            lines.push(format!("watch stopped: subscription ended ({reason}) and resubscribing failed"));
                            task.flush(&mut lines);
                            task.log("error", "watch.failed", json!({"reason": reason}));
                            return;
                        }
                    }
                    None => {
                        let reason = "subscription channel closed".to_string();
                        if !task.resubscribe(&reason, &mut lines).await {
                            lines.push(format!("watch stopped: {reason}"));
                            task.flush(&mut lines);
                            task.log("error", "watch.failed", json!({"reason": reason}));
                            return;
                        }
                    }
                }
                // Trailing debounce: each update pushes the flush out by
                // `debounce`, but never past `max_wait` from the burst start.
                if !lines.is_empty() {
                    let now = Instant::now();
                    let started = *burst_started.get_or_insert(now);
                    flush_at = Some((now + task.limits.debounce).min(started + task.limits.max_wait));
                }
            }
            _ = sleep_until(flush_at), if flush_at.is_some() => {
                flush_at = None;
                burst_started = None;
                task.flush(&mut lines);
            }
        }
    }
}

/// Python that defines `event` and runs the filter.
fn filter_program(event: &Value, code: &str) -> String {
    let literal = serde_json::to_string(&event.to_string()).unwrap_or_else(|_| "\"{}\"".into());
    format!("event = json.loads({literal})\n{code}")
}

async fn sleep_until(deadline: Option<Instant>) {
    match deadline {
        Some(deadline) => tokio::time::sleep_until(deadline).await,
        None => std::future::pending().await,
    }
}