Agentsd/plugins/ai_test/src/main.rs

use agentsd_plugin_sdk::{export, plugin_info, request_with_plugin};
use agentsd_proto::agentsd::{
    display_client::DisplayClient, fs_client::FsClient, memory_client::MemoryClient,
    network_client::NetworkClient, plugin_bridge_client::PluginBridgeClient,
    process_client::ProcessClient, timer_client::TimerClient, CallRequest, DisplayTextRequest,
    Empty, FsPathRequest, FsReadRequest, FsRenameRequest, FsWriteRequest, MemoryKeyRequest,
    MemoryListRequest, MemoryReadRequest, MemoryWriteRequest, PidRequest, SpawnRequest,
    TimerOnceRequest,
};
use anyhow::{anyhow, bail, ensure, Context, Result};
use serde::Serialize;
use serde_json::Value;
use std::collections::HashMap;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::time::{timeout, Duration};
use tokio_stream::StreamExt;
use tonic::transport::Channel;
use tonic::Code;

const PLUGIN_ID: &str = "ai-test";
const PROCESS_OK: &str = "AI_TEST_PROCESS_OK";
const TIMER_CALLBACK_ID: &str = "ai-test-callback";

#[derive(Debug, Serialize)]
struct CheckResult {
    name: String,
    ok: bool,
    details: String,
}

#[derive(Debug, Serialize)]
struct Summary<'a> {
    plugin: &'a str,
    ok: bool,
    results: &'a [CheckResult],
}

struct CheckRunner {
    channel: Channel,
    results: Vec<CheckResult>,
    run_id: String,
    scratch_dir: PathBuf,
}

impl CheckRunner {
    fn new(channel: Channel) -> Self {
        let run_id = now_millis().to_string();
        let scratch_dir = executable_data_dir().join("ai_test").join(&run_id);

        Self {
            channel,
            results: Vec::new(),
            run_id,
            scratch_dir,
        }
    }

    async fn run(&mut self) -> i32 {
        println!(
            "AI Test Plugin connecting to {}",
            agentsd_plugin_sdk::agentsd_addr()
        );
        println!("Scratch dir: {}", self.scratch_dir.display());

        let result = self.register().await;
        self.record_result("register", result);
        let result = self.display_text().await;
        self.record_result("display.text", result);
        let result = self.memory_roundtrip().await;
        self.record_result("memory roundtrip", result);
        let result = self.fs_roundtrip().await;
        self.record_result("fs roundtrip", result);
        let result = self.process_roundtrip().await;
        self.record_result("process roundtrip", result);
        let result = self.timer_event().await;
        self.record_result("timer event", result);
        let result = self.permission_denial().await;
        self.record_result("permission denial", result);

        let ok = self.results.iter().all(|r| r.ok);
        let summary = Summary {
            plugin: PLUGIN_ID,
            ok,
            results: &self.results,
        };
        let summary_json = serde_json::to_string(&summary)
            .unwrap_or_else(|e| format!(r#"{{"plugin":"{PLUGIN_ID}","ok":false,"error":"{e}"}}"#));
        println!("AI_TEST_SUMMARY {summary_json}");

        if ok {
            0
        } else {
            1
        }
    }

    fn record_result(&mut self, name: &str, result: Result<String>) {
        match result {
            Ok(details) => {
                self.record(name, true, if details.is_empty() { "ok" } else { &details })
            }
            Err(err) => self.record(name, false, &err.to_string()),
        }
    }

    fn record(&mut self, name: &str, ok: bool, details: &str) {
        let status = if ok { "PASS" } else { "FAIL" };
        println!("[{status}] {name}: {details}");
        self.results.push(CheckResult {
            name: name.to_string(),
            ok,
            details: details.to_string(),
        });
    }

    async fn register(&mut self) -> Result<String> {
        let mut bridge = PluginBridgeClient::new(self.channel.clone());
        let resp = bridge
            .register(plugin_info(
                PLUGIN_ID,
                "AI Test Plugin",
                "0.1.0",
                "basic",
                &[
                    "display.text",
                    "memory.write",
                    "memory.read",
                    "memory.list",
                    "memory.delete",
                    "fs.write",
                    "fs.read",
                    "fs.stat",
                    "fs.list",
                    "fs.rename",
                    "fs.delete",
                    "process.spawn",
                    "process.stdout",
                    "process.wait",
                    "timer.once",
                ],
                &[],
                vec![export("run_checks", "Run Agentsd verification checks")],
            ))
            .await
            .context("register RPC")?
            .into_inner();

        ensure!(resp.ok, "{}", resp.error);
        Ok(format!("session={}", resp.session_id))
    }

    async fn display_text(&mut self) -> Result<String> {
        let mut display = DisplayClient::new(self.channel.clone());
        let resp = display
            .text(request_with_plugin(
                PLUGIN_ID,
                DisplayTextRequest {
                    content: format!("[ai-test] display check run={}", self.run_id),
                },
            )?)
            .await
            .context("display.text RPC")?
            .into_inner();

        ensure!(resp.ok, "{}", resp.error);
        Ok("display.text ok".to_string())
    }

    async fn memory_roundtrip(&mut self) -> Result<String> {
        let mut memory = MemoryClient::new(self.channel.clone());
        let key = format!("ai_test.{}.greeting", self.run_id);
        let value = format!("hello-{}", self.run_id).into_bytes();

        let write = memory
            .write(request_with_plugin(
                PLUGIN_ID,
                MemoryWriteRequest {
                    key: key.clone(),
                    value: value.clone(),
                },
            )?)
            .await
            .context("memory.write RPC")?
            .into_inner();
        ensure!(write.ok, "write: {}", write.error);

        let read = memory
            .read(request_with_plugin(
                PLUGIN_ID,
                MemoryReadRequest { key: key.clone() },
            )?)
            .await
            .context("memory.read RPC")?
            .into_inner();
        ensure!(
            read.found && read.value == value,
            "read mismatch found={} value={:?} error={:?}",
            read.found,
            read.value,
            read.error
        );

        let list = memory
            .list(request_with_plugin(
                PLUGIN_ID,
                MemoryListRequest {
                    prefix: format!("ai_test.{}", self.run_id),
                    limit: 10,
                },
            )?)
            .await
            .context("memory.list RPC")?
            .into_inner();
        ensure!(
            list.error.is_empty() && list.keys.iter().any(|k| k == &key),
            "list missing key; keys={:?} error={:?}",
            list.keys,
            list.error
        );

        let delete = memory
            .delete(request_with_plugin(
                PLUGIN_ID,
                MemoryKeyRequest { key: key.clone() },
            )?)
            .await
            .context("memory.delete RPC")?
            .into_inner();
        ensure!(delete.ok, "delete: {}", delete.error);

        Ok(format!("key={key}"))
    }

    async fn fs_roundtrip(&mut self) -> Result<String> {
        let mut fs = FsClient::new(self.channel.clone());
        let base = self.scratch_dir.clone();
        let file_a = base.join("sample.txt");
        let file_b = base.join("renamed.txt");
        let data = format!("AI_TEST_FS_OK {}", self.run_id).into_bytes();

        let write = fs
            .write(request_with_plugin(
                PLUGIN_ID,
                FsWriteRequest {
                    path: path_string(&file_a),
                    data: data.clone(),
                    append: false,
                },
            )?)
            .await
            .context("fs.write RPC")?
            .into_inner();
        ensure!(write.ok, "write: {}", write.error);

        let read = fs
            .read(request_with_plugin(
                PLUGIN_ID,
                FsReadRequest {
                    path: path_string(&file_a),
                    offset: 0,
                    length: 0,
                },
            )?)
            .await
            .context("fs.read RPC")?
            .into_inner();
        ensure!(
            read.ok && read.data == data,
            "read mismatch ok={} data={:?} error={:?}",
            read.ok,
            read.data,
            read.error
        );

        let stat = fs
            .stat(request_with_plugin(
                PLUGIN_ID,
                FsPathRequest {
                    path: path_string(&file_a),
                },
            )?)
            .await
            .context("fs.stat RPC")?
            .into_inner();
        ensure!(
            stat.ok && stat.is_file && stat.size == data.len() as i64,
            "stat mismatch ok={} is_file={} size={} error={:?}",
            stat.ok,
            stat.is_file,
            stat.size,
            stat.error
        );

        let list = fs
            .list(request_with_plugin(
                PLUGIN_ID,
                FsPathRequest {
                    path: path_string(&base),
                },
            )?)
            .await
            .context("fs.list RPC")?
            .into_inner();
        let names: Vec<&str> = list.entries.iter().map(|e| e.name.as_str()).collect();
        ensure!(
            list.ok && names.contains(&"sample.txt"),
            "list mismatch ok={} names={:?} error={:?}",
            list.ok,
            names,
            list.error
        );

        let rename = fs
            .rename(request_with_plugin(
                PLUGIN_ID,
                FsRenameRequest {
                    from: path_string(&file_a),
                    to: path_string(&file_b),
                },
            )?)
            .await
            .context("fs.rename RPC")?
            .into_inner();
        ensure!(rename.ok, "rename: {}", rename.error);

        let delete = fs
            .delete(request_with_plugin(
                PLUGIN_ID,
                FsPathRequest {
                    path: path_string(&base),
                },
            )?)
            .await
            .context("fs.delete RPC")?
            .into_inner();
        ensure!(delete.ok, "delete: {}", delete.error);

        Ok(path_string(&base))
    }

    async fn process_roundtrip(&mut self) -> Result<String> {
        let mut process = ProcessClient::new(self.channel.clone());
        let (cmd, args) = process_command();
        let resp = process
            .spawn(request_with_plugin(
                PLUGIN_ID,
                SpawnRequest {
                    cmd,
                    args,
                    env: HashMap::new(),
                    cwd: executable_dir()
                        .map(|path| path_string(&path))
                        .unwrap_or_default(),
                },
            )?)
            .await
            .context("process.spawn RPC")?
            .into_inner();
        ensure!(resp.ok, "{}", resp.error);
        let pid = resp.pid;

        let mut stdout = process
            .stdout(request_with_plugin(PLUGIN_ID, PidRequest { pid })?)
            .await
            .context("process.stdout RPC")?
            .into_inner();
        let mut chunks = Vec::new();
        while let Some(chunk) = stdout.next().await {
            chunks.extend(chunk.context("process.stdout stream")?.data);
        }
        let output = String::from_utf8_lossy(&chunks).trim().to_string();

        let wait = process
            .wait(request_with_plugin(PLUGIN_ID, PidRequest { pid })?)
            .await
            .context("process.wait RPC")?
            .into_inner();
        ensure!(
            wait.ok && wait.exit_code == 0,
            "wait failed ok={} code={} error={:?}",
            wait.ok,
            wait.exit_code,
            wait.error
        );
        ensure!(output == PROCESS_OK, "stdout mismatch: {output:?}");

        Ok(format!("pid={pid} stdout={output}"))
    }

    async fn timer_event(&mut self) -> Result<String> {
        let mut bridge = PluginBridgeClient::new(self.channel.clone());
        let mut timer = TimerClient::new(self.channel.clone());

        let mut stream = bridge
            .stream_call(request_with_plugin(
                PLUGIN_ID,
                CallRequest {
                    target: PLUGIN_ID.to_string(),
                    r#fn: "events".to_string(),
                    args: Vec::new(),
                },
            )?)
            .await
            .context("plugin_bridge.stream_call RPC")?
            .into_inner();

        let once = timer
            .once(request_with_plugin(
                PLUGIN_ID,
                TimerOnceRequest {
                    delay_ms: 100,
                    callback_id: TIMER_CALLBACK_ID.to_string(),
                },
            )?)
            .await
            .context("timer.once RPC")?
            .into_inner();
        ensure!(once.ok, "{}", once.error);

        let event = timeout(Duration::from_secs(5), stream.message())
            .await
            .context("timed out waiting for timer event")?
            .context("timer event stream")?
            .ok_or_else(|| anyhow!("timer event stream ended before event"))?;
        ensure!(event.ok, "stream event failed: {}", event.error);

        let payload: Value =
            serde_json::from_slice(&event.result).context("decode timer event JSON")?;
        ensure!(
            payload.get("event").and_then(Value::as_str) == Some("timer")
                && payload.get("callback_id").and_then(Value::as_str) == Some(TIMER_CALLBACK_ID),
            "unexpected payload: {payload}"
        );

        Ok(serde_json::to_string(&payload)?)
    }

    async fn permission_denial(&mut self) -> Result<String> {
        let mut network = NetworkClient::new(self.channel.clone());
        match network
            .available(request_with_plugin(PLUGIN_ID, Empty {})?)
            .await
        {
            Ok(_) => bail!("network.available unexpectedly succeeded"),
            Err(status) if status.code() == Code::PermissionDenied => {
                Ok(status.message().to_string())
            }
            Err(status) => bail!("wrong error: {} {}", status.code(), status.message()),
        }
    }
}

#[tokio::main]
async fn main() {
    let code = match agentsd_plugin_sdk::connect().await {
        Ok(channel) => {
            let mut runner = CheckRunner::new(channel);
            runner.run().await
        }
        Err(err) => {
            let result = CheckResult {
                name: "connect".to_string(),
                ok: false,
                details: err.to_string(),
            };
            println!(
                "AI Test Plugin connecting to {}",
                agentsd_plugin_sdk::agentsd_addr()
            );
            println!("[FAIL] connect: {}", result.details);
            let results = vec![result];
            let summary = Summary {
                plugin: PLUGIN_ID,
                ok: false,
                results: &results,
            };
            let summary_json = serde_json::to_string(&summary).unwrap_or_else(|e| {
                format!(r#"{{"plugin":"{PLUGIN_ID}","ok":false,"error":"{e}"}}"#)
            });
            println!("AI_TEST_SUMMARY {summary_json}");
            1
        }
    };

    std::process::exit(code);
}

fn now_millis() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis()
}

fn executable_dir() -> Option<PathBuf> {
    std::env::current_exe()
        .ok()
        .and_then(|exe| exe.parent().map(|parent| parent.to_path_buf()))
}

fn executable_data_dir() -> PathBuf {
    executable_dir()
        .unwrap_or_else(|| PathBuf::from("."))
        .join("data")
}

fn process_command() -> (String, Vec<String>) {
    if cfg!(windows) {
        (
            "cmd".to_string(),
            vec!["/C".to_string(), format!("echo {PROCESS_OK}")],
        )
    } else {
        (
            "sh".to_string(),
            vec!["-c".to_string(), format!("printf %s {PROCESS_OK}")],
        )
    }
}

fn path_string(path: &std::path::Path) -> String {
    path.to_string_lossy().to_string()
}