Initial commit: aether-tools-rs workspace

examples/stress_test/src/main.rs

@@ -0,0 +1,160 @@
+//! Stress test for aether-recall-mcp.
+//!
+//! Spawns one MCP server, keeps stdin/stdout open, sends many requests
+//! over the lifetime of the binary. Each request uses a UNIQUE failure
+//! text so the canonicalizer can't collapse records to duplicates.
+//!
+//! Usage: `cargo run --release -p stress_test -- <N>`
+//! where N is the number of records (default 50_000). Each invocation
+//! prints throughput stats on exit.
+
+use serde_json::{json, Value};
+use std::env;
+use std::io::{BufRead, BufReader, Read, Write};
+use std::path::PathBuf;
+use std::process::{Command, Stdio};
+use std::time::{Instant, SystemTime, UNIX_EPOCH};
+
+fn main() {
+    let n: usize = env::args().nth(1)
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(50_000);
+
+    // Use an isolated tmp dir so the stress run never collides with the
+    // user's real ledger.
+    let stamp = SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos();
+    let data_dir = std::env::temp_dir().join(format!("stress-{stamp}"));
+    std::fs::create_dir_all(&data_dir).unwrap();
+
+    // Resolve the binary at runtime. env!("CARGO_BIN_EXE_<bin>") only
+    // works when this crate *owns* the binary; here we depend on it
+    // from a sibling crate, so we probe the workspace target/ dir.
+    let manifest = env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR");
+    let examples_dir = PathBuf::from(manifest);
+    let workspace_dir = examples_dir.parent().and_then(|p| p.parent()).unwrap().to_path_buf();
+    let bin = ["release", "debug"].iter()
+        .map(|p| workspace_dir.join("target").join(p).join("aether-recall-mcp"))
+        .find(|p| p.exists())
+        .expect("aether-recall-mcp binary missing — run `cargo build --release --workspace` first");
+    println!("spawning {:?} for {n} ops", bin);
+    let mut child = Command::new(bin)
+        .env("AETHER_RECALL_DIR", &data_dir)
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::null())
+        .spawn()
+        .expect("spawn mcp");
+    // Take stdin/stdout once each. stdin becomes our writer; stdout
+    // becomes a buffered reader for incoming responses.
+    let stdout = child.stdout.take().unwrap();
+    let mut writer = child.stdin.take().unwrap();
+    let mut out = BufReader::new(stdout);
+
+    // PHASE 1: seed N unique failure records
+    let t = Instant::now();
+    for i in 0..n {
+        let req = json!({
+            "jsonrpc":"2.0","id": i,
+            "method":"tools/call",
+            "params":{"name":"record_outcome","arguments":{
+                "failure": format!("UNIQ_X{i:08}_error_{i}: cannot resolve type 'Z{i}' to inferred type 'W{i}'"),
+                "resolution": format!("fix {}: type annotation `W`", i),
+                "worked": true
+            }}
+        });
+        write_line(&mut writer, &req.to_string());
+        drain_response(&mut out);
+    }
+    let seed_elapsed = t.elapsed();
+    println!(
+        "seed: {n} records in {:.2?} ({:.1} records/s)",
+        seed_elapsed,
+        n as f64 / seed_elapsed.as_secs_f64()
+    );
+
+    // PHASE 2: look each one up via MCP
+    let t = Instant::now();
+    let mut hits = 0usize;
+    for i in 0..n {
+        let req = json!({
+            "jsonrpc":"2.0","id": i + n,
+            "method":"tools/call",
+            "params":{"name":"lookup_resolution","arguments":{
+                "failure": format!("UNIQ_X{i:08}_error_{i}: cannot resolve type 'Z{i}' to inferred type 'W{i}'")
+            }}
+        });
+        write_line(&mut writer, &req.to_string());
+        let resp = drain_response(&mut out);
+        if resp_hit(&resp) {
+            hits += 1;
+        }
+    }
+    let lookup_elapsed = t.elapsed();
+    println!(
+        "lookups: {n} in {:.2?} ({:.1}/s) — hits={hits} ({:.1}%)",
+        lookup_elapsed,
+        n as f64 / lookup_elapsed.as_secs_f64(),
+        hits as f64 * 100.0 / n as f64
+    );
+
+    // PHASE 3: stats
+    write_line(&mut writer, &json!({
+        "jsonrpc":"2.0","id": n*2+1,
+        "method":"tools/call",
+        "params":{"name":"memory_stats","arguments":{}}
+    }).to_string());
+    let stats_resp = drain_response(&mut out);
+    println!("stats: {}", stats_resp.get("result").and_then(|r| r.get("content")).and_then(|c| c.as_array()).and_then(|a| a.first()).and_then(|t| t.get("text")).and_then(|s| s.as_str()).unwrap_or("?"));
+
+    let _ = child.kill();
+    let _ = child.wait();
+}
+
+fn write_line(w: &mut std::process::ChildStdin, line: &str) {
+    w.write_all(line.as_bytes()).unwrap();
+    w.write_all(b"\n").unwrap();
+    w.flush().unwrap();
+}
+
+fn drain_response(r: &mut BufReader<std::process::ChildStdout>) -> Value {
+    let mut line = String::new();
+    r.read_line(&mut line).unwrap();
+    let mut s = &line[..];
+    while !s.trim().is_empty() && !s.starts_with('{') && !s.starts_with('[') {
+        line.clear();
+        r.read_line(&mut line).unwrap();
+        s = &line[..];
+    }
+    serde_json::from_str(s.trim()).unwrap_or_else(|_| {
+        let mut all = String::new();
+        r.read_to_string(&mut all).ok();
+        panic!("non-JSON response from MCP: {:?}\nnext bytes: {:?}", s.trim(), &all[..all.len().min(2048)]);
+    })
+}
+
+fn resp_hit(v: &Value) -> bool {
+    v.get("result")
+        .and_then(|r| r.get("content"))
+        .and_then(|c| c.as_array())
+        .and_then(|a| a.first())
+        .and_then(|t| t.get("text"))
+        .and_then(|s| s.as_str())
+        .map(|s| s.contains("\"hit\":true"))
+        .unwrap_or(false)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn resp_hit_recognises_true_and_false() {
+        let yes: Value = serde_json::from_str(
+            r#"{"result":{"content":[{"text":"{\"hit\":true, \"signature\":\"abc\"}"}]}}"#
+        ).unwrap();
+        let no:  Value = serde_json::from_str(
+            r#"{"result":{"content":[{"text":"{\"hit\":false}"}]}}"#
+        ).unwrap();
+        assert!(resp_hit(&yes));
+        assert!(!resp_hit(&no));
+    }
+}