johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use async_trait::async_trait;
use asyncai::code_mode::{CodeModeOptions, WatchLimits, WatchNotice};
use asyncai::resources::{
    ResourceContent, ResourceInfo, ResourceSource, ResourceWatch, WatchSignal, EVENT_LOG_MIME,
};
use asyncai::{CodeModeToolExecutor, ToolCallOutput, ToolDefinition, ToolError, ToolExecutor};
use serde_json::{json, Value};
use tokio::sync::mpsc;

const WAIT: Duration = Duration::from_secs(20);

/// In-memory resources that push updates to subscribers.
#[derive(Default)]
struct FakeResources {
    contents: Mutex<HashMap<String, ResourceContent>>,
    subscribers: Mutex<HashMap<String, Vec<mpsc::Sender<WatchSignal>>>>,
    subscribe_count: Mutex<usize>,
    /// When set, `watch` fails (server down).
    fail_watch: Mutex<bool>,
}

impl FakeResources {
    fn set(&self, uri: &str, text: Value, mime: &str) {
        self.contents.lock().unwrap().insert(
            uri.to_string(),
            ResourceContent {
                text: text.to_string(),
                mime_type: Some(mime.to_string()),
            },
        );
    }

    async fn push(&self, uri: &str, text: Value, mime: &str) {
        self.set(uri, text, mime);
        let senders = self
            .subscribers
            .lock()
            .unwrap()
            .get(uri)
            .cloned()
            .unwrap_or_default();
        for sender in senders {
            let _ = sender.send(WatchSignal::Updated).await;
        }
    }

    async fn end(&self, uri: &str) {
        let senders = self
            .subscribers
            .lock()
            .unwrap()
            .remove(uri)
            .unwrap_or_default();
        for sender in senders {
            let _ = sender
                .send(WatchSignal::Ended("server restarted".to_string()))
                .await;
        }
    }
}

#[async_trait]
impl ResourceSource for FakeResources {
    async fn list_resources(&self) -> Result<Vec<ResourceInfo>, ToolError> {
        Ok(vec![ResourceInfo {
            uri: "memo://notes".to_string(),
            template: false,
            name: "notes".to_string(),
            description: None,
            example: None,
            mime_type: Some("application/json".to_string()),
            server: "fake".to_string(),
            subscribable: true,
        }])
    }

    async fn read_resource(&self, uri: &str) -> Result<ResourceContent, ToolError> {
        self.contents
            .lock()
            .unwrap()
            .get(uri)
            .cloned()
            .ok_or_else(|| ToolError::new(format!("unknown resource {uri}")))
    }

    async fn watch(&self, uri: &str) -> Result<ResourceWatch, ToolError> {
        if *self.fail_watch.lock().unwrap() {
            return Err(ToolError::new("MCP server is down"));
        }
        let (sender, signals) = mpsc::channel(16);
        self.subscribers
            .lock()
            .unwrap()
            .entry(uri.to_string())
            .or_default()
            .push(sender);
        *self.subscribe_count.lock().unwrap() += 1;
        Ok(ResourceWatch {
            signals,
            guard: Box::new(()),
        })
    }
}

/// One read-only tool and one mutating tool.
struct Tools;

#[async_trait]
impl ToolExecutor for Tools {
    async fn list_tools(&self) -> Result<Vec<ToolDefinition>, ToolError> {
        let mut lookup = ToolDefinition::new(
            "owner_of",
            None::<String>,
            json!({"type": "object", "properties": {"key": {"type": "string"}}, "required": ["key"]}),
        );
        lookup.read_only = Some(true);
        let mut delete = ToolDefinition::new(
            "delete_note",
            None::<String>,
            json!({"type": "object", "properties": {"key": {"type": "string"}}, "required": ["key"]}),
        );
        delete.destructive = Some(true);
        Ok(vec![lookup, delete])
    }

    async fn call_tool(&self, name: &str, arguments: Value) -> Result<ToolCallOutput, ToolError> {
        match name {
            "owner_of" => Ok(ToolCallOutput::text(
                json!({"owner": format!("owner-of-{}", arguments["key"].as_str().unwrap_or(""))})
                    .to_string(),
            )),
            _ => Ok(ToolCallOutput::text("deleted")),
        }
    }
}

struct Harness {
    resources: Arc<FakeResources>,
    executor: CodeModeToolExecutor,
    notices: mpsc::UnboundedReceiver<WatchNotice>,
}

fn fast_limits() -> WatchLimits {
    WatchLimits {
        debounce: Duration::from_millis(50),
        resubscribe_backoff: Duration::from_millis(10),
        ..WatchLimits::default()
    }
}

async fn harness() -> Harness {
    harness_with(fast_limits(), None, Arc::new(FakeResources::default())).await
}

async fn harness_with(
    watch_limits: WatchLimits,
    watch_state: Option<std::path::PathBuf>,
    resources: Arc<FakeResources>,
) -> Harness {
    if resources.contents.lock().unwrap().is_empty() {
        resources.set("memo://notes", json!({"a": "1"}), "application/json");
    }
    let executor = CodeModeToolExecutor::with_options(
        Arc::new(Tools),
        CodeModeOptions {
            resources: Some(resources.clone()),
            watch_limits,
            watch_state,
            ..CodeModeOptions::default()
        },
    )
    .await
    .unwrap();
    let notices = executor.take_watch_notices().unwrap();
    Harness {
        resources,
        executor,
        notices,
    }
}

impl Harness {
    async fn run(&self, code: &str) -> ToolCallOutput {
        self.executor
            .call_tool("run_code", json!({"code": code}))
            .await
            .unwrap()
    }

    async fn next_notice(&mut self) -> WatchNotice {
        tokio::time::timeout(WAIT, self.notices.recv())
            .await
            .expect("no watch notice")
            .expect("notice channel closed")
    }
}

#[tokio::test]
async fn catalog_exposes_resource_and_watch_builtins() {
    let harness = harness().await;
    let description = harness.executor.list_tools().await.unwrap()[0]
        .description
        .clone()
        .unwrap();
    for stub in [
        "async def list_resources() -> list[dict[str, Any]]:",
        "async def read_resource(*, uri: str) -> Any:",
        "async def watch_start(*, name: str, uri: str, on_event: str | None = None) -> dict[str, Any]:",
        "async def watch_stop(*, name: str) -> bool:",
        "async def watch_list() -> list[dict[str, Any]]:",
    ] {
        assert!(description.contains(stub), "missing {stub}");
    }

    let listed = harness
        .run("[r['uri'] for r in await list_resources() if r['subscribable']]")
        .await;
    assert_eq!(listed.content, r#"["memo://notes"]"#);
    assert_eq!(
        harness
            .run("(await read_resource(uri='memo://notes'))['a']")
            .await
            .content,
        "1"
    );
}

#[tokio::test]
async fn state_watch_pushes_a_debounced_notice() {
    let mut harness = harness().await;
    let started = harness
        .run("await watch_start(name='notes', uri='memo://notes')")
        .await;
    assert!(!started.is_error, "{}", started.content);
    assert_eq!(
        serde_json::from_str::<Value>(&started.content).unwrap(),
        json!({"name": "notes", "uri": "memo://notes", "shape": "state"})
    );

    harness
        .resources
        .push(
            "memo://notes",
            json!({"a": "1", "b": "2"}),
            "application/json",
        )
        .await;
    harness
        .resources
        .push("memo://notes", json!({"b": "3"}), "application/json")
        .await;

    let notice = harness.next_notice().await;
    assert_eq!(notice.name, "notes");
    // Both changes landed before the first read: a state resource coalesces
    // them into one diff against the last read.
    assert_eq!(notice.lines, ["DELETED a", "ADDED b"]);
    let text = notice.render();
    assert!(text.starts_with("<watch-event name=\"notes\" uri=\"memo://notes\">"));
    assert!(text.contains("reply exactly NO_REPLY"));

    let listed = harness.run("(await watch_list())[0]['events']").await;
    assert_eq!(listed.content, "2");
}

#[tokio::test]
async fn on_event_filters_transform_drop_and_use_read_only_tools() {
    let mut harness = harness().await;
    let started = harness
        .run(
            "await watch_start(name='notes', uri='memo://notes', on_event='''\n\
             if event['type'] == 'ADDED':\n    \
                 owner = (await owner_of(key=event['key']))['owner']\n    \
                 result = f\"new note {event['key']} by {owner}\"\n\
             else:\n    \
                 result = None\n\
             result\n''')",
        )
        .await;
    assert!(!started.is_error, "{}", started.content);

    harness
        .resources
        .push("memo://notes", json!({"b": "2"}), "application/json")
        .await;
    let notice = harness.next_notice().await;
    assert_eq!(
        notice.lines,
        ["new note b by owner-of-b"],
        "DELETED a must be dropped"
    );
}

#[tokio::test]
async fn filters_cannot_call_mutating_tools_or_start_watches() {
    let harness = harness().await;
    // Mutating tools and watch built-ins are absent from the filter scope,
    // so the dry run rejects them before any subscription.
    for filter in [
        "await delete_note(key=event['key'])",
        "await watch_start(name='x', uri='memo://notes')",
    ] {
        let started = harness
            .run(&format!(
                "await watch_start(name='notes', uri='memo://notes', on_event={filter:?})"
            ))
            .await;
        assert!(started.is_error, "{filter}: {}", started.content);
        assert!(
            started.content.contains("undefined name"),
            "{}",
            started.content
        );
    }
    assert_eq!(*harness.resources.subscribe_count.lock().unwrap(), 0);
}

#[tokio::test]
async fn watch_stop_list_duplicates_and_limits() {
    let harness = harness().await;
    harness
        .run("await watch_start(name='w0', uri='memo://notes')")
        .await;
    let duplicate = harness
        .run("await watch_start(name='w0', uri='memo://notes')")
        .await;
    assert!(duplicate.is_error);
    assert!(
        duplicate.content.contains("already exists"),
        "{}",
        duplicate.content
    );

    for index in 1..5 {
        let started = harness
            .run(&format!(
                "await watch_start(name='w{index}', uri='memo://notes')"
            ))
            .await;
        assert!(!started.is_error, "{}", started.content);
    }
    let over = harness
        .run("await watch_start(name='w5', uri='memo://notes')")
        .await;
    assert!(
        over.content.contains("at most 5 watches"),
        "{}",
        over.content
    );

    assert_eq!(
        harness.run("await watch_stop(name='w0')").await.content,
        "true"
    );
    assert_eq!(
        harness.run("await watch_stop(name='w0')").await.content,
        "false"
    );
    assert_eq!(
        harness
            .run("sorted(w['name'] for w in await watch_list())")
            .await
            .content,
        r#"["w1","w2","w3","w4"]"#
    );
}

#[tokio::test]
async fn event_log_watch_delivers_by_seq() {
    let mut harness = harness().await;
    let uri = "kcp-events://watch/root:app/configmaps";
    harness.resources.set(
        uri,
        json!({"events": [{"seq": 1, "type": "ADDED"}]}),
        EVENT_LOG_MIME,
    );
    let started = harness
        .run(&format!("await watch_start(name='cm', uri='{uri}')"))
        .await;
    assert!(started.content.contains("event-log"), "{}", started.content);

    harness
        .resources
        .push(
            uri,
            json!({"events": [
                {"seq": 1, "type": "ADDED"},
                {"seq": 2, "type": "ADDED", "object": {"metadata": {"name": "x", "namespace": "default"}}},
            ]}),
            EVENT_LOG_MIME,
        )
        .await;
    let notice = harness.next_notice().await;
    assert_eq!(notice.lines, ["ADDED default/x"]);
}

#[tokio::test]
async fn ended_subscriptions_resubscribe_and_reconcile() {
    let mut harness = harness().await;
    harness
        .run("await watch_start(name='notes', uri='memo://notes')")
        .await;
    assert_eq!(*harness.resources.subscribe_count.lock().unwrap(), 1);

    harness.resources.set(
        "memo://notes",
        json!({"a": "1", "z": "26"}),
        "application/json",
    );
    harness.resources.end("memo://notes").await;

    let notice = harness.next_notice().await;
    assert_eq!(
        notice.lines,
        ["ADDED z"],
        "change during the gap is reconciled"
    );
    assert_eq!(*harness.resources.subscribe_count.lock().unwrap(), 2);

    harness
        .resources
        .push(
            "memo://notes",
            json!({"a": "2", "z": "26"}),
            "application/json",
        )
        .await;
    assert_eq!(harness.next_notice().await.lines, ["MODIFIED a (1→2)"]);
}

fn event_log(events: Value) -> Value {
    json!({ "events": events })
}

fn bid_event(seq: u64, kind: &str, phase: &str, version: &str) -> Value {
    json!({"seq": seq, "type": kind, "object": {"metadata": {"name": "bid-1", "uid": "u1", "resourceVersion": version}, "status": {"phase": phase}}})
}

#[tokio::test]
async fn repeated_lines_per_object_are_dropped_and_bursts_merge() {
    let mut harness = harness_with(
        WatchLimits {
            debounce: Duration::from_millis(300),
            max_wait: Duration::from_secs(5),
            ..fast_limits()
        },
        None,
        Arc::new(FakeResources::default()),
    )
    .await;
    harness
        .resources
        .set("kcp-events://bids", event_log(json!([])), EVENT_LOG_MIME);
    let started = harness
        .run("await watch_start(name='bids', uri='kcp-events://bids', on_event=\"'bid ' + event['object']['metadata']['name'] + ' ' + event['object']['status']['phase']\")")
        .await;
    assert!(!started.is_error, "{}", started.content);

    let mut events = Vec::new();
    // Three updates 100 ms apart: under the 300 ms quiet window, so one notice.
    for (seq, phase) in [(1, "Deciding"), (2, "Deciding"), (3, "Published")] {
        let kind = if seq == 1 { "ADDED" } else { "MODIFIED" };
        events.push(bid_event(seq, kind, phase, &seq.to_string()));
        harness
            .resources
            .push(
                "kcp-events://bids",
                event_log(json!(events)),
                EVENT_LOG_MIME,
            )
            .await;
        tokio::time::sleep(Duration::from_millis(100)).await;
    }
    let notice = harness.next_notice().await;
    assert_eq!(notice.lines, ["bid bid-1 Deciding", "bid bid-1 Published"]);
}

#[tokio::test]
async fn max_wait_caps_a_continuous_burst() {
    let mut harness = harness_with(
        WatchLimits {
            debounce: Duration::from_millis(200),
            max_wait: Duration::from_millis(300),
            ..fast_limits()
        },
        None,
        Arc::new(FakeResources::default()),
    )
    .await;
    harness
        .run("await watch_start(name='notes', uri='memo://notes')")
        .await;
    let started = tokio::time::Instant::now();
    for value in 2..12 {
        harness
            .resources
            .push(
                "memo://notes",
                json!({"a": value.to_string()}),
                "application/json",
            )
            .await;
        tokio::time::sleep(Duration::from_millis(60)).await;
        if let Ok(notice) = harness.notices.try_recv() {
            assert!(
                started.elapsed() < Duration::from_millis(550),
                "flushed late"
            );
            assert!(notice.lines.len() > 1, "{:?}", notice.lines);
            return;
        }
    }
    panic!("burst never flushed while updates kept coming");
}

#[tokio::test]
async fn watch_specs_persist_and_restore() {
    let dir = tempfile::tempdir().unwrap();
    let state = dir.path().join("watches.json");
    {
        let harness = harness_with(
            fast_limits(),
            Some(state.clone()),
            Arc::new(FakeResources::default()),
        )
        .await;
        harness
            .run("await watch_start(name='notes', uri='memo://notes', on_event=\"event['type']\")")
            .await;
        harness
            .run("await watch_start(name='gone', uri='memo://notes')")
            .await;
        harness.run("await watch_stop(name='gone')").await;
    }
    let saved: Value = serde_json::from_str(&std::fs::read_to_string(&state).unwrap()).unwrap();
    assert_eq!(
        saved,
        json!([{"name": "notes", "uri": "memo://notes", "on_event": "event['type']"}])
    );

    // A spec that can no longer start is reported and dropped from the file.
    let mut specs = saved.as_array().unwrap().clone();
    specs.push(json!({"name": "broken", "uri": "memo://notes", "on_event": "fmt(event)"}));
    std::fs::write(&state, serde_json::to_vec(&specs).unwrap()).unwrap();

    let mut harness = harness_with(
        fast_limits(),
        Some(state.clone()),
        Arc::new(FakeResources::default()),
    )
    .await;
    assert_eq!(harness.executor.restore_watches().await, 1);
    let restored = harness.next_notice().await;
    assert_eq!(restored.name, "restore");
    assert!(restored.lines[0].starts_with("restored notes memo://notes"));
    assert!(
        restored.lines[1].starts_with("dropped broken (invalid"),
        "{:?}",
        restored.lines
    );
    let listed = harness.run("[w['name'] for w in await watch_list()]").await;
    assert_eq!(listed.content, r#"["notes"]"#);
    let resaved: Value = serde_json::from_str(&std::fs::read_to_string(&state).unwrap()).unwrap();
    assert_eq!(resaved, saved, "failed spec dropped from the file");

    harness
        .resources
        .push("memo://notes", json!({"a": "2"}), "application/json")
        .await;
    assert_eq!(harness.next_notice().await.lines, ["MODIFIED"]);
}

#[tokio::test]
async fn on_event_syntax_errors_fail_watch_start() {
    let harness = harness().await;
    let started = harness
        .run("await watch_start(name='bad', uri='memo://notes', on_event='if event[')")
        .await;
    assert!(started.is_error, "{}", started.content);
    assert!(
        started.content.contains("on_event does not parse"),
        "{}",
        started.content
    );
    assert_eq!(harness.run("await watch_list()").await.content, "[]");
    assert_eq!(
        *harness.resources.subscribe_count.lock().unwrap(),
        0,
        "no subscription on failure"
    );
}

#[tokio::test]
async fn function_and_undefined_name_filters_fail_watch_start() {
    let harness = harness().await;
    // Live: the model passed a function object; it crosses as its repr.
    let function = harness
        .run("def fmt(event):\n    return event['type']\nawait watch_start(name='f', uri='memo://notes', on_event=fmt)")
        .await;
    assert!(function.is_error, "{}", function.content);
    assert!(
        function.content.contains("must be Python source text"),
        "{}",
        function.content
    );

    let undefined = harness
        .run("await watch_start(name='u', uri='memo://notes', on_event='fmt(event)')")
        .await;
    assert!(undefined.is_error, "{}", undefined.content);
    assert!(
        undefined.content.contains("undefined name"),
        "{}",
        undefined.content
    );

    // Missing keys on the sample event are fine.
    let keyed = harness
        .run("await watch_start(name='k', uri='memo://notes', on_event=\"event['object']['spec']['x']\")")
        .await;
    assert!(!keyed.is_error, "{}", keyed.content);
    assert_eq!(*harness.resources.subscribe_count.lock().unwrap(), 1);
}

#[tokio::test]
async fn transiently_failed_restores_stay_saved_until_started_or_stopped() {
    let dir = tempfile::tempdir().unwrap();
    let state = dir.path().join("watches.json");
    let specs = json!([
        {"name": "a", "uri": "memo://notes"},
        {"name": "b", "uri": "memo://notes"},
    ]);
    std::fs::write(&state, specs.to_string()).unwrap();

    let resources = Arc::new(FakeResources::default());
    *resources.fail_watch.lock().unwrap() = true;
    let mut harness = harness_with(fast_limits(), Some(state.clone()), resources).await;
    assert_eq!(harness.executor.restore_watches().await, 0);
    let notice = harness.next_notice().await;
    assert!(
        notice.lines[0].starts_with("could not restore a yet (kept"),
        "{:?}",
        notice.lines
    );
    let saved =
        || -> Value { serde_json::from_str(&std::fs::read_to_string(&state).unwrap()).unwrap() };
    assert_eq!(saved(), specs, "transient failures keep their specs");
    let listed = harness
        .run("[(w['name'], w.get('restore_failed')) for w in await watch_list()]")
        .await;
    assert_eq!(listed.content, r#"[["a",true],["b",true]]"#);

    // Server back: starting `a` again runs it; stopping `b` forgets it.
    *harness.resources.fail_watch.lock().unwrap() = false;
    let started = harness
        .run("await watch_start(name='a', uri='memo://notes')")
        .await;
    assert!(!started.is_error, "{}", started.content);
    assert_eq!(
        harness.run("await watch_stop(name='b')").await.content,
        "true"
    );
    assert_eq!(saved(), json!([{"name": "a", "uri": "memo://notes"}]));
}

/// The example in prompts/atirc-socialweb-computer.md must run as written.
#[tokio::test]
async fn socialweb_prompt_example_runs_and_reports_phase_changes() {
    let prompt = include_str!("../prompts/atirc-socialweb-computer.md");
    let code = prompt
        .split("```python\n")
        .nth(1)
        .and_then(|rest| rest.split("```").next())
        .expect("python example in the prompt");
    let mut harness = harness().await;
    for kind in ["bids", "computevms", "statesyncs"] {
        harness.resources.set(
            &format!("kcp-events://watch/root/{kind}"),
            event_log(json!([])),
            EVENT_LOG_MIME,
        );
    }
    let started = harness.run(code).await;
    assert!(!started.is_error, "{}", started.content);
    assert!(started.content.contains("\"bids\""), "{}", started.content);

    let bid = |version: &str, phase: &str, labels: Value| json!({"kind": "Bid", "metadata": {"name": "bid-1", "uid": "u1", "resourceVersion": version, "labels": labels}, "status": {"phase": phase}});
    harness
        .resources
        .push(
            "kcp-events://watch/root/bids",
            event_log(json!([
                {"seq": 1, "type": "MODIFIED", "object": bid("2", "Published", json!({"x": "1"})), "old": bid("1", "Published", json!({}))},
                {"seq": 2, "type": "MODIFIED", "object": bid("3", "Won", json!({"x": "1"})), "old": bid("2", "Published", json!({"x": "1"}))},
            ])),
            EVENT_LOG_MIME,
        )
        .await;
    let notice = harness.next_notice().await;
    assert_eq!(notice.lines, ["MODIFIED Bid bid-1: Published -> Won"]);

    harness
        .resources
        .push(
            "kcp-events://watch/root/bids",
            event_log(json!([
                {"seq": 3, "type": "DELETED", "object": bid("4", "Releasing", json!({}))},
            ])),
            EVENT_LOG_MIME,
        )
        .await;
    assert_eq!(
        harness.next_notice().await.lines,
        ["DELETED Bid bid-1: gone"]
    );
}