Signed-off-by: 未知时光 <732857315@qq.com>
This commit is contained in:
509
plugins/ai_test/src/main.rs
Normal file
509
plugins/ai_test/src/main.rs
Normal file
@@ -0,0 +1,509 @@
|
||||
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.iter().any(|name| *name == "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 resp = process
|
||||
.spawn(request_with_plugin(
|
||||
PLUGIN_ID,
|
||||
SpawnRequest {
|
||||
cmd: "cmd".to_string(),
|
||||
args: vec!["/C".to_string(), format!("echo {PROCESS_OK}")],
|
||||
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 path_string(path: &std::path::Path) -> String {
|
||||
path.to_string_lossy().to_string()
|
||||
}
|
||||
Reference in New Issue
Block a user