~ knotless Search/
↑↓ move↵ open⇧↵ all results

johnandersen777.bsky.social/asyncai

0
0
main
696 lines · 24.7 KB Raw History
//! Generic MCP resources and watches (ADR 0005).
//!
//! `ResourceSource` is the seam: `McpToolExecutor` implements it over any
//! MCP server; tests use fakes. Event extraction turns a resource read into
//! events, for both resource shapes:
//!
//! - **event-log** (`application/vnd.asyncai.events+json`): `{"events":
//!   [{"seq": n, ...}]}`, delivered by `seq`.
//! - **state** (anything else): diffed against the previous read.

use std::collections::{BTreeMap, BTreeSet};

use async_trait::async_trait;
use serde::Serialize;
use serde_json::{json, Map, Value};
use tokio::sync::mpsc;

use crate::tools::ToolError;

/// Mime type of event-log resources.
pub const EVENT_LOG_MIME: &str = "application/vnd.asyncai.events+json";

/// A resource or resource template advertised by a server.
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct ResourceInfo {
    /// The resource URI, or the RFC 6570 template when `template` is true.
    pub uri: String,
    pub template: bool,
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    /// The description's `Example:` line, if any.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub example: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mime_type: Option<String>,
    pub server: String,
    pub subscribable: bool,
}

/// The text after `Example:` on the description's first such line.
pub fn example_of(description: Option<&str>) -> Option<String> {
    description?.lines().find_map(|line| {
        line.trim()
            .strip_prefix("Example:")
            .map(|rest| rest.trim().to_string())
            .filter(|rest| !rest.is_empty())
    })
}

/// The text of one resource read.
#[derive(Clone, Debug, PartialEq)]
pub struct ResourceContent {
    pub text: String,
    pub mime_type: Option<String>,
}

/// What a watch subscription reports.
#[derive(Clone, Debug, PartialEq)]
pub enum WatchSignal {
    /// The resource changed; read it.
    Updated,
    /// The subscription ended (reason).
    Ended(String),
}

/// A live subscription. Dropping it unsubscribes.
pub struct ResourceWatch {
    pub signals: mpsc::Receiver<WatchSignal>,
    /// Dropped with the watch; cancels the subscription.
    pub guard: Box<dyn Send + Sync>,
}

#[async_trait]
pub trait ResourceSource: Send + Sync {
    async fn list_resources(&self) -> Result<Vec<ResourceInfo>, ToolError>;
    async fn read_resource(&self, uri: &str) -> Result<ResourceContent, ToolError>;
    async fn watch(&self, uri: &str) -> Result<ResourceWatch, ToolError>;
}

/// Read shape of a watched resource.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum Shape {
    State,
    EventLog,
}

impl Shape {
    pub fn of(content: &ResourceContent) -> Self {
        match content.mime_type.as_deref() {
            Some(mime) if mime.eq_ignore_ascii_case(EVENT_LOG_MIME) => Self::EventLog,
            _ => Self::State,
        }
    }
}

/// Turns successive reads of one resource into events.
#[derive(Clone, Debug)]
pub struct EventExtractor {
    shape: Shape,
    last_seq: Option<u64>,
    previous: Option<Value>,
    max_events: usize,
}

impl EventExtractor {
    /// Starts from a baseline read: nothing before it is delivered.
    pub fn baseline(content: &ResourceContent, max_events: usize) -> Self {
        let shape = Shape::of(content);
        let mut extractor = Self {
            shape,
            last_seq: None,
            previous: None,
            max_events,
        };
        match shape {
            Shape::EventLog => {
                extractor.last_seq = log_events(&content.text)
                    .ok()
                    .and_then(|events| events.iter().filter_map(event_seq).max())
                    .or(Some(0));
            }
            Shape::State => extractor.previous = Some(parse_state(&content.text)),
        }
        extractor
    }

    pub fn shape(&self) -> Shape {
        self.shape
    }

    /// Events since the previous read (at most `max_events`, then a
    /// `TRUNCATED` marker).
    pub fn extract(&mut self, content: &ResourceContent) -> Vec<Value> {
        let mut events = match Shape::of(content) {
            Shape::EventLog => {
                self.shape = Shape::EventLog;
                self.extract_log(&content.text)
            }
            Shape::State => {
                self.shape = Shape::State;
                let current = parse_state(&content.text);
                let events = match &self.previous {
                    Some(previous) => state_diff(previous, &current),
                    None => Vec::new(),
                };
                self.previous = Some(current);
                events
            }
        };
        events.retain_mut(enrich);
        if events.len() > self.max_events {
            let dropped = events.len() - self.max_events;
            events.truncate(self.max_events);
            events.push(json!({"type": "TRUNCATED", "dropped": dropped}));
        }
        events
    }

    fn extract_log(&mut self, text: &str) -> Vec<Value> {
        let mut events = match log_events(text) {
            Ok(events) => events,
            Err(message) => return vec![json!({"type": "ERROR", "message": message})],
        };
        let last = self.last_seq.unwrap_or(0);
        events.retain(|event| event_seq(event).is_some_and(|seq| seq > last));
        events.sort_by_key(|event| event_seq(event).unwrap_or(0));

        let mut out = Vec::new();
        if let Some(first) = events.first().and_then(event_seq) {
            if self.last_seq.is_some() && first > last + 1 {
                out.push(json!({"type": "GAP", "missed_from": last + 1, "missed_to": first - 1}));
            }
        }
        if let Some(max) = events.iter().filter_map(event_seq).max() {
            self.last_seq = Some(max);
        }
        out.extend(events);
        out
    }
}

/// Metadata that changes on every write; ignored when computing `changed`.
const VOLATILE_PATHS: &[&str] = &[
    "metadata.resourceVersion",
    "metadata.generation",
    "metadata.managedFields",
];
/// Depth of the dotted paths in `changed`.
const CHANGED_DEPTH: usize = 3;

/// Adds `old` (previous object or null) and, when both sides are known,
/// `changed` (dotted paths). Returns false for a `MODIFIED` event that
/// changed nothing but volatile metadata.
fn enrich(event: &mut Value) -> bool {
    let Some(map) = event.as_object_mut() else {
        return true;
    };
    if !matches!(
        map.get("type").and_then(Value::as_str),
        Some("ADDED" | "MODIFIED" | "DELETED")
    ) {
        return true;
    }
    let new = map
        .get("new")
        .or_else(|| map.get("object"))
        .or_else(|| map.get("item"))
        .cloned();
    map.entry("old").or_insert(Value::Null);
    if map.get("type").and_then(Value::as_str) != Some("MODIFIED") {
        return true;
    }
    let (Some(old), Some(new)) = (map.get("old").filter(|old| !old.is_null()), new) else {
        return true;
    };
    let mut changed = Vec::new();
    changed_paths(old, &new, "", 0, &mut changed);
    if changed.is_empty() {
        return false;
    }
    map.insert("changed".to_string(), json!(changed));
    true
}

/// Dotted paths where `old` and `new` differ, down to `CHANGED_DEPTH`.
pub fn changed_paths(old: &Value, new: &Value, prefix: &str, depth: usize, out: &mut Vec<String>) {
    if VOLATILE_PATHS.contains(&prefix) || old == new {
        return;
    }
    match (old, new) {
        (Value::Object(before), Value::Object(after)) if depth < CHANGED_DEPTH => {
            let keys: BTreeSet<&String> = before.keys().chain(after.keys()).collect();
            for key in keys {
                let path = if prefix.is_empty() {
                    key.clone()
                } else {
                    format!("{prefix}.{key}")
                };
                let null = Value::Null;
                changed_paths(
                    before.get(key).unwrap_or(&null),
                    after.get(key).unwrap_or(&null),
                    &path,
                    depth + 1,
                    out,
                );
            }
        }
        _ => out.push(if prefix.is_empty() {
            ".".to_string()
        } else {
            prefix.to_string()
        }),
    }
}

fn event_seq(event: &Value) -> Option<u64> {
    event.get("seq").and_then(Value::as_u64)
}

fn log_events(text: &str) -> Result<Vec<Value>, String> {
    let value: Value = serde_json::from_str(text)
        .map_err(|error| format!("event-log resource is not JSON: {error}"))?;
    value
        .get("events")
        .and_then(Value::as_array)
        .cloned()
        .ok_or_else(|| "event-log resource has no `events` array".to_string())
}

fn parse_state(text: &str) -> Value {
    serde_json::from_str(text).unwrap_or_else(|_| Value::String(text.to_string()))
}

/// Identity of a list item: `metadata.uid`, `uid`, `id`,
/// `metadata.namespace/metadata.name`, `name`; else the item itself.
pub fn item_key(item: &Value) -> String {
    let pointer = |path: &str| item.pointer(path).and_then(scalar_text);
    if let Some(uid) = pointer("/metadata/uid").or_else(|| pointer("/uid")) {
        return uid;
    }
    if let Some(id) = pointer("/id") {
        return id;
    }
    if let Some(name) = pointer("/metadata/name") {
        return match pointer("/metadata/namespace") {
            Some(namespace) => format!("{namespace}/{name}"),
            None => name,
        };
    }
    if let Some(name) = pointer("/name") {
        return name;
    }
    item.to_string()
}

fn scalar_text(value: &Value) -> Option<String> {
    match value {
        Value::String(text) => Some(text.clone()),
        Value::Number(number) => Some(number.to_string()),
        _ => None,
    }
}

/// Diff of two state reads as events.
pub fn state_diff(old: &Value, new: &Value) -> Vec<Value> {
    match (old, new) {
        (Value::Array(old_items), Value::Array(new_items)) => {
            let index = |items: &[Value]| -> BTreeMap<String, Value> {
                items
                    .iter()
                    .map(|item| (item_key(item), item.clone()))
                    .collect()
            };
            let (before, after) = (index(old_items), index(new_items));
            let mut events = Vec::new();
            for (key, item) in &after {
                match before.get(key) {
                    None => events.push(json!({"type": "ADDED", "key": key, "item": item})),
                    Some(previous) if previous != item => events.push(
                        json!({"type": "MODIFIED", "key": key, "old": previous, "new": item}),
                    ),
                    Some(_) => {}
                }
            }
            for (key, item) in &before {
                if !after.contains_key(key) {
                    events.push(json!({"type": "DELETED", "key": key, "item": item}));
                }
            }
            events
        }
        (Value::Object(before), Value::Object(after)) => {
            let keys: BTreeSet<&String> = before.keys().chain(after.keys()).collect();
            keys.into_iter()
                .filter_map(|key| match (before.get(key), after.get(key)) {
                    (None, Some(value)) => Some(json!({"type": "ADDED", "key": key, "new": value})),
                    (Some(value), None) => {
                        Some(json!({"type": "DELETED", "key": key, "old": value}))
                    }
                    (Some(old), Some(new)) if old != new => {
                        Some(json!({"type": "MODIFIED", "key": key, "old": old, "new": new}))
                    }
                    _ => None,
                })
                .collect()
        }
        (old, new) if old != new => vec![json!({"type": "CHANGED", "old": old, "new": new})],
        _ => Vec::new(),
    }
}

/// Default one-line summary of an event when no filter is given.
pub fn default_notice(event: &Value) -> String {
    let kind = event.get("type").and_then(Value::as_str).unwrap_or("EVENT");
    let object = event
        .get("object")
        .or_else(|| event.get("item"))
        .or_else(|| event.get("new"));
    if let Some(metadata) = object.and_then(|object| object.get("metadata")) {
        let name = metadata.get("name").and_then(Value::as_str).unwrap_or("?");
        let line = match metadata.get("namespace").and_then(Value::as_str) {
            Some(namespace) => format!("{kind} {namespace}/{name}"),
            None => format!("{kind} {name}"),
        };
        return line + &change_detail(event, object);
    }
    if let Some(key) = event.get("key").and_then(Value::as_str) {
        return format!("{kind} {key}") + &change_detail(event, object);
    }
    let mut rest = Map::new();
    if let Some(object) = event.as_object() {
        for (key, value) in object {
            if !matches!(key.as_str(), "watch" | "uri") {
                rest.insert(key.clone(), value.clone());
            }
        }
    }
    truncate(&Value::Object(rest).to_string(), 300)
}

/// ` (status.phase: Deciding→Published, spec)` for up to three changed paths.
fn change_detail(event: &Value, new: Option<&Value>) -> String {
    const SHOWN: usize = 3;
    let Some(changed) = event.get("changed").and_then(Value::as_array) else {
        return String::new();
    };
    let old = event.get("old");
    let at = |root: Option<&Value>, path: &str| -> Option<String> {
        let pointer = if path == "." {
            String::new()
        } else {
            format!("/{}", path.replace('.', "/"))
        };
        // Non-scalars are shown too (compact, truncated): without the
        // values, adding and removing a label render the same line and the
        // second is deduped away (seen live).
        match root?.pointer(&pointer)? {
            value @ (Value::Object(_) | Value::Array(_)) => Some(truncate(&value.to_string(), 60)),
            Value::String(text) => Some(truncate(text, 40)),
            other => Some(other.to_string()),
        }
    };
    let mut parts: Vec<String> = changed
        .iter()
        .filter_map(Value::as_str)
        .take(SHOWN)
        .map(|path| {
            let label = if path == "." {
                String::new()
            } else {
                format!("{path}: ")
            };
            let non_scalar = |root: Option<&Value>| {
                let pointer = format!("/{}", path.replace('.', "/"));
                root.and_then(|root| root.pointer(&pointer))
                    .is_some_and(|value| value.is_object() || value.is_array())
            };
            if path != "." && non_scalar(old) && non_scalar(new) {
                // Truncated arrays/objects on both sides read the same; the
                // path says enough.
                return path.to_string();
            }
            match (at(old, path), at(new, path)) {
                (Some(before), Some(after)) => format!("{label}{before}→{after}"),
                (None, Some(after)) => format!("{label}→{after}"),
                (Some(before), None) => format!("{label}{before}→"),
                (None, None) => path.to_string(),
            }
        })
        .collect();
    if changed.len() > SHOWN {
        parts.push(format!("+{}", changed.len() - SHOWN));
    }
    format!(" ({})", parts.join(", "))
}

pub fn truncate(text: &str, max_chars: usize) -> String {
    if text.chars().count() <= max_chars {
        return text.to_string();
    }
    let mut out: String = text.chars().take(max_chars).collect();
    out.push('…');
    out
}

/// RFC 6570 level-1 matching plus `{?a,b}` query expansions: `{var}`
/// matches one or more characters other than `/` and `?`; a `{?...}`
/// suffix accepts any query string.
pub fn template_matches(template: &str, uri: &str) -> bool {
    let (template, query_ok) = match template.find("{?") {
        Some(index) => (&template[..index], true),
        None => (template, false),
    };
    let (path, query) = match uri.split_once('?') {
        Some((path, query)) => (path, Some(query)),
        None => (uri, None),
    };
    if query.is_some() && !query_ok {
        return false;
    }
    matches_segments(template, path)
}

fn matches_segments(template: &str, text: &str) -> bool {
    match template.find('{') {
        None => template == text,
        Some(open) => {
            let (literal, rest) = template.split_at(open);
            let Some(text) = text.strip_prefix(literal) else {
                return false;
            };
            let Some(close) = rest.find('}') else {
                return false;
            };
            let after = &rest[close + 1..];
            let limit = text.find(['/', '?']).unwrap_or(text.len());
            (1..=limit).any(|end| matches_segments(after, &text[end..]))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn content(text: Value, mime: Option<&str>) -> ResourceContent {
        ResourceContent {
            text: text.to_string(),
            mime_type: mime.map(str::to_string),
        }
    }

    #[test]
    fn state_lists_diff_by_identity() {
        let old = json!([{"metadata": {"name": "a", "namespace": "default"}, "v": 1}, {"metadata": {"name": "b", "namespace": "default"}}]);
        let new = json!([{"metadata": {"name": "a", "namespace": "default"}, "v": 2}, {"metadata": {"name": "c", "namespace": "default"}}]);
        let events = state_diff(&old, &new);
        let summary = events
            .iter()
            .map(|event| {
                format!(
                    "{} {}",
                    event["type"].as_str().unwrap(),
                    event["key"].as_str().unwrap()
                )
            })
            .collect::<Vec<_>>();
        assert_eq!(
            summary,
            ["MODIFIED default/a", "ADDED default/c", "DELETED default/b"]
        );
    }

    #[test]
    fn state_objects_and_text_diff() {
        let events = state_diff(&json!({"a": 1, "b": 2}), &json!({"a": 1, "b": 3, "c": 4}));
        assert_eq!(events.len(), 2);
        assert_eq!(
            events[0],
            json!({"type": "MODIFIED", "key": "b", "old": 2, "new": 3})
        );
        assert_eq!(events[1]["type"], "ADDED");
        assert_eq!(
            state_diff(&json!("x"), &json!("y")),
            vec![json!({"type": "CHANGED", "old": "x", "new": "y"})]
        );
        assert!(state_diff(&json!([1]), &json!([1])).is_empty());
    }

    #[test]
    fn state_extractor_uses_baseline() {
        let mut extractor =
            EventExtractor::baseline(&content(json!({"a": 1}), Some("application/json")), 100);
        assert_eq!(extractor.shape(), Shape::State);
        let events = extractor.extract(&content(json!({"a": 2}), Some("application/json")));
        assert_eq!(
            events,
            vec![json!({"type": "MODIFIED", "key": "a", "old": 1, "new": 2, "changed": ["."]})]
        );
        assert_eq!(default_notice(&events[0]), "MODIFIED a (1→2)");
        assert!(extractor
            .extract(&content(json!({"a": 2}), None))
            .is_empty());
    }

    #[test]
    fn event_log_delivers_new_seqs_and_reports_gaps() {
        let log = |seqs: &[u64]| {
            content(
                json!({"events": seqs.iter().map(|seq| json!({"seq": seq, "type": "ADDED"})).collect::<Vec<_>>()}),
                Some(EVENT_LOG_MIME),
            )
        };
        let mut extractor = EventExtractor::baseline(&log(&[1, 2]), 100);
        assert_eq!(extractor.shape(), Shape::EventLog);
        let events = extractor.extract(&log(&[1, 2, 3]));
        assert_eq!(events.len(), 1);
        assert_eq!(events[0]["seq"], 3);
        let gapped = extractor.extract(&log(&[6, 7]));
        assert_eq!(
            gapped[0],
            json!({"type": "GAP", "missed_from": 4, "missed_to": 5})
        );
        assert_eq!(gapped.len(), 3);
        assert!(extractor.extract(&log(&[6, 7])).is_empty());
    }

    #[test]
    fn empty_event_log_baseline_delivers_first_event() {
        let mut extractor =
            EventExtractor::baseline(&content(json!({"events": []}), Some(EVENT_LOG_MIME)), 100);
        let events = extractor.extract(&content(
            json!({"events": [{"seq": 1, "type": "ADDED"}]}),
            Some(EVENT_LOG_MIME),
        ));
        assert_eq!(events.len(), 1);
    }

    #[test]
    fn extraction_is_capped() {
        let mut extractor = EventExtractor::baseline(&content(json!([]), None), 2);
        let events = extractor.extract(&content(json!([1, 2, 3, 4]), None));
        assert_eq!(events.len(), 3);
        assert_eq!(events[2], json!({"type": "TRUNCATED", "dropped": 2}));
    }

    #[test]
    fn default_notices() {
        assert_eq!(
            default_notice(
                &json!({"type": "ADDED", "object": {"metadata": {"name": "x", "namespace": "default"}}})
            ),
            "ADDED default/x"
        );
        assert_eq!(
            default_notice(&json!({"type": "MODIFIED", "key": "k", "old": 1, "new": 2})),
            "MODIFIED k"
        );
        assert_eq!(
            default_notice(&json!({"type": "GAP", "watch": "w", "missed_from": 4})),
            r#"{"missed_from":4,"type":"GAP"}"#
        );
    }

    #[test]
    fn event_log_events_gain_old_and_changed_and_noise_is_dropped() {
        let bid = |phase: &str, version: &str| json!({"metadata": {"name": "bid-1", "resourceVersion": version}, "status": {"phase": phase}});
        let log = content(
            json!({"events": [
                {"seq": 1, "type": "ADDED", "object": bid("Deciding", "1")},
                {"seq": 2, "type": "MODIFIED", "object": bid("Deciding", "2"), "old": bid("Deciding", "1")},
                {"seq": 3, "type": "MODIFIED", "object": bid("Published", "3"), "old": bid("Deciding", "2")},
                {"seq": 4, "type": "MODIFIED", "object": bid("Won", "4")},
            ]}),
            Some(EVENT_LOG_MIME),
        );
        let mut extractor =
            EventExtractor::baseline(&content(json!({"events": []}), Some(EVENT_LOG_MIME)), 100);
        let events = extractor.extract(&log);
        let seqs: Vec<u64> = events.iter().filter_map(event_seq).collect();
        assert_eq!(seqs, [1, 3, 4], "resourceVersion-only change dropped");
        assert_eq!(events[0]["old"], Value::Null);
        assert_eq!(events[1]["changed"], json!(["status.phase"]));
        assert_eq!(
            default_notice(&events[1]),
            "MODIFIED bid-1 (status.phase: Deciding→Published)"
        );
        assert_eq!(events[2].get("changed"), None, "no old, no changed");
        assert_eq!(default_notice(&events[2]), "MODIFIED bid-1");
    }

    #[test]
    fn label_add_and_remove_render_different_lines() {
        let bid = |labels: Option<Value>| {
            let mut metadata = json!({"name": "bid-1"});
            if let Some(labels) = labels {
                metadata["labels"] = labels;
            }
            json!({"metadata": metadata})
        };
        let mut added =
            json!({"type": "MODIFIED", "object": bid(Some(json!({"t": "1"}))), "old": bid(None)});
        let mut removed =
            json!({"type": "MODIFIED", "object": bid(None), "old": bid(Some(json!({"t": "1"})))});
        assert!(enrich(&mut added) && enrich(&mut removed));
        assert_eq!(
            default_notice(&added),
            r#"MODIFIED bid-1 (metadata.labels: →{"t":"1"})"#
        );
        assert_eq!(
            default_notice(&removed),
            r#"MODIFIED bid-1 (metadata.labels: {"t":"1"}→)"#
        );
    }

    #[test]
    fn non_scalar_on_both_sides_shows_only_the_path() {
        let bid = |message: &str| json!({"metadata": {"name": "b"}, "status": {"conditions": [{"message": message}], "phase": "Won"}});
        let mut event = json!({"type": "MODIFIED", "object": bid("new"), "old": bid("old")});
        assert!(enrich(&mut event));
        assert_eq!(default_notice(&event), "MODIFIED b (status.conditions)");
    }

    #[test]
    fn changed_paths_stop_at_depth_and_skip_volatile_metadata() {
        let mut out = Vec::new();
        changed_paths(
            &json!({"metadata": {"generation": 1, "labels": {"a": "1"}}, "spec": {"a": {"b": {"c": 1}}}}),
            &json!({"metadata": {"generation": 2, "labels": {"a": "2"}}, "spec": {"a": {"b": {"c": 2}}}}),
            "",
            0,
            &mut out,
        );
        assert_eq!(out, ["metadata.labels.a", "spec.a.b"]);
    }

    #[test]
    fn templates_match() {
        let template = "kcp-events://watch/{workspace}/{resource}{?namespace,kcp}";
        assert!(template_matches(
            template,
            "kcp-events://watch/root:app/configmaps"
        ));
        assert!(template_matches(
            template,
            "kcp-events://watch/root:app/configmaps?namespace=*"
        ));
        assert!(!template_matches(template, "kcp-events://watch/root:app"));
        assert!(!template_matches(template, "memo://notes"));
        assert!(template_matches("memo://notes", "memo://notes"));
        assert!(!template_matches("memo://notes", "memo://notes?x=1"));
    }
}