johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::fs;
use std::path::{Path, PathBuf};

use asyncai::mcp::{McpToolExecutor, StdioServerConfig};
use asyncai::resources::{EventExtractor, ResourceSource, WatchSignal, EVENT_LOG_MIME};
use asyncai::ToolExecutor;

struct Fixture {
    _directory: tempfile::TempDir,
    kubectl: PathBuf,
    admin_kubeconfig: PathBuf,
    log: PathBuf,
    stdin_log: PathBuf,
    events: PathBuf,
}

impl Fixture {
    fn new() -> Self {
        let directory = tempfile::tempdir().unwrap();
        let server = loopback_url();
        let workspaces = format!(
            r#"{{"items":[{{"metadata":{{"name":"app"}},"spec":{{"type":{{"name":"socialweb-application"}}}},"status":{{"phase":"Ready","url":"{server}/clusters/root:app"}}}}]}}"#
        );
        let view = directory.path().join("view.json");
        let workspace_list = directory.path().join("workspaces.json");
        let kubectl = directory.path().join("kubectl");
        let admin_kubeconfig = directory.path().join("admin.kubeconfig");
        let log = directory.path().join("kubectl.log");
        let stdin_log = directory.path().join("kubectl.stdin.log");
        let events = directory.path().join("watch-events.json");
        let seed = directory.path().join("seed.json");
        fs::write(&events, "").unwrap();
        fs::write(
            &seed,
            r#"{"kind":"List","items":[{"metadata":{"name":"old","namespace":"default","uid":"uid-old","resourceVersion":"1"}}]}"#,
        )
        .unwrap();

        fs::write(
            &view,
            format!(
                r#"{{"current-context":"root","clusters":[{{"name":"root","cluster":{{"server":"{server}/clusters/root","certificate-authority-data":"Y2E=","tls-server-name":"127.0.0.1"}}}}],"contexts":[{{"name":"root","context":{{"cluster":"root","user":"shard-admin"}}}}],"users":[{{"name":"shard-admin","user":{{"token":"admin-secret"}}}}]}}"#
            ),
        )
        .unwrap();
        fs::write(&workspace_list, workspaces).unwrap();
        fs::write(&admin_kubeconfig, "{}").unwrap();

        let script = format!(
            r#"#!/bin/sh
printf '%s\n' "$*" >> "{}"
case "$*" in
  *"config view"*)
    cat "{}"
    ;;
  *"--raw=/livez"*)
    echo ok
    ;;
  *"create token"*)
    echo fake-readonly-token
    ;;
  *"apply --server-side"*)
    cat >> "{}"
    echo '{{}}'
    ;;
  *"get workspaces"*)
    case "$*" in
      *"/clusters/root --request-timeout=30s get workspaces"*) cat "{}" ;;
      *) echo '{{"items":[]}}' ;;
    esac
    ;;
  *"get configmaps"*"--output-watch-events"*)
    exec tail -n +1 -f "{}"
    ;;
  *"get configmaps"*)
    cat "{}"
    ;;
  *"get computevms"*)
    echo '{{"kind":"List","items":[{{"metadata":{{"name":"probe"}}}}]}}'
    ;;
  *"get computevm "*)
    echo 'error: the server doesn'"'"'t have a resource type "computevm"' >&2
    exit 1
    ;;
  *"api-resources"*)
    printf '%-16s%-16s%-32s%-16s%-16s%s\n' NAME SHORTNAMES APIVERSION NAMESPACED KIND VERBS
    printf '%-16s%-16s%-32s%-16s%-16s%s\n' computevms '' socialweb.computer/v1alpha1 false ComputeVM get,list
    ;;
  *)
    echo '{{"kind":"List","items":[]}}'
    ;;
esac
"#,
            log.display(),
            view.display(),
            stdin_log.display(),
            workspace_list.display(),
            events.display(),
            seed.display(),
        );
        fs::write(&kubectl, script).unwrap();
        set_executable(&kubectl);

        Self {
            _directory: directory,
            kubectl,
            admin_kubeconfig,
            log,
            stdin_log,
            events,
        }
    }

    fn push_event(&self, line: &str) {
        use std::io::Write;
        let mut file = fs::OpenOptions::new()
            .append(true)
            .open(&self.events)
            .unwrap();
        writeln!(file, "{line}").unwrap();
    }

    fn config(&self, binary: &'static str) -> StdioServerConfig {
        StdioServerConfig {
            command: binary.to_string(),
            args: vec![
                "--kubectl".to_string(),
                self.kubectl.display().to_string(),
                "--kubeconfig".to_string(),
                self.admin_kubeconfig.display().to_string(),
            ],
        }
    }

    fn invocations(&self) -> String {
        fs::read_to_string(&self.log).unwrap_or_default()
    }

    fn token_kubeconfig(&self) -> PathBuf {
        let invocations = self.invocations();
        let line = invocations
            .lines()
            .find(|line| line.contains("get computevms"))
            .expect("list_objects never called kubectl");
        let mut arguments = line.split_whitespace();
        while let Some(argument) = arguments.next() {
            if argument == "--kubeconfig" {
                return PathBuf::from(arguments.next().expect("--kubeconfig has no value"));
            }
        }
        panic!("list_objects did not pass --kubeconfig: {line}");
    }
}

#[tokio::test]
async fn readonly_server_discovers_and_reads_with_a_minted_token() {
    let fixture = Fixture::new();
    let executor = McpToolExecutor::connect_stdio(vec![
        fixture.config(env!("CARGO_BIN_EXE_socialweb-computer-kcp-readonly"))
    ])
    .await
    .unwrap();

    let tools = executor.list_tools().await.unwrap();
    let names = tools
        .iter()
        .map(|tool| tool.name.as_str())
        .collect::<Vec<_>>();
    assert!(names.contains(&"list_kcps"));
    assert!(names.contains(&"readonly_credentials"));
    assert!(!names.contains(&"apply_object"));
    assert!(!names.contains(&"delete_object"));

    let discovered = executor
        .call_tool("list_kcps", serde_json::json!({}))
        .await
        .unwrap();
    assert!(!discovered.is_error, "{}", discovered.content);
    assert!(discovered
        .content
        .contains(&fixture.admin_kubeconfig.display().to_string()));

    let workspaces = executor
        .call_tool(
            "list_workspaces",
            serde_json::json!({"kcp": fixture.admin_kubeconfig.display().to_string(), "recursive": true}),
        )
        .await
        .unwrap();
    assert!(!workspaces.is_error, "{}", workspaces.content);
    assert!(workspaces.content.contains("root:app"));
    assert!(workspaces.content.contains("socialweb-application"));

    let listed = executor
        .call_tool(
            "list_objects",
            serde_json::json!({
                "kcp": fixture.admin_kubeconfig.display().to_string(),
                "workspace": "root:app",
                "resource": "computevms",
            }),
        )
        .await
        .unwrap();
    assert!(!listed.is_error, "{}", listed.content);
    assert!(listed.content.contains("probe"));

    let resources = executor
        .call_tool(
            "get_api_resources",
            serde_json::json!({
                "kcp": fixture.admin_kubeconfig.display().to_string(),
                "workspace": "root:app",
            }),
        )
        .await
        .unwrap();
    assert!(!resources.is_error, "{}", resources.content);
    let resources: serde_json::Value =
        serde_json::from_str(&resources.content).expect("api resources are JSON");
    let resources = &resources["resources"];
    assert_eq!(resources[0]["name"], "computevms");
    assert_eq!(resources[0]["kind"], "ComputeVM");
    assert_eq!(resources[0]["api_version"], "socialweb.computer/v1alpha1");

    let defaulted = executor
        .call_tool(
            "list_objects",
            serde_json::json!({"workspace": "root:app", "resource": "computevms"}),
        )
        .await
        .unwrap();
    assert!(!defaulted.is_error, "kcp omitted: {}", defaulted.content);
    let defaulted: serde_json::Value = serde_json::from_str(&defaulted.content).unwrap();
    assert_eq!(defaulted["items"][0]["metadata"]["name"], "probe");

    let definitions = executor.list_tools().await.unwrap();
    for definition in &definitions {
        assert_eq!(definition.read_only, Some(true), "{}", definition.name);
        assert!(!definition.is_mutating(), "{}", definition.name);
    }
    let list_objects = definitions
        .iter()
        .find(|definition| definition.name == "list_objects")
        .unwrap();
    assert!(list_objects.output_schema.as_ref().unwrap()["properties"]["items"].is_object());
    assert!(list_objects
        .description
        .as_deref()
        .unwrap()
        .contains("Example: names(await list_objects("));
    assert!(list_objects.input_schema["required"]
        .as_array()
        .is_some_and(|required| !required.iter().any(|name| name == "kcp")));

    let invocations = fixture.invocations();
    let admin = fixture.admin_kubeconfig.display().to_string();
    let mut read_calls = 0;
    for line in invocations.lines() {
        if ["get workspaces", "--raw=/livez", "get computevms"]
            .iter()
            .any(|needle| line.contains(needle))
        {
            read_calls += 1;
            assert!(
                !line.contains(&admin),
                "read call used the admin kubeconfig: {line}"
            );
        }
    }
    assert!(read_calls >= 3, "{invocations}");
    assert!(!invocations.contains("-o yaml"), "{invocations}");
    assert!(
        invocations
            .lines()
            .any(|line| line.contains("get computevms") && line.contains("-o json")),
        "{invocations}"
    );
    assert!(invocations.contains("apply -f -"), "{invocations}");
    assert!(invocations.contains("create token socialweb-computer-kcp-readonly"));
    assert!(
        !invocations.contains("fake-readonly-token"),
        "{invocations}"
    );

    let token_kubeconfig = fixture.token_kubeconfig();
    assert!(
        token_kubeconfig.starts_with(std::env::temp_dir()),
        "token kubeconfig not private temp file: {}",
        token_kubeconfig.display()
    );
    let contents = fs::read_to_string(&token_kubeconfig).unwrap();
    assert!(contents.contains("fake-readonly-token"));
    assert!(contents.contains("certificate-authority-data"));
    assert!(!contents.contains("client-certificate"));
    assert!(!contents.contains("client-key"));
    assert!(!contents.contains("admin-secret"));

    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let mode = fs::metadata(&token_kubeconfig)
            .unwrap()
            .permissions()
            .mode();
        assert_eq!(mode & 0o777, 0o600);
    }

    let credential = executor
        .call_tool(
            "readonly_credentials",
            serde_json::json!({"kcp": fixture.admin_kubeconfig.display().to_string(), "workspace": "root:app"}),
        )
        .await
        .unwrap();
    assert!(!credential.is_error, "{}", credential.content);
    assert!(credential.content.contains("token_expires_in_seconds"));
    assert!(!credential.content.contains("fake-readonly-token"));
}

#[tokio::test]
async fn readwrite_server_applies_with_admin_credentials() {
    let fixture = Fixture::new();
    let executor = McpToolExecutor::connect_stdio(vec![
        fixture.config(env!("CARGO_BIN_EXE_socialweb-computer-kcp-readwrite"))
    ])
    .await
    .unwrap();

    let tools = executor.list_tools().await.unwrap();
    let names = tools
        .iter()
        .map(|tool| tool.name.as_str())
        .collect::<Vec<_>>();
    assert!(names.contains(&"apply_object"));
    assert!(names.contains(&"delete_object"));
    assert!(names.contains(&"patch_object"));
    assert!(names.contains(&"scale_object"));
    let mutating = tools
        .iter()
        .filter(|tool| tool.is_mutating())
        .map(|tool| tool.name.as_str())
        .collect::<Vec<_>>();
    assert_eq!(
        mutating.len(),
        4,
        "apply/delete/patch/scale are mutating: {mutating:?}"
    );
    let delete = tools
        .iter()
        .find(|tool| tool.name == "delete_object")
        .unwrap();
    assert_eq!(delete.destructive, Some(true));
    let apply = tools
        .iter()
        .find(|tool| tool.name == "apply_object")
        .unwrap();
    assert_eq!(apply.destructive, Some(false));
    assert_eq!(apply.read_only, Some(false));

    let applied = executor
        .call_tool(
            "apply_object",
            serde_json::json!({
                "kcp": fixture.admin_kubeconfig.display().to_string(),
                "workspace": "root:app",
                "manifest": {
                    "apiVersion": "v1",
                    "kind": "ConfigMap",
                    "metadata": {"name": "probe", "namespace": "default"},
                    "data": {"key": "value"},
                },
                "dry_run": true,
            }),
        )
        .await
        .unwrap();
    assert!(!applied.is_error, "{}", applied.content);

    let invocations = fixture.invocations();
    let apply_line = invocations
        .lines()
        .find(|line| line.contains("apply --server-side"))
        .expect("apply_object did not run kubectl apply");
    assert!(
        apply_line.contains(&fixture.admin_kubeconfig.display().to_string()),
        "{apply_line}"
    );
    assert!(!apply_line.contains("readonly"), "{apply_line}");
    assert!(apply_line.contains("-o json"), "{apply_line}");
    assert!(
        fs::read_to_string(&fixture.stdin_log)
            .unwrap()
            .contains("\"ConfigMap\""),
        "manifest was not sent on stdin"
    );
}

#[tokio::test]
async fn apply_rejects_a_manifest_without_a_name() {
    let fixture = Fixture::new();
    let executor = McpToolExecutor::connect_stdio(vec![
        fixture.config(env!("CARGO_BIN_EXE_socialweb-computer-kcp-readwrite"))
    ])
    .await
    .unwrap();

    let applied = executor
        .call_tool(
            "apply_object",
            serde_json::json!({
                "kcp": fixture.admin_kubeconfig.display().to_string(),
                "workspace": "root:app",
                "manifest": {"apiVersion": "v1", "kind": "ConfigMap", "metadata": {}},
            }),
        )
        .await
        .unwrap();
    assert!(
        applied.is_error,
        "failed tool must set is_error: {}",
        applied.content
    );
    assert!(applied.content.contains("metadata.name"));
    assert!(!fixture.invocations().contains("apply --server-side"));
}

#[cfg(unix)]
fn loopback_url() -> String {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    let port = listener.local_addr().unwrap().port();
    format!("http://127.0.0.1:{port}")
}

fn set_executable(path: &Path) {
    use std::os::unix::fs::PermissionsExt;
    let mut permissions = fs::metadata(path).unwrap().permissions();
    permissions.set_mode(0o700);
    fs::set_permissions(path, permissions).unwrap();
}

#[cfg(not(unix))]
fn set_executable(_path: &Path) {}

#[tokio::test]
async fn kcp_watch_pushes_event_log_updates_without_polling() {
    watch_roundtrip(env!("CARGO_BIN_EXE_socialweb-computer-kcp-readonly")).await;
}

#[tokio::test]
async fn kcp_readwrite_server_watches_too() {
    watch_roundtrip(env!("CARGO_BIN_EXE_socialweb-computer-kcp-readwrite")).await;
}

async fn watch_roundtrip(binary: &'static str) {
    let fixture = Fixture::new();
    let executor = McpToolExecutor::connect_stdio(vec![fixture.config(binary)])
        .await
        .unwrap();

    let templates = executor.list_resources().await.unwrap();
    let events = templates
        .iter()
        .find(|resource| resource.name == "kcp-events")
        .expect("kcp-events template");
    assert_eq!(
        events.uri,
        "kcp-events://watch/{workspace}/{resource}{?namespace,labelSelector,fieldSelector,kcp}"
    );
    assert!(events.template && events.subscribable);
    assert!(events
        .example
        .as_deref()
        .is_some_and(|example| example.contains("watch_start")));

    let uri = "kcp-events://watch/root:app/configmaps?namespace=*";
    let mut watch = executor.watch(uri).await.unwrap();
    let baseline = executor.read_resource(uri).await.unwrap();
    assert_eq!(baseline.mime_type.as_deref(), Some(EVENT_LOG_MIME));
    let mut extractor = EventExtractor::baseline(&baseline, 100);

    // A replay of the seeded object is suppressed; a new one is delivered.
    fixture.push_event(
        r#"{"type":"ADDED","object":{"metadata":{"name":"old","namespace":"default","uid":"uid-old","resourceVersion":"1"}}}"#,
    );
    fixture.push_event(
        r#"{"type":"ADDED","object":{"metadata":{"name":"fresh","namespace":"default","uid":"uid-fresh","resourceVersion":"7","managedFields":[{"manager":"x"}]}}}"#,
    );
    let signal = tokio::time::timeout(std::time::Duration::from_secs(10), watch.signals.recv())
        .await
        .expect("no resources/updated notification")
        .unwrap();
    assert_eq!(signal, WatchSignal::Updated);

    let events = extractor.extract(&executor.read_resource(uri).await.unwrap());
    assert_eq!(events.len(), 1, "{events:?}");
    assert_eq!(events[0]["type"], "ADDED");
    assert_eq!(events[0]["seq"], 1);
    assert_eq!(events[0]["object"]["metadata"]["name"], "fresh");
    assert!(events[0]["object"]["metadata"]
        .get("managedFields")
        .is_none());

    // A MODIFIED event carries the previous object; asyncai adds `changed`.
    fixture.push_event(
        r#"{"type":"MODIFIED","object":{"metadata":{"name":"fresh","namespace":"default","uid":"uid-fresh","resourceVersion":"8"},"data":{"k":"v"}}}"#,
    );
    tokio::time::timeout(std::time::Duration::from_secs(10), watch.signals.recv())
        .await
        .expect("no resources/updated notification")
        .unwrap();
    let events = extractor.extract(&executor.read_resource(uri).await.unwrap());
    assert_eq!(events.len(), 1, "{events:?}");
    assert_eq!(events[0]["old"]["metadata"]["resourceVersion"], "7");
    assert_eq!(events[0]["changed"], serde_json::json!(["data"]));

    // kcp-api://: readable, JSON, not watchable.
    let api = executor
        .read_resource("kcp-api://root:app/resources")
        .await
        .unwrap();
    assert_eq!(api.mime_type.as_deref(), Some("application/json"));
    let api: serde_json::Value = serde_json::from_str(&api.text).unwrap();
    assert!(
        api["resources"]
            .as_array()
            .is_some_and(|list| !list.is_empty()),
        "{api}"
    );
    assert!(executor
        .watch("kcp-api://root:app/resources")
        .await
        .is_err());

    // A bad resource fails the first read (seed list), with a suggestion.
    let bad = "kcp-events://watch/root:app/computevm";
    let _bad_watch = executor.watch(bad).await.unwrap();
    let error = executor.read_resource(bad).await.unwrap_err().to_string();
    assert!(error.contains("doesn't have a resource type"), "{error}");
    assert!(error.contains("did you mean: computevms"), "{error}");

    let invocations = fixture.invocations();
    let watch_line = invocations
        .lines()
        .find(|line| line.contains("--output-watch-events"))
        .expect("kubectl watch never started");
    assert!(watch_line.contains("--request-timeout=0"), "{watch_line}");
    assert!(watch_line.contains("--all-namespaces"), "{watch_line}");
    assert!(
        !watch_line.contains(&fixture.admin_kubeconfig.display().to_string()),
        "read-only watch used the admin kubeconfig: {watch_line}"
    );
}