mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Illia PolosukhinandClaude Opus 4.6
|
2e62d71567 |
Generated
+169
-24
@@ -170,6 +170,18 @@ version = "1.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1"
|
||||
|
||||
[[package]]
|
||||
name = "argon2"
|
||||
version = "0.5.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3c3610892ee6e0cbce8ae2700349fcf8f98adb0dbfbee85aec3c9179d29cc072"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"blake2",
|
||||
"cpufeatures",
|
||||
"password-hash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "arrayref"
|
||||
version = "0.3.9"
|
||||
@@ -403,6 +415,12 @@ version = "0.22.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
|
||||
|
||||
[[package]]
|
||||
name = "base64ct"
|
||||
version = "1.8.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06"
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "1.3.2"
|
||||
@@ -427,6 +445,15 @@ dependencies = [
|
||||
"wyz",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "blake2"
|
||||
version = "0.10.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe"
|
||||
dependencies = [
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "blake3"
|
||||
version = "1.8.3"
|
||||
@@ -545,6 +572,15 @@ dependencies = [
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bs58"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4"
|
||||
dependencies = [
|
||||
"tinyvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bumpalo"
|
||||
version = "3.19.1"
|
||||
@@ -796,6 +832,12 @@ dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "const-oid"
|
||||
version = "0.9.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8"
|
||||
|
||||
[[package]]
|
||||
name = "constant_time_eq"
|
||||
version = "0.4.2"
|
||||
@@ -1115,6 +1157,33 @@ dependencies = [
|
||||
"cipher",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "curve25519-dalek"
|
||||
version = "4.1.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures",
|
||||
"curve25519-dalek-derive",
|
||||
"digest",
|
||||
"fiat-crypto",
|
||||
"rustc_version",
|
||||
"subtle",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "curve25519-dalek-derive"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "darling"
|
||||
version = "0.21.3"
|
||||
@@ -1200,6 +1269,16 @@ dependencies = [
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "der"
|
||||
version = "0.7.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb"
|
||||
dependencies = [
|
||||
"const-oid",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "deranged"
|
||||
version = "0.5.5"
|
||||
@@ -1354,6 +1433,31 @@ version = "1.0.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555"
|
||||
|
||||
[[package]]
|
||||
name = "ed25519"
|
||||
version = "2.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53"
|
||||
dependencies = [
|
||||
"pkcs8",
|
||||
"signature",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ed25519-dalek"
|
||||
version = "2.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9"
|
||||
dependencies = [
|
||||
"curve25519-dalek",
|
||||
"ed25519",
|
||||
"rand_core 0.6.4",
|
||||
"serde",
|
||||
"sha2",
|
||||
"subtle",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "either"
|
||||
version = "1.15.0"
|
||||
@@ -1497,6 +1601,12 @@ dependencies = [
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fiat-crypto"
|
||||
version = "0.2.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d"
|
||||
|
||||
[[package]]
|
||||
name = "filetime"
|
||||
version = "0.2.27"
|
||||
@@ -2181,11 +2291,14 @@ dependencies = [
|
||||
"aes-gcm",
|
||||
"aho-corasick",
|
||||
"anyhow",
|
||||
"argon2",
|
||||
"async-trait",
|
||||
"axum",
|
||||
"base64 0.22.1",
|
||||
"blake3",
|
||||
"bollard",
|
||||
"borsh",
|
||||
"bs58",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"clap",
|
||||
@@ -2193,6 +2306,7 @@ dependencies = [
|
||||
"deadpool-postgres",
|
||||
"dirs 6.0.0",
|
||||
"dotenvy",
|
||||
"ed25519-dalek",
|
||||
"futures",
|
||||
"hkdf",
|
||||
"http-body-util",
|
||||
@@ -2234,30 +2348,7 @@ dependencies = [
|
||||
"wasmtime",
|
||||
"wasmtime-wasi",
|
||||
"zbus",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ironclaw-bench"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"async-trait",
|
||||
"chrono",
|
||||
"clap",
|
||||
"futures",
|
||||
"ironclaw",
|
||||
"regex",
|
||||
"rust_decimal",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"toml",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"uuid",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -2798,6 +2889,17 @@ dependencies = [
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "password-hash"
|
||||
version = "0.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "346f04948ba92c43e8469c1ee6736c7563d71012b17d40745260fe106aac2166"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"rand_core 0.6.4",
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "paste"
|
||||
version = "1.0.15"
|
||||
@@ -2868,6 +2970,16 @@ dependencies = [
|
||||
"futures-io",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pkcs8"
|
||||
version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
|
||||
dependencies = [
|
||||
"der",
|
||||
"spki",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pkg-config"
|
||||
version = "0.3.32"
|
||||
@@ -3949,6 +4061,15 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "signature"
|
||||
version = "2.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de"
|
||||
dependencies = [
|
||||
"rand_core 0.6.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "simdutf8"
|
||||
version = "0.1.5"
|
||||
@@ -3986,6 +4107,16 @@ dependencies = [
|
||||
"windows-sys 0.60.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "spki"
|
||||
version = "0.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d"
|
||||
dependencies = [
|
||||
"base64ct",
|
||||
"der",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sptr"
|
||||
version = "0.3.2"
|
||||
@@ -5970,6 +6101,20 @@ name = "zeroize"
|
||||
version = "1.8.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0"
|
||||
dependencies = [
|
||||
"zeroize_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zeroize_derive"
|
||||
version = "1.4.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "85a5b4158499876c763cb03bc4e49185d3cccbabb15b33c627f7884f43db852e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.114",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "zerotrie"
|
||||
|
||||
+7
-3
@@ -1,6 +1,3 @@
|
||||
[workspace]
|
||||
members = [".", "benchmarks"]
|
||||
|
||||
[package]
|
||||
name = "ironclaw"
|
||||
version = "0.1.0"
|
||||
@@ -93,6 +90,13 @@ sha2 = "0.10"
|
||||
blake3 = "1"
|
||||
rand = "0.8"
|
||||
|
||||
# NEAR key management (ed25519 signing, borsh serialization, base58 encoding)
|
||||
ed25519-dalek = { version = "2", features = ["rand_core", "zeroize"] }
|
||||
borsh = { version = "1", features = ["derive"] }
|
||||
bs58 = "0.5"
|
||||
argon2 = "0.5"
|
||||
zeroize = { version = "1", features = ["derive"] }
|
||||
|
||||
# Docker sandbox
|
||||
bollard = "0.18"
|
||||
|
||||
|
||||
@@ -1,49 +0,0 @@
|
||||
[package]
|
||||
name = "ironclaw-bench"
|
||||
version = "0.1.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.85"
|
||||
description = "Benchmarking harness for IronClaw agent"
|
||||
license = "MIT OR Apache-2.0"
|
||||
|
||||
[[bin]]
|
||||
name = "ironclaw-bench"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
ironclaw = { path = ".." }
|
||||
|
||||
# Async runtime
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
tokio-stream = { version = "0.1", features = ["sync"] }
|
||||
futures = "0.3"
|
||||
|
||||
# Serialization
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
toml = "0.8"
|
||||
|
||||
# CLI
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
|
||||
# Core types
|
||||
uuid = { version = "1", features = ["v4", "serde"] }
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
rust_decimal = { version = "1", features = ["serde", "serde-with-str"] }
|
||||
|
||||
# Error handling
|
||||
thiserror = "2"
|
||||
anyhow = "1"
|
||||
|
||||
# Async traits
|
||||
async-trait = "0.1"
|
||||
|
||||
# Logging
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
|
||||
# Scoring
|
||||
regex = "1"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
@@ -1,237 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::scoring;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskSubmission};
|
||||
|
||||
/// A single entry in the custom JSONL format.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct CustomEntry {
|
||||
id: String,
|
||||
prompt: String,
|
||||
#[serde(default)]
|
||||
context: Option<String>,
|
||||
#[serde(default)]
|
||||
tags: Vec<String>,
|
||||
#[serde(default)]
|
||||
expected: Option<String>,
|
||||
#[serde(default)]
|
||||
expected_contains: Option<String>,
|
||||
#[serde(default)]
|
||||
expected_regex: Option<String>,
|
||||
/// "exact", "contains", "regex", or "llm" (default: "exact")
|
||||
#[serde(default = "default_scorer")]
|
||||
scorer: String,
|
||||
}
|
||||
|
||||
fn default_scorer() -> String {
|
||||
"exact".to_string()
|
||||
}
|
||||
|
||||
/// Custom JSONL benchmark suite.
|
||||
///
|
||||
/// Each line of the JSONL file is a task with `id`, `prompt`, and scoring
|
||||
/// criteria (`expected`, `expected_contains`, `expected_regex`).
|
||||
pub struct CustomSuite {
|
||||
dataset_path: PathBuf,
|
||||
}
|
||||
|
||||
impl CustomSuite {
|
||||
pub fn new(dataset_path: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for CustomSuite {
|
||||
fn name(&self) -> &str {
|
||||
"Custom JSONL"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"custom"
|
||||
}
|
||||
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError> {
|
||||
let file = std::fs::File::open(&self.dataset_path).map_err(BenchError::Io)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: CustomEntry = serde_json::from_str(trimmed)
|
||||
.map_err(|e| BenchError::Config(format!("line {}: {}", line_num + 1, e)))?;
|
||||
|
||||
let mut metadata = serde_json::json!({
|
||||
"scorer": entry.scorer,
|
||||
});
|
||||
if let Some(ref expected) = entry.expected {
|
||||
metadata["expected"] = serde_json::Value::String(expected.clone());
|
||||
}
|
||||
if let Some(ref expected_contains) = entry.expected_contains {
|
||||
metadata["expected_contains"] =
|
||||
serde_json::Value::String(expected_contains.clone());
|
||||
}
|
||||
if let Some(ref expected_regex) = entry.expected_regex {
|
||||
metadata["expected_regex"] = serde_json::Value::String(expected_regex.clone());
|
||||
}
|
||||
|
||||
tasks.push(BenchTask {
|
||||
id: entry.id,
|
||||
prompt: entry.prompt,
|
||||
context: entry.context,
|
||||
resources: vec![],
|
||||
tags: entry.tags,
|
||||
expected_turns: None,
|
||||
timeout: None,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError> {
|
||||
let scorer = task
|
||||
.metadata
|
||||
.get("scorer")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("exact");
|
||||
|
||||
match scorer {
|
||||
"exact" => {
|
||||
if let Some(expected) = task.metadata.get("expected").and_then(|v| v.as_str()) {
|
||||
Ok(scoring::exact_match(expected, &submission.response))
|
||||
} else {
|
||||
Err(BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "no 'expected' field for exact scoring".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
"contains" => {
|
||||
if let Some(expected) = task
|
||||
.metadata
|
||||
.get("expected_contains")
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
Ok(scoring::contains_match(expected, &submission.response))
|
||||
} else {
|
||||
Err(BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "no 'expected_contains' field for contains scoring".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
"regex" => {
|
||||
if let Some(pattern) = task.metadata.get("expected_regex").and_then(|v| v.as_str())
|
||||
{
|
||||
Ok(scoring::regex_match(pattern, &submission.response))
|
||||
} else {
|
||||
Err(BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "no 'expected_regex' field for regex scoring".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
"llm" => {
|
||||
// TODO: LLM-as-judge scoring
|
||||
Ok(BenchScore::partial(0.5, "LLM scoring not yet implemented"))
|
||||
}
|
||||
other => Err(BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("unknown scorer: {other}"),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_custom_load_tasks() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "prompt": "What is 2+2?", "expected": "4"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t2", "prompt": "Say hello", "expected_contains": "hello", "scorer": "contains"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = CustomSuite::new(&path);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 2);
|
||||
assert_eq!(tasks[0].id, "t1");
|
||||
assert_eq!(tasks[1].id, "t2");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_custom_exact_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "prompt": "What is 2+2?", "expected": "4"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = CustomSuite::new(&path);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
let submission = TaskSubmission {
|
||||
response: "4".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
assert_eq!(score.label, "pass");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_custom_contains_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "prompt": "Greet me", "expected_contains": "hello", "scorer": "contains"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = CustomSuite::new(&path);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
let submission = TaskSubmission {
|
||||
response: "Hello there!".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
}
|
||||
@@ -1,183 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::scoring;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskResource, TaskSubmission};
|
||||
|
||||
/// GAIA dataset entry (Hugging Face JSONL format).
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GaiaEntry {
|
||||
task_id: String,
|
||||
#[serde(alias = "Question")]
|
||||
question: String,
|
||||
#[serde(alias = "Final answer", alias = "final_answer")]
|
||||
final_answer: String,
|
||||
#[serde(alias = "Level", default)]
|
||||
level: Option<u32>,
|
||||
#[serde(alias = "file_name", default)]
|
||||
file_name: Option<String>,
|
||||
#[serde(alias = "Annotator Metadata", default)]
|
||||
annotator_metadata: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// GAIA benchmark suite.
|
||||
///
|
||||
/// Tasks are loaded from HuggingFace JSONL exports. Scoring uses normalized
|
||||
/// exact match against the `final_answer` field.
|
||||
pub struct GaiaSuite {
|
||||
dataset_path: PathBuf,
|
||||
attachments_dir: Option<PathBuf>,
|
||||
}
|
||||
|
||||
impl GaiaSuite {
|
||||
pub fn new(
|
||||
dataset_path: impl Into<PathBuf>,
|
||||
attachments_dir: Option<impl Into<PathBuf>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
attachments_dir: attachments_dir.map(|d| d.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for GaiaSuite {
|
||||
fn name(&self) -> &str {
|
||||
"GAIA"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"gaia"
|
||||
}
|
||||
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError> {
|
||||
let file = std::fs::File::open(&self.dataset_path)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: GaiaEntry = serde_json::from_str(trimmed)
|
||||
.map_err(|e| BenchError::Config(format!("GAIA line {}: {}", line_num + 1, e)))?;
|
||||
|
||||
let mut resources = Vec::new();
|
||||
if let Some(ref file_name) = entry.file_name {
|
||||
if !file_name.is_empty() {
|
||||
if let Some(ref dir) = self.attachments_dir {
|
||||
resources.push(TaskResource {
|
||||
name: file_name.clone(),
|
||||
path: dir.join(file_name).to_string_lossy().to_string(),
|
||||
resource_type: crate::suite::ResourceType::File,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut tags = Vec::new();
|
||||
if let Some(level) = entry.level {
|
||||
tags.push(format!("level-{level}"));
|
||||
}
|
||||
|
||||
let metadata = serde_json::json!({
|
||||
"expected": entry.final_answer,
|
||||
"level": entry.level,
|
||||
});
|
||||
|
||||
tasks.push(BenchTask {
|
||||
id: entry.task_id,
|
||||
prompt: entry.question,
|
||||
context: None,
|
||||
resources,
|
||||
tags,
|
||||
expected_turns: None,
|
||||
timeout: None,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError> {
|
||||
let expected = task
|
||||
.metadata
|
||||
.get("expected")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "missing expected answer in metadata".to_string(),
|
||||
})?;
|
||||
|
||||
Ok(scoring::exact_match(expected, &submission.response))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_gaia_load_tasks() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("gaia.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"task_id": "g1", "question": "What is the capital of France?", "final_answer": "Paris", "Level": 1}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = GaiaSuite::new(&path, None::<PathBuf>);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks[0].id, "g1");
|
||||
assert!(tasks[0].tags.contains(&"level-1".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_gaia_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("gaia.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"task_id": "g1", "question": "Capital of France?", "final_answer": "Paris"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = GaiaSuite::new(&path, None::<PathBuf>);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
// Exact match (case insensitive)
|
||||
let submission = TaskSubmission {
|
||||
response: "paris".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
|
||||
// Wrong answer
|
||||
let submission = TaskSubmission {
|
||||
response: "London".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
}
|
||||
@@ -1,110 +0,0 @@
|
||||
pub mod custom;
|
||||
pub mod gaia;
|
||||
pub mod swe_bench;
|
||||
pub mod tau_bench;
|
||||
|
||||
use crate::config::BenchConfig;
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::BenchSuite;
|
||||
|
||||
/// List of all known suite IDs.
|
||||
pub const KNOWN_SUITES: &[(&str, &str)] = &[
|
||||
("custom", "Custom JSONL tasks"),
|
||||
("gaia", "GAIA benchmark (knowledge & reasoning)"),
|
||||
("tau_bench", "Tau-bench (multi-turn tool use)"),
|
||||
("swe_bench", "SWE-bench Pro (software engineering)"),
|
||||
];
|
||||
|
||||
/// Create a suite adapter by name.
|
||||
pub fn create_suite(name: &str, config: &BenchConfig) -> Result<Box<dyn BenchSuite>, BenchError> {
|
||||
let suite_map = config.suite_config_map();
|
||||
match name {
|
||||
"custom" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'custom' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
Ok(Box::new(custom::CustomSuite::new(dataset_path)))
|
||||
}
|
||||
"gaia" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'gaia' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
let attachments_dir = suite_map
|
||||
.get("attachments_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
Ok(Box::new(gaia::GaiaSuite::new(
|
||||
dataset_path,
|
||||
attachments_dir,
|
||||
)))
|
||||
}
|
||||
"tau_bench" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'tau_bench' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
let domain = suite_map
|
||||
.get("domain")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("retail")
|
||||
.to_string();
|
||||
Ok(Box::new(tau_bench::TauBenchSuite::new(
|
||||
dataset_path,
|
||||
domain,
|
||||
)))
|
||||
}
|
||||
"swe_bench" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'swe_bench' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
let workspace_dir = suite_map
|
||||
.get("workspace_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("/tmp/swe-bench")
|
||||
.to_string();
|
||||
let use_docker = suite_map
|
||||
.get("use_docker")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
Ok(Box::new(swe_bench::SweBenchSuite::new(
|
||||
dataset_path,
|
||||
workspace_dir,
|
||||
use_docker,
|
||||
)))
|
||||
}
|
||||
_ => {
|
||||
let available = KNOWN_SUITES
|
||||
.iter()
|
||||
.map(|(id, _)| *id)
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
Err(BenchError::SuiteNotFound {
|
||||
name: name.to_string(),
|
||||
available,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,291 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskSubmission};
|
||||
|
||||
/// SWE-bench dataset entry.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SweBenchEntry {
|
||||
instance_id: String,
|
||||
repo: String,
|
||||
base_commit: String,
|
||||
#[serde(default)]
|
||||
problem_statement: String,
|
||||
#[serde(default)]
|
||||
hints_text: Option<String>,
|
||||
#[serde(default)]
|
||||
test_patch: Option<String>,
|
||||
#[serde(default)]
|
||||
patch: Option<String>,
|
||||
}
|
||||
|
||||
/// SWE-bench Pro: real-world software engineering tasks.
|
||||
///
|
||||
/// Each task clones a repo at a specific commit, presents the problem statement,
|
||||
/// and expects the agent to produce a patch. Scoring runs the test suite.
|
||||
pub struct SweBenchSuite {
|
||||
dataset_path: PathBuf,
|
||||
workspace_dir: PathBuf,
|
||||
use_docker: bool,
|
||||
}
|
||||
|
||||
impl SweBenchSuite {
|
||||
pub fn new(
|
||||
dataset_path: impl Into<PathBuf>,
|
||||
workspace_dir: impl Into<PathBuf>,
|
||||
use_docker: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
workspace_dir: workspace_dir.into(),
|
||||
use_docker,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for SweBenchSuite {
|
||||
fn name(&self) -> &str {
|
||||
"SWE-bench Pro"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"swe_bench"
|
||||
}
|
||||
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError> {
|
||||
let file = std::fs::File::open(&self.dataset_path)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: SweBenchEntry = serde_json::from_str(trimmed).map_err(|e| {
|
||||
BenchError::Config(format!("swe_bench line {}: {}", line_num + 1, e))
|
||||
})?;
|
||||
|
||||
let metadata = serde_json::json!({
|
||||
"repo": entry.repo,
|
||||
"base_commit": entry.base_commit,
|
||||
"test_patch": entry.test_patch,
|
||||
"gold_patch": entry.patch,
|
||||
"use_docker": self.use_docker,
|
||||
"workspace_dir": self.workspace_dir.to_string_lossy(),
|
||||
});
|
||||
|
||||
let prompt = if let Some(ref hints) = entry.hints_text {
|
||||
format!("{}\n\nHints:\n{}", entry.problem_statement, hints)
|
||||
} else {
|
||||
entry.problem_statement
|
||||
};
|
||||
|
||||
tasks.push(BenchTask {
|
||||
id: entry.instance_id,
|
||||
prompt,
|
||||
context: Some(format!(
|
||||
"Repository: {}, Commit: {}",
|
||||
entry.repo, entry.base_commit
|
||||
)),
|
||||
resources: vec![],
|
||||
tags: vec![format!("repo-{}", entry.repo.replace('/', "-"))],
|
||||
expected_turns: None,
|
||||
timeout: None,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn setup_task(&self, task: &BenchTask) -> Result<(), BenchError> {
|
||||
let repo = task
|
||||
.metadata
|
||||
.get("repo")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: "missing repo in metadata".to_string(),
|
||||
})?;
|
||||
let base_commit = task
|
||||
.metadata
|
||||
.get("base_commit")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: "missing base_commit in metadata".to_string(),
|
||||
})?;
|
||||
|
||||
let task_dir = self.workspace_dir.join(&task.id);
|
||||
|
||||
// Clone repo if not already present
|
||||
if !task_dir.exists() {
|
||||
let repo_url = format!("https://github.com/{}.git", repo);
|
||||
let output = tokio::process::Command::new("git")
|
||||
.args([
|
||||
"clone",
|
||||
"--depth",
|
||||
"1",
|
||||
&repo_url,
|
||||
&task_dir.to_string_lossy(),
|
||||
])
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git clone failed: {e}"),
|
||||
})?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
return Err(BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git clone failed: {stderr}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Checkout the base commit
|
||||
let output = tokio::process::Command::new("git")
|
||||
.args(["checkout", base_commit])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git checkout failed: {e}"),
|
||||
})?;
|
||||
|
||||
if !output.status.success() {
|
||||
// Shallow clone might not have the commit; fetch more history
|
||||
let _ = tokio::process::Command::new("git")
|
||||
.args(["fetch", "--unshallow"])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await;
|
||||
|
||||
let output = tokio::process::Command::new("git")
|
||||
.args(["checkout", base_commit])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git checkout retry failed: {e}"),
|
||||
})?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
return Err(BenchError::TaskFailed {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("git checkout failed: {stderr}"),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn teardown_task(&self, task: &BenchTask) -> Result<(), BenchError> {
|
||||
let task_dir = self.workspace_dir.join(&task.id);
|
||||
if task_dir.exists() {
|
||||
// Reset any changes
|
||||
let _ = tokio::process::Command::new("git")
|
||||
.args(["checkout", "."])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await;
|
||||
let _ = tokio::process::Command::new("git")
|
||||
.args(["clean", "-fdx"])
|
||||
.current_dir(&task_dir)
|
||||
.output()
|
||||
.await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError> {
|
||||
// For SWE-bench, scoring requires running the test patch against the agent's changes.
|
||||
// This is a simplified version that checks if the agent produced any code changes.
|
||||
|
||||
let test_patch = task.metadata.get("test_patch").and_then(|v| v.as_str());
|
||||
|
||||
if submission.response.is_empty() {
|
||||
return Ok(BenchScore::fail("no response from agent"));
|
||||
}
|
||||
|
||||
// If we have a test patch, try to verify the submission
|
||||
if let Some(_test_patch) = test_patch {
|
||||
// TODO: Apply agent's patch, then apply test patch, then run tests.
|
||||
// For now, give partial credit if the agent produced some output.
|
||||
Ok(BenchScore::partial(
|
||||
0.25,
|
||||
"test execution not yet implemented; partial credit for response",
|
||||
))
|
||||
} else {
|
||||
// No test patch available; can't automatically score
|
||||
Ok(BenchScore::partial(
|
||||
0.25,
|
||||
"no test_patch available for automated scoring",
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_swe_bench_load() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("swe.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"instance_id": "django__django-12345", "repo": "django/django", "base_commit": "abc123", "problem_statement": "Fix the ORM bug"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = SweBenchSuite::new(&path, "/tmp/swe-test", false);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks[0].id, "django__django-12345");
|
||||
assert!(tasks[0].tags.contains(&"repo-django-django".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_swe_bench_scoring_no_response() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("swe.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"instance_id": "s1", "repo": "org/repo", "base_commit": "abc", "problem_statement": "Fix bug"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = SweBenchSuite::new(&path, "/tmp/swe-test", false);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
let submission = TaskSubmission {
|
||||
response: String::new(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec![],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
}
|
||||
@@ -1,227 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, ConversationTurn, TaskSubmission};
|
||||
|
||||
/// Tau-bench task entry.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TauBenchEntry {
|
||||
id: String,
|
||||
#[serde(default)]
|
||||
domain: String,
|
||||
instruction: String,
|
||||
#[serde(default)]
|
||||
user_persona: Option<String>,
|
||||
#[serde(default)]
|
||||
expected_state: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
expected_actions: Vec<String>,
|
||||
#[serde(default)]
|
||||
max_turns: Option<usize>,
|
||||
}
|
||||
|
||||
/// Tau-bench: multi-turn tool-calling dialog benchmark.
|
||||
///
|
||||
/// Tests agent ability to handle customer service scenarios with simulated
|
||||
/// domain APIs (retail, airline). Scoring compares final state against expected.
|
||||
pub struct TauBenchSuite {
|
||||
dataset_path: PathBuf,
|
||||
domain: String,
|
||||
}
|
||||
|
||||
impl TauBenchSuite {
|
||||
pub fn new(dataset_path: impl Into<PathBuf>, domain: impl Into<String>) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
domain: domain.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for TauBenchSuite {
|
||||
fn name(&self) -> &str {
|
||||
"Tau-bench"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"tau_bench"
|
||||
}
|
||||
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError> {
|
||||
let file = std::fs::File::open(&self.dataset_path)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for (line_num, line) in reader.lines().enumerate() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let entry: TauBenchEntry = serde_json::from_str(trimmed).map_err(|e| {
|
||||
BenchError::Config(format!("tau_bench line {}: {}", line_num + 1, e))
|
||||
})?;
|
||||
|
||||
let domain = if entry.domain.is_empty() {
|
||||
self.domain.clone()
|
||||
} else {
|
||||
entry.domain.clone()
|
||||
};
|
||||
|
||||
let metadata = serde_json::json!({
|
||||
"domain": domain,
|
||||
"user_persona": entry.user_persona,
|
||||
"expected_state": entry.expected_state,
|
||||
"expected_actions": entry.expected_actions,
|
||||
});
|
||||
|
||||
tasks.push(BenchTask {
|
||||
id: entry.id,
|
||||
prompt: entry.instruction,
|
||||
context: entry.user_persona.clone(),
|
||||
resources: vec![],
|
||||
tags: vec![format!("domain-{domain}")],
|
||||
expected_turns: entry.max_turns,
|
||||
timeout: None,
|
||||
metadata,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(tasks)
|
||||
}
|
||||
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError> {
|
||||
// Score based on expected actions completion
|
||||
let expected_actions: Vec<String> = task
|
||||
.metadata
|
||||
.get("expected_actions")
|
||||
.and_then(|v| serde_json::from_value(v.clone()).ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
if expected_actions.is_empty() {
|
||||
// No expected actions defined; score based on whether agent responded
|
||||
if submission.response.is_empty() {
|
||||
return Ok(BenchScore::fail("no response"));
|
||||
}
|
||||
return Ok(BenchScore::partial(
|
||||
0.5,
|
||||
"no expected_actions to evaluate against",
|
||||
));
|
||||
}
|
||||
|
||||
// Check which expected actions were actually called
|
||||
let called: std::collections::HashSet<&str> =
|
||||
submission.tool_calls.iter().map(|s| s.as_str()).collect();
|
||||
let matched = expected_actions
|
||||
.iter()
|
||||
.filter(|a| called.contains(a.as_str()))
|
||||
.count();
|
||||
|
||||
let ratio = matched as f64 / expected_actions.len() as f64;
|
||||
if ratio >= 1.0 {
|
||||
Ok(BenchScore::pass())
|
||||
} else if ratio > 0.0 {
|
||||
Ok(BenchScore::partial(
|
||||
ratio,
|
||||
format!(
|
||||
"{}/{} expected actions completed",
|
||||
matched,
|
||||
expected_actions.len()
|
||||
),
|
||||
))
|
||||
} else {
|
||||
Ok(BenchScore::fail(format!(
|
||||
"0/{} expected actions completed",
|
||||
expected_actions.len()
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_user_message(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
conversation: &[ConversationTurn],
|
||||
) -> Result<Option<String>, BenchError> {
|
||||
// Check if we've exceeded max turns
|
||||
if let Some(max) = task.expected_turns {
|
||||
let user_turns = conversation
|
||||
.iter()
|
||||
.filter(|t| matches!(t.role, crate::suite::TurnRole::User))
|
||||
.count();
|
||||
if user_turns >= max {
|
||||
return Ok(None);
|
||||
}
|
||||
}
|
||||
|
||||
// For now, multi-turn simulation requires an LLM (not implemented yet).
|
||||
// Return None to end after the first turn.
|
||||
// TODO: Use LLM to simulate customer based on user_persona.
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tau_bench_load() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tau.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "instruction": "Return my order", "expected_actions": ["lookup_order", "process_return"], "max_turns": 3}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = TauBenchSuite::new(&path, "retail");
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 1);
|
||||
assert_eq!(tasks[0].expected_turns, Some(3));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_tau_bench_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("tau.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "t1", "instruction": "Return order", "expected_actions": ["lookup_order", "process_return"]}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = TauBenchSuite::new(&path, "retail");
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
// Partial completion
|
||||
let submission = TaskSubmission {
|
||||
response: "I found your order.".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec!["lookup_order".to_string()],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 0.5);
|
||||
assert_eq!(score.label, "partial");
|
||||
|
||||
// Full completion
|
||||
let submission = TaskSubmission {
|
||||
response: "Return processed.".to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: vec!["lookup_order".to_string(), "process_return".to_string()],
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
}
|
||||
@@ -1,229 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::{Mutex, mpsc};
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
use ironclaw::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate};
|
||||
use ironclaw::error::ChannelError;
|
||||
|
||||
use crate::results::TraceToolCall;
|
||||
use crate::suite::ConversationTurn;
|
||||
|
||||
/// Captured state from a benchmark channel run.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct ChannelCapture {
|
||||
/// All responses the agent sent back.
|
||||
pub responses: Vec<String>,
|
||||
/// Tool calls observed (name, success, duration_ms).
|
||||
pub tool_calls: Vec<TraceToolCall>,
|
||||
/// Full conversation turns for multi-turn scoring.
|
||||
pub conversation: Vec<ConversationTurn>,
|
||||
/// Status messages (for debugging).
|
||||
pub status_log: Vec<String>,
|
||||
}
|
||||
|
||||
/// A headless Channel implementation for benchmarking.
|
||||
///
|
||||
/// Modeled after `ReplChannel`: uses mpsc to inject messages and captures
|
||||
/// all responses and tool status events. Auto-approves tool execution
|
||||
/// so benchmarks run without user interaction.
|
||||
pub struct BenchChannel {
|
||||
/// Sender to inject messages into the agent loop.
|
||||
msg_tx: mpsc::Sender<IncomingMessage>,
|
||||
/// Receiver the agent loop reads from (taken once by `start()`).
|
||||
msg_rx: Mutex<Option<mpsc::Receiver<IncomingMessage>>>,
|
||||
/// Accumulated capture data.
|
||||
capture: Arc<Mutex<ChannelCapture>>,
|
||||
}
|
||||
|
||||
impl BenchChannel {
|
||||
pub fn new() -> (Self, mpsc::Sender<IncomingMessage>) {
|
||||
let (tx, rx) = mpsc::channel(64);
|
||||
let channel = Self {
|
||||
msg_tx: tx.clone(),
|
||||
msg_rx: Mutex::new(Some(rx)),
|
||||
capture: Arc::new(Mutex::new(ChannelCapture::default())),
|
||||
};
|
||||
(channel, tx)
|
||||
}
|
||||
|
||||
/// Get a handle to the capture data.
|
||||
pub fn capture(&self) -> Arc<Mutex<ChannelCapture>> {
|
||||
Arc::clone(&self.capture)
|
||||
}
|
||||
|
||||
/// Get a clone of the message sender for injecting follow-up messages.
|
||||
pub fn sender(&self) -> mpsc::Sender<IncomingMessage> {
|
||||
self.msg_tx.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for BenchChannel {
|
||||
fn name(&self) -> &str {
|
||||
"bench"
|
||||
}
|
||||
|
||||
async fn start(&self) -> Result<MessageStream, ChannelError> {
|
||||
let rx = self
|
||||
.msg_rx
|
||||
.lock()
|
||||
.await
|
||||
.take()
|
||||
.ok_or_else(|| ChannelError::StartupFailed {
|
||||
name: "bench".to_string(),
|
||||
reason: "start() already called".to_string(),
|
||||
})?;
|
||||
Ok(Box::pin(ReceiverStream::new(rx)))
|
||||
}
|
||||
|
||||
async fn respond(
|
||||
&self,
|
||||
_msg: &IncomingMessage,
|
||||
response: OutgoingResponse,
|
||||
) -> Result<(), ChannelError> {
|
||||
let mut cap = self.capture.lock().await;
|
||||
cap.responses.push(response.content.clone());
|
||||
cap.conversation.push(ConversationTurn {
|
||||
role: crate::suite::TurnRole::Assistant,
|
||||
content: response.content,
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_status(
|
||||
&self,
|
||||
status: StatusUpdate,
|
||||
_metadata: &serde_json::Value,
|
||||
) -> Result<(), ChannelError> {
|
||||
let mut cap = self.capture.lock().await;
|
||||
|
||||
match status {
|
||||
StatusUpdate::ToolCompleted { ref name, success } => {
|
||||
cap.tool_calls.push(TraceToolCall {
|
||||
name: name.clone(),
|
||||
duration_ms: 0, // We don't have precise per-tool timing here
|
||||
success,
|
||||
});
|
||||
cap.status_log
|
||||
.push(format!("tool_completed: {name} success={success}"));
|
||||
}
|
||||
StatusUpdate::ApprovalNeeded { ref request_id, .. } => {
|
||||
// Auto-approve all tools during benchmarks
|
||||
cap.status_log.push(format!("auto_approved: {request_id}"));
|
||||
drop(cap); // Release lock before sending
|
||||
let approval = IncomingMessage::new("bench", "bench-user", "always");
|
||||
let _ = self.msg_tx.send(approval).await;
|
||||
return Ok(());
|
||||
}
|
||||
StatusUpdate::Thinking(ref msg) => {
|
||||
cap.status_log.push(format!("thinking: {msg}"));
|
||||
}
|
||||
StatusUpdate::ToolStarted { ref name } => {
|
||||
cap.status_log.push(format!("tool_started: {name}"));
|
||||
}
|
||||
StatusUpdate::ToolResult {
|
||||
ref name,
|
||||
ref preview,
|
||||
} => {
|
||||
cap.status_log.push(format!(
|
||||
"tool_result: {name} -> {}",
|
||||
&preview[..preview.len().min(100)]
|
||||
));
|
||||
}
|
||||
StatusUpdate::StreamChunk(_) => {}
|
||||
StatusUpdate::Status(ref msg) => {
|
||||
cap.status_log.push(format!("status: {msg}"));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn broadcast(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
response: OutgoingResponse,
|
||||
) -> Result<(), ChannelError> {
|
||||
let mut cap = self.capture.lock().await;
|
||||
cap.status_log.push(format!(
|
||||
"broadcast: {}",
|
||||
&response.content[..response.content.len().min(100)]
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> Result<(), ChannelError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn shutdown(&self) -> Result<(), ChannelError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_captures_responses() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
let capture = channel.capture();
|
||||
|
||||
let msg = IncomingMessage::new("bench", "user", "hello");
|
||||
let response = OutgoingResponse::text("world");
|
||||
channel.respond(&msg, response).await.unwrap();
|
||||
|
||||
let cap = capture.lock().await;
|
||||
assert_eq!(cap.responses.len(), 1);
|
||||
assert_eq!(cap.responses[0], "world");
|
||||
assert_eq!(cap.conversation.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_auto_approves() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
// start() to consume the receiver
|
||||
let _stream = channel.start().await.unwrap();
|
||||
|
||||
let status = StatusUpdate::ApprovalNeeded {
|
||||
request_id: "req-1".to_string(),
|
||||
tool_name: "shell".to_string(),
|
||||
description: "run ls".to_string(),
|
||||
parameters: serde_json::json!({}),
|
||||
};
|
||||
channel
|
||||
.send_status(status, &serde_json::Value::Null)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// The approval message was sent through msg_tx,
|
||||
// which means the stream would receive it.
|
||||
// We can't easily read from the stream in this test without
|
||||
// consuming it, but we can verify the status log.
|
||||
let capture_arc = channel.capture();
|
||||
let cap = capture_arc.lock().await;
|
||||
assert!(cap.status_log.iter().any(|s| s.contains("auto_approved")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_captures_tool_events() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
|
||||
let status = StatusUpdate::ToolCompleted {
|
||||
name: "echo".to_string(),
|
||||
success: true,
|
||||
};
|
||||
channel
|
||||
.send_status(status, &serde_json::Value::Null)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let capture_arc = channel.capture();
|
||||
let cap = capture_arc.lock().await;
|
||||
assert_eq!(cap.tool_calls.len(), 1);
|
||||
assert_eq!(cap.tool_calls[0].name, "echo");
|
||||
assert!(cap.tool_calls[0].success);
|
||||
}
|
||||
}
|
||||
@@ -1,203 +0,0 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
|
||||
/// Top-level bench configuration, loaded from TOML.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct BenchConfig {
|
||||
/// Where to write results. Default: "./bench-results".
|
||||
#[serde(default = "default_results_dir")]
|
||||
pub results_dir: PathBuf,
|
||||
|
||||
/// Per-task timeout. Default: "300s".
|
||||
#[serde(
|
||||
default = "default_task_timeout",
|
||||
deserialize_with = "deserialize_duration"
|
||||
)]
|
||||
pub task_timeout: Duration,
|
||||
|
||||
/// Max agent iterations per task. Default: 15.
|
||||
#[serde(default = "default_max_iterations")]
|
||||
pub max_iterations: usize,
|
||||
|
||||
/// How many tasks to run in parallel. Default: 1.
|
||||
#[serde(default = "default_parallelism")]
|
||||
pub parallelism: usize,
|
||||
|
||||
/// Model/config matrix entries. At least one required.
|
||||
#[serde(default)]
|
||||
pub matrix: Vec<MatrixEntry>,
|
||||
|
||||
/// Suite-specific configuration (passed through to adapter).
|
||||
#[serde(default = "default_suite_config")]
|
||||
pub suite_config: toml::Value,
|
||||
}
|
||||
|
||||
/// A single model/config combination to benchmark.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct MatrixEntry {
|
||||
/// Label for this configuration (used in results).
|
||||
pub label: String,
|
||||
|
||||
/// Model identifier.
|
||||
#[serde(default)]
|
||||
pub model: Option<String>,
|
||||
|
||||
/// Optional tool allowlist. If set, only these tools are available.
|
||||
#[serde(default)]
|
||||
pub tools: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl BenchConfig {
|
||||
/// Load from a TOML file.
|
||||
pub fn from_file(path: &Path) -> Result<Self, BenchError> {
|
||||
if !path.exists() {
|
||||
return Err(BenchError::ConfigNotFound {
|
||||
path: path.to_path_buf(),
|
||||
});
|
||||
}
|
||||
let content = std::fs::read_to_string(path)?;
|
||||
let config: BenchConfig = toml::from_str(&content)?;
|
||||
Ok(config)
|
||||
}
|
||||
|
||||
/// Create a minimal config for when no config file is provided.
|
||||
/// Uses defaults and optional CLI overrides.
|
||||
pub fn minimal(model: Option<String>) -> Self {
|
||||
let label = model.as_deref().unwrap_or("default").to_string();
|
||||
Self {
|
||||
results_dir: default_results_dir(),
|
||||
task_timeout: default_task_timeout(),
|
||||
max_iterations: default_max_iterations(),
|
||||
parallelism: default_parallelism(),
|
||||
matrix: vec![MatrixEntry {
|
||||
label,
|
||||
model,
|
||||
tools: None,
|
||||
}],
|
||||
suite_config: toml::Value::Table(toml::map::Map::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the suite_config as a generic map for adapter use.
|
||||
pub fn suite_config_map(&self) -> toml::map::Map<String, toml::Value> {
|
||||
match &self.suite_config {
|
||||
toml::Value::Table(map) => map.clone(),
|
||||
_ => toml::map::Map::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a string value from suite_config.
|
||||
pub fn suite_config_str(&self, key: &str) -> Option<String> {
|
||||
self.suite_config_map()
|
||||
.get(key)
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn default_suite_config() -> toml::Value {
|
||||
toml::Value::Table(toml::map::Map::new())
|
||||
}
|
||||
|
||||
fn default_results_dir() -> PathBuf {
|
||||
PathBuf::from("./bench-results")
|
||||
}
|
||||
|
||||
fn default_task_timeout() -> Duration {
|
||||
Duration::from_secs(300)
|
||||
}
|
||||
|
||||
fn default_max_iterations() -> usize {
|
||||
15
|
||||
}
|
||||
|
||||
fn default_parallelism() -> usize {
|
||||
1
|
||||
}
|
||||
|
||||
/// Deserialize a duration from a string like "300s", "5m", etc.
|
||||
fn deserialize_duration<'de, D>(deserializer: D) -> Result<Duration, D::Error>
|
||||
where
|
||||
D: serde::Deserializer<'de>,
|
||||
{
|
||||
let s = String::deserialize(deserializer)?;
|
||||
parse_duration(&s).map_err(serde::de::Error::custom)
|
||||
}
|
||||
|
||||
fn parse_duration(s: &str) -> Result<Duration, String> {
|
||||
let s = s.trim();
|
||||
if let Some(secs) = s.strip_suffix('s') {
|
||||
secs.trim()
|
||||
.parse::<u64>()
|
||||
.map(Duration::from_secs)
|
||||
.map_err(|e| format!("invalid seconds: {e}"))
|
||||
} else if let Some(mins) = s.strip_suffix('m') {
|
||||
mins.trim()
|
||||
.parse::<u64>()
|
||||
.map(|m| Duration::from_secs(m * 60))
|
||||
.map_err(|e| format!("invalid minutes: {e}"))
|
||||
} else {
|
||||
// Assume seconds if no suffix
|
||||
s.parse::<u64>()
|
||||
.map(Duration::from_secs)
|
||||
.map_err(|e| format!("invalid duration '{s}': {e}"))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_duration() {
|
||||
assert_eq!(parse_duration("300s").unwrap(), Duration::from_secs(300));
|
||||
assert_eq!(parse_duration("5m").unwrap(), Duration::from_secs(300));
|
||||
assert_eq!(parse_duration("60").unwrap(), Duration::from_secs(60));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_minimal_config() {
|
||||
let config = BenchConfig::minimal(Some("test-model".to_string()));
|
||||
assert_eq!(config.matrix.len(), 1);
|
||||
assert_eq!(config.matrix[0].label, "test-model");
|
||||
assert_eq!(config.max_iterations, 15);
|
||||
assert_eq!(config.parallelism, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_from_toml() {
|
||||
let toml_str = r#"
|
||||
results_dir = "./my-results"
|
||||
task_timeout = "60s"
|
||||
max_iterations = 10
|
||||
parallelism = 2
|
||||
|
||||
[[matrix]]
|
||||
label = "fast"
|
||||
model = "gpt-4o-mini"
|
||||
|
||||
[[matrix]]
|
||||
label = "full"
|
||||
model = "claude-3-5-sonnet"
|
||||
tools = ["echo", "time"]
|
||||
|
||||
[suite_config]
|
||||
dataset_path = "./data/test.jsonl"
|
||||
"#;
|
||||
let config: BenchConfig = toml::from_str(toml_str).unwrap();
|
||||
assert_eq!(config.results_dir, PathBuf::from("./my-results"));
|
||||
assert_eq!(config.task_timeout, Duration::from_secs(60));
|
||||
assert_eq!(config.max_iterations, 10);
|
||||
assert_eq!(config.parallelism, 2);
|
||||
assert_eq!(config.matrix.len(), 2);
|
||||
assert_eq!(config.matrix[1].tools.as_ref().unwrap().len(), 2);
|
||||
assert_eq!(
|
||||
config.suite_config_str("dataset_path").unwrap(),
|
||||
"./data/test.jsonl"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum BenchError {
|
||||
#[error("Config error: {0}")]
|
||||
Config(String),
|
||||
|
||||
#[error("Config file not found: {path}")]
|
||||
ConfigNotFound { path: PathBuf },
|
||||
|
||||
#[error("Suite {name} not found. Available: {available}")]
|
||||
SuiteNotFound { name: String, available: String },
|
||||
|
||||
#[error("Task {task_id} failed: {reason}")]
|
||||
TaskFailed { task_id: String, reason: String },
|
||||
|
||||
#[error("Timeout after {seconds}s for task {task_id}")]
|
||||
Timeout { task_id: String, seconds: u64 },
|
||||
|
||||
#[error("Scoring error for task {task_id}: {reason}")]
|
||||
Scoring { task_id: String, reason: String },
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
#[error("JSON error: {0}")]
|
||||
Json(#[from] serde_json::Error),
|
||||
|
||||
#[error("TOML parse error: {0}")]
|
||||
Toml(#[from] toml::de::Error),
|
||||
|
||||
#[error("Agent error: {0}")]
|
||||
Agent(#[from] ironclaw::Error),
|
||||
|
||||
#[error("Results directory error: {0}")]
|
||||
ResultsDir(String),
|
||||
|
||||
#[error("Resume failed: no completed tasks found in {path}")]
|
||||
ResumeEmpty { path: PathBuf },
|
||||
}
|
||||
@@ -1,248 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU32, Ordering};
|
||||
use std::time::Instant;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use rust_decimal::Decimal;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use ironclaw::error::LlmError;
|
||||
use ironclaw::llm::{
|
||||
CompletionRequest, CompletionResponse, LlmProvider, ToolCompletionRequest,
|
||||
ToolCompletionResponse,
|
||||
};
|
||||
|
||||
/// Recorded metrics from a single LLM call.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LlmCallRecord {
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub duration_ms: u64,
|
||||
pub had_tool_calls: bool,
|
||||
}
|
||||
|
||||
/// Wraps an `LlmProvider` to record per-call metrics.
|
||||
///
|
||||
/// The wrapper is transparent to the agent: it delegates every call
|
||||
/// to the inner provider and captures token counts and timings.
|
||||
pub struct InstrumentedLlm {
|
||||
inner: Arc<dyn LlmProvider>,
|
||||
records: Mutex<Vec<LlmCallRecord>>,
|
||||
total_input_tokens: AtomicU32,
|
||||
total_output_tokens: AtomicU32,
|
||||
call_count: AtomicU32,
|
||||
}
|
||||
|
||||
impl InstrumentedLlm {
|
||||
pub fn new(inner: Arc<dyn LlmProvider>) -> Self {
|
||||
Self {
|
||||
inner,
|
||||
records: Mutex::new(Vec::new()),
|
||||
total_input_tokens: AtomicU32::new(0),
|
||||
total_output_tokens: AtomicU32::new(0),
|
||||
call_count: AtomicU32::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Take all recorded call metrics, clearing the internal buffer.
|
||||
pub async fn take_records(&self) -> Vec<LlmCallRecord> {
|
||||
let mut records = self.records.lock().await;
|
||||
std::mem::take(&mut *records)
|
||||
}
|
||||
|
||||
/// Snapshot of total tokens without clearing.
|
||||
pub fn total_input_tokens(&self) -> u32 {
|
||||
self.total_input_tokens.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn total_output_tokens(&self) -> u32 {
|
||||
self.total_output_tokens.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn call_count(&self) -> u32 {
|
||||
self.call_count.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
/// Estimated cost using the inner provider's cost-per-token rates.
|
||||
pub fn estimated_cost(&self) -> f64 {
|
||||
let (input_rate, output_rate) = self.inner.cost_per_token();
|
||||
let input_cost =
|
||||
input_rate * Decimal::from(self.total_input_tokens.load(Ordering::Relaxed));
|
||||
let output_cost =
|
||||
output_rate * Decimal::from(self.total_output_tokens.load(Ordering::Relaxed));
|
||||
let total = input_cost + output_cost;
|
||||
// Convert Decimal to f64 for the trace (benchmarks don't need exact precision)
|
||||
total.to_string().parse::<f64>().unwrap_or(0.0)
|
||||
}
|
||||
|
||||
/// Reset all counters and records.
|
||||
pub async fn reset(&self) {
|
||||
self.records.lock().await.clear();
|
||||
self.total_input_tokens.store(0, Ordering::Relaxed);
|
||||
self.total_output_tokens.store(0, Ordering::Relaxed);
|
||||
self.call_count.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
async fn record(
|
||||
&self,
|
||||
input_tokens: u32,
|
||||
output_tokens: u32,
|
||||
duration_ms: u64,
|
||||
had_tool_calls: bool,
|
||||
) {
|
||||
self.total_input_tokens
|
||||
.fetch_add(input_tokens, Ordering::Relaxed);
|
||||
self.total_output_tokens
|
||||
.fetch_add(output_tokens, Ordering::Relaxed);
|
||||
self.call_count.fetch_add(1, Ordering::Relaxed);
|
||||
self.records.lock().await.push(LlmCallRecord {
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
duration_ms,
|
||||
had_tool_calls,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for InstrumentedLlm {
|
||||
fn model_name(&self) -> &str {
|
||||
self.inner.model_name()
|
||||
}
|
||||
|
||||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||
self.inner.cost_per_token()
|
||||
}
|
||||
|
||||
async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
let start = Instant::now();
|
||||
let response = self.inner.complete(request).await?;
|
||||
let elapsed = start.elapsed().as_millis() as u64;
|
||||
self.record(
|
||||
response.input_tokens,
|
||||
response.output_tokens,
|
||||
elapsed,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
let start = Instant::now();
|
||||
let response = self.inner.complete_with_tools(request).await?;
|
||||
let elapsed = start.elapsed().as_millis() as u64;
|
||||
let had_tool_calls = !response.tool_calls.is_empty();
|
||||
self.record(
|
||||
response.input_tokens,
|
||||
response.output_tokens,
|
||||
elapsed,
|
||||
had_tool_calls,
|
||||
)
|
||||
.await;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
|
||||
self.inner.list_models().await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use ironclaw::llm::{ChatMessage, CompletionRequest, CompletionResponse, FinishReason};
|
||||
|
||||
/// Fake LLM that returns a canned response with known token counts.
|
||||
struct FakeLlm;
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for FakeLlm {
|
||||
fn model_name(&self) -> &str {
|
||||
"fake-model"
|
||||
}
|
||||
|
||||
fn cost_per_token(&self) -> (Decimal, Decimal) {
|
||||
(
|
||||
Decimal::new(3, 6), // $0.000003 per input token
|
||||
Decimal::new(15, 6), // $0.000015 per output token
|
||||
)
|
||||
}
|
||||
|
||||
async fn complete(
|
||||
&self,
|
||||
_request: CompletionRequest,
|
||||
) -> Result<CompletionResponse, LlmError> {
|
||||
Ok(CompletionResponse {
|
||||
content: "test response".to_string(),
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
finish_reason: FinishReason::Stop,
|
||||
})
|
||||
}
|
||||
|
||||
async fn complete_with_tools(
|
||||
&self,
|
||||
_request: ToolCompletionRequest,
|
||||
) -> Result<ToolCompletionResponse, LlmError> {
|
||||
Ok(ToolCompletionResponse {
|
||||
content: Some("tool response".to_string()),
|
||||
tool_calls: vec![],
|
||||
input_tokens: 200,
|
||||
output_tokens: 100,
|
||||
finish_reason: FinishReason::Stop,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_records_metrics() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
|
||||
assert_eq!(instrumented.call_count(), 1);
|
||||
assert_eq!(instrumented.total_input_tokens(), 100);
|
||||
assert_eq!(instrumented.total_output_tokens(), 50);
|
||||
|
||||
let records = instrumented.take_records().await;
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].input_tokens, 100);
|
||||
assert!(!records[0].had_tool_calls);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_cost_calculation() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
|
||||
// 100 * 0.000003 + 50 * 0.000015 = 0.0003 + 0.00075 = 0.00105
|
||||
let cost = instrumented.estimated_cost();
|
||||
assert!((cost - 0.00105).abs() < 0.0001);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_reset() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
assert_eq!(instrumented.call_count(), 1);
|
||||
|
||||
instrumented.reset().await;
|
||||
assert_eq!(instrumented.call_count(), 0);
|
||||
assert_eq!(instrumented.total_input_tokens(), 0);
|
||||
|
||||
let records = instrumented.take_records().await;
|
||||
assert!(records.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -1,314 +0,0 @@
|
||||
mod adapters;
|
||||
mod channel;
|
||||
mod config;
|
||||
mod error;
|
||||
mod instrumented_llm;
|
||||
mod results;
|
||||
mod runner;
|
||||
mod scoring;
|
||||
mod suite;
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::{Parser, Subcommand};
|
||||
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::config::BenchConfig;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "ironclaw-bench", about = "IronClaw benchmarking harness")]
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
command: Commands,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum Commands {
|
||||
/// Run a benchmark suite.
|
||||
Run {
|
||||
/// Suite to run (custom, gaia, tau_bench, swe_bench).
|
||||
#[arg(long)]
|
||||
suite: String,
|
||||
|
||||
/// Path to bench config TOML.
|
||||
#[arg(long)]
|
||||
config: Option<PathBuf>,
|
||||
|
||||
/// Override model for all matrix entries.
|
||||
#[arg(long)]
|
||||
model: Option<String>,
|
||||
|
||||
/// Max tasks to run in parallel.
|
||||
#[arg(long)]
|
||||
parallelism: Option<usize>,
|
||||
|
||||
/// Sample N tasks from the suite (for quick testing).
|
||||
#[arg(long)]
|
||||
sample: Option<usize>,
|
||||
|
||||
/// Only run these task IDs (comma-separated).
|
||||
#[arg(long, value_delimiter = ',')]
|
||||
task_ids: Option<Vec<String>>,
|
||||
|
||||
/// Only run tasks with these tags (comma-separated).
|
||||
#[arg(long, value_delimiter = ',')]
|
||||
tags: Option<Vec<String>>,
|
||||
|
||||
/// Per-task timeout in seconds.
|
||||
#[arg(long)]
|
||||
timeout_secs: Option<u64>,
|
||||
|
||||
/// Override results directory.
|
||||
#[arg(long)]
|
||||
results_dir: Option<PathBuf>,
|
||||
|
||||
/// Resume a previous run by ID.
|
||||
#[arg(long)]
|
||||
resume: Option<Uuid>,
|
||||
},
|
||||
|
||||
/// Show results for a run.
|
||||
Results {
|
||||
/// Run ID or "latest".
|
||||
#[arg(default_value = "latest")]
|
||||
run_id: String,
|
||||
|
||||
/// Output format.
|
||||
#[arg(long, default_value = "table")]
|
||||
format: ResultsFormat,
|
||||
|
||||
/// Override results directory.
|
||||
#[arg(long)]
|
||||
results_dir: Option<PathBuf>,
|
||||
},
|
||||
|
||||
/// Compare two runs.
|
||||
Compare {
|
||||
/// Baseline run ID.
|
||||
baseline: Uuid,
|
||||
|
||||
/// Comparison run ID.
|
||||
comparison: Uuid,
|
||||
|
||||
/// Override results directory.
|
||||
#[arg(long)]
|
||||
results_dir: Option<PathBuf>,
|
||||
},
|
||||
|
||||
/// List available benchmark suites.
|
||||
List,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, clap::ValueEnum)]
|
||||
enum ResultsFormat {
|
||||
Table,
|
||||
Json,
|
||||
Csv,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let cli = Cli::parse();
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("ironclaw_bench=info,ironclaw=warn")),
|
||||
)
|
||||
.with(tracing_subscriber::fmt::layer().with_target(false))
|
||||
.init();
|
||||
|
||||
match cli.command {
|
||||
Commands::List => {
|
||||
println!("Available benchmark suites:\n");
|
||||
for (id, desc) in adapters::KNOWN_SUITES {
|
||||
println!(" {:<15} {}", id, desc);
|
||||
}
|
||||
println!();
|
||||
}
|
||||
Commands::Run {
|
||||
suite,
|
||||
config: config_path,
|
||||
model,
|
||||
parallelism,
|
||||
sample,
|
||||
task_ids,
|
||||
tags,
|
||||
timeout_secs,
|
||||
results_dir,
|
||||
resume,
|
||||
} => {
|
||||
// Load or create config
|
||||
let mut bench_config = if let Some(ref path) = config_path {
|
||||
BenchConfig::from_file(path)?
|
||||
} else {
|
||||
BenchConfig::minimal(model.clone())
|
||||
};
|
||||
|
||||
// Apply CLI overrides
|
||||
if let Some(p) = parallelism {
|
||||
bench_config.parallelism = p;
|
||||
}
|
||||
if let Some(t) = timeout_secs {
|
||||
bench_config.task_timeout = std::time::Duration::from_secs(t);
|
||||
}
|
||||
if let Some(ref dir) = results_dir {
|
||||
bench_config.results_dir = dir.clone();
|
||||
}
|
||||
|
||||
// If model override specified and we have matrix entries, update them
|
||||
if let Some(ref m) = model {
|
||||
for entry in &mut bench_config.matrix {
|
||||
entry.model = Some(m.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Create suite
|
||||
let bench_suite = adapters::create_suite(&suite, &bench_config)?;
|
||||
|
||||
// Initialize ironclaw LLM provider
|
||||
let ironclaw_config = ironclaw::Config::from_env().map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"Failed to load ironclaw config: {}. Make sure .env is configured.",
|
||||
e
|
||||
)
|
||||
})?;
|
||||
|
||||
let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig {
|
||||
auth_base_url: ironclaw_config.llm.nearai.auth_base_url.clone(),
|
||||
session_path: ironclaw_config.llm.nearai.session_path.clone(),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
session.ensure_authenticated().await?;
|
||||
|
||||
let llm = ironclaw::llm::create_llm_provider(&ironclaw_config.llm, session)?;
|
||||
let safety = Arc::new(ironclaw::safety::SafetyLayer::new(&ironclaw_config.safety));
|
||||
|
||||
let runner = runner::BenchRunner::new(bench_suite, bench_config.clone(), llm, safety);
|
||||
|
||||
// Run for each matrix entry
|
||||
for matrix_entry in &bench_config.matrix {
|
||||
let run_id = runner
|
||||
.run(
|
||||
matrix_entry,
|
||||
sample,
|
||||
task_ids.as_deref(),
|
||||
tags.as_deref(),
|
||||
resume,
|
||||
)
|
||||
.await?;
|
||||
println!("Run complete: {}", run_id);
|
||||
}
|
||||
}
|
||||
Commands::Results {
|
||||
run_id,
|
||||
format,
|
||||
results_dir,
|
||||
} => {
|
||||
let base = results_dir.unwrap_or_else(|| PathBuf::from("./bench-results"));
|
||||
let uuid = if run_id == "latest" {
|
||||
results::find_latest_run(&base)?
|
||||
.ok_or_else(|| anyhow::anyhow!("No runs found in {}", base.display()))?
|
||||
} else {
|
||||
Uuid::parse_str(&run_id)?
|
||||
};
|
||||
|
||||
let json_path = results::run_json_path(&base, uuid);
|
||||
let jsonl_path = results::tasks_jsonl_path(&base, uuid);
|
||||
|
||||
let run = results::read_run_result(&json_path)?;
|
||||
let tasks = results::read_task_results(&jsonl_path)?;
|
||||
|
||||
match format {
|
||||
ResultsFormat::Table => {
|
||||
results::print_results_table(&tasks, &run);
|
||||
}
|
||||
ResultsFormat::Json => {
|
||||
let output = serde_json::json!({
|
||||
"run": run,
|
||||
"tasks": tasks,
|
||||
});
|
||||
println!("{}", serde_json::to_string_pretty(&output)?);
|
||||
}
|
||||
ResultsFormat::Csv => {
|
||||
println!("task_id,score,label,tokens,cost,turns,time_s");
|
||||
for task in &tasks {
|
||||
println!(
|
||||
"{},{:.3},{},{},{:.4},{},{:.1}",
|
||||
task.task_id,
|
||||
task.score.value,
|
||||
task.score.label,
|
||||
task.trace.input_tokens + task.trace.output_tokens,
|
||||
task.trace.estimated_cost_usd,
|
||||
task.trace.turns,
|
||||
task.trace.wall_time_ms as f64 / 1000.0,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Commands::Compare {
|
||||
baseline,
|
||||
comparison,
|
||||
results_dir,
|
||||
} => {
|
||||
let base = results_dir.unwrap_or_else(|| PathBuf::from("./bench-results"));
|
||||
|
||||
let baseline_run = results::read_run_result(&results::run_json_path(&base, baseline))?;
|
||||
let comparison_run =
|
||||
results::read_run_result(&results::run_json_path(&base, comparison))?;
|
||||
|
||||
println!("\nComparison: {} vs {}\n", baseline, comparison);
|
||||
println!(
|
||||
"{:<20} {:>12} {:>12} {:>10}",
|
||||
"Metric", "Baseline", "Comparison", "Delta"
|
||||
);
|
||||
println!("{}", "-".repeat(58));
|
||||
|
||||
let pass_delta = comparison_run.pass_rate - baseline_run.pass_rate;
|
||||
println!(
|
||||
"{:<20} {:>11.1}% {:>11.1}% {:>+9.1}%",
|
||||
"Pass rate",
|
||||
baseline_run.pass_rate * 100.0,
|
||||
comparison_run.pass_rate * 100.0,
|
||||
pass_delta * 100.0,
|
||||
);
|
||||
|
||||
let score_delta = comparison_run.avg_score - baseline_run.avg_score;
|
||||
println!(
|
||||
"{:<20} {:>12.3} {:>12.3} {:>+10.3}",
|
||||
"Avg score", baseline_run.avg_score, comparison_run.avg_score, score_delta,
|
||||
);
|
||||
|
||||
let cost_delta = comparison_run.total_cost_usd - baseline_run.total_cost_usd;
|
||||
println!(
|
||||
"{:<20} {:>11.4}$ {:>11.4}$ {:>+9.4}$",
|
||||
"Total cost",
|
||||
baseline_run.total_cost_usd,
|
||||
comparison_run.total_cost_usd,
|
||||
cost_delta,
|
||||
);
|
||||
|
||||
let time_b = baseline_run.total_wall_time_ms as f64 / 1000.0;
|
||||
let time_c = comparison_run.total_wall_time_ms as f64 / 1000.0;
|
||||
println!(
|
||||
"{:<20} {:>11.1}s {:>11.1}s {:>+9.1}s",
|
||||
"Total time",
|
||||
time_b,
|
||||
time_c,
|
||||
time_c - time_b,
|
||||
);
|
||||
|
||||
println!(
|
||||
"{:<20} {:>12} {:>12}",
|
||||
"Model", baseline_run.model, comparison_run.model,
|
||||
);
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,392 +0,0 @@
|
||||
use std::collections::HashSet;
|
||||
use std::io::{BufRead, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::BenchScore;
|
||||
|
||||
/// Metrics from a single task run: LLM usage, timing, tool calls.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct Trace {
|
||||
pub wall_time_ms: u64,
|
||||
pub llm_calls: u32,
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub estimated_cost_usd: f64,
|
||||
pub tool_calls: Vec<TraceToolCall>,
|
||||
pub turns: u32,
|
||||
pub hit_iteration_limit: bool,
|
||||
pub hit_timeout: bool,
|
||||
}
|
||||
|
||||
impl Trace {
|
||||
pub fn wall_time(&self) -> Duration {
|
||||
Duration::from_millis(self.wall_time_ms)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct TraceToolCall {
|
||||
pub name: String,
|
||||
pub duration_ms: u64,
|
||||
pub success: bool,
|
||||
}
|
||||
|
||||
/// Result of running a single benchmark task.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct TaskResult {
|
||||
pub task_id: String,
|
||||
pub suite_id: String,
|
||||
pub score: BenchScore,
|
||||
pub trace: Trace,
|
||||
pub response: String,
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub finished_at: DateTime<Utc>,
|
||||
pub config_label: String,
|
||||
#[serde(default)]
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Aggregate results for a full benchmark run.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct RunResult {
|
||||
pub run_id: Uuid,
|
||||
pub suite_id: String,
|
||||
pub config_label: String,
|
||||
pub model: String,
|
||||
pub pass_rate: f64,
|
||||
pub avg_score: f64,
|
||||
pub total_tasks: usize,
|
||||
pub completed_tasks: usize,
|
||||
pub total_cost_usd: f64,
|
||||
pub total_wall_time_ms: u64,
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub finished_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl RunResult {
|
||||
/// Build aggregate from individual task results.
|
||||
pub fn from_tasks(
|
||||
run_id: Uuid,
|
||||
suite_id: &str,
|
||||
config_label: &str,
|
||||
model: &str,
|
||||
total_tasks: usize,
|
||||
tasks: &[TaskResult],
|
||||
started_at: DateTime<Utc>,
|
||||
) -> Self {
|
||||
let pass_count = tasks.iter().filter(|t| t.score.value >= 1.0).count();
|
||||
let pass_rate = if tasks.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
pass_count as f64 / tasks.len() as f64
|
||||
};
|
||||
let avg_score = if tasks.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
tasks.iter().map(|t| t.score.value).sum::<f64>() / tasks.len() as f64
|
||||
};
|
||||
let total_cost: f64 = tasks.iter().map(|t| t.trace.estimated_cost_usd).sum();
|
||||
let total_wall: u64 = tasks.iter().map(|t| t.trace.wall_time_ms).sum();
|
||||
|
||||
Self {
|
||||
run_id,
|
||||
suite_id: suite_id.to_string(),
|
||||
config_label: config_label.to_string(),
|
||||
model: model.to_string(),
|
||||
pass_rate,
|
||||
avg_score,
|
||||
total_tasks,
|
||||
completed_tasks: tasks.len(),
|
||||
total_cost_usd: total_cost,
|
||||
total_wall_time_ms: total_wall,
|
||||
started_at,
|
||||
finished_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn total_wall_time(&self) -> Duration {
|
||||
Duration::from_millis(self.total_wall_time_ms)
|
||||
}
|
||||
}
|
||||
|
||||
/// Append a single task result as one JSON line to the JSONL file.
|
||||
pub fn append_task_result(path: &Path, result: &TaskResult) -> Result<(), BenchError> {
|
||||
let mut file = std::fs::OpenOptions::new()
|
||||
.create(true)
|
||||
.append(true)
|
||||
.open(path)?;
|
||||
let line = serde_json::to_string(result)?;
|
||||
writeln!(file, "{line}")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read all task results from a JSONL file.
|
||||
pub fn read_task_results(path: &Path) -> Result<Vec<TaskResult>, BenchError> {
|
||||
if !path.exists() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let file = std::fs::File::open(path)?;
|
||||
let reader = std::io::BufReader::new(file);
|
||||
let mut results = Vec::new();
|
||||
for line in reader.lines() {
|
||||
let line = line?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let result: TaskResult = serde_json::from_str(trimmed)?;
|
||||
results.push(result);
|
||||
}
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Write the aggregate run result as JSON.
|
||||
pub fn write_run_result(path: &Path, result: &RunResult) -> Result<(), BenchError> {
|
||||
let json = serde_json::to_string_pretty(result)?;
|
||||
std::fs::write(path, json)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Read the aggregate run result from JSON.
|
||||
pub fn read_run_result(path: &Path) -> Result<RunResult, BenchError> {
|
||||
let json = std::fs::read_to_string(path)?;
|
||||
let result: RunResult = serde_json::from_str(&json)?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Get the set of already-completed task IDs from a JSONL file (for resume).
|
||||
pub fn completed_task_ids(path: &Path) -> Result<HashSet<String>, BenchError> {
|
||||
let results = read_task_results(path)?;
|
||||
Ok(results.into_iter().map(|r| r.task_id).collect())
|
||||
}
|
||||
|
||||
/// Get the results directory for a specific run.
|
||||
pub fn run_dir(base: &Path, run_id: Uuid) -> PathBuf {
|
||||
base.join(run_id.to_string())
|
||||
}
|
||||
|
||||
/// Get the tasks JSONL path for a run.
|
||||
pub fn tasks_jsonl_path(base: &Path, run_id: Uuid) -> PathBuf {
|
||||
run_dir(base, run_id).join("tasks.jsonl")
|
||||
}
|
||||
|
||||
/// Get the run JSON path for a run.
|
||||
pub fn run_json_path(base: &Path, run_id: Uuid) -> PathBuf {
|
||||
run_dir(base, run_id).join("run.json")
|
||||
}
|
||||
|
||||
/// Find the latest run directory (by modification time).
|
||||
pub fn find_latest_run(base: &Path) -> Result<Option<Uuid>, BenchError> {
|
||||
if !base.exists() {
|
||||
return Ok(None);
|
||||
}
|
||||
let mut entries: Vec<_> = std::fs::read_dir(base)?
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| e.file_type().map(|ft| ft.is_dir()).unwrap_or(false))
|
||||
.filter_map(|e| {
|
||||
let name = e.file_name().to_string_lossy().to_string();
|
||||
let uuid = Uuid::parse_str(&name).ok()?;
|
||||
let modified = e.metadata().ok()?.modified().ok()?;
|
||||
Some((uuid, modified))
|
||||
})
|
||||
.collect();
|
||||
entries.sort_by(|a, b| b.1.cmp(&a.1));
|
||||
Ok(entries.first().map(|(uuid, _)| *uuid))
|
||||
}
|
||||
|
||||
/// Print a summary table of task results.
|
||||
pub fn print_results_table(tasks: &[TaskResult], run: &RunResult) {
|
||||
println!();
|
||||
println!(
|
||||
"Run: {} | Suite: {} | Config: {} | Model: {}",
|
||||
run.run_id, run.suite_id, run.config_label, run.model
|
||||
);
|
||||
println!(
|
||||
"Pass rate: {:.1}% | Avg score: {:.3} | Tasks: {}/{} | Cost: ${:.4} | Time: {:.1}s",
|
||||
run.pass_rate * 100.0,
|
||||
run.avg_score,
|
||||
run.completed_tasks,
|
||||
run.total_tasks,
|
||||
run.total_cost_usd,
|
||||
run.total_wall_time_ms as f64 / 1000.0,
|
||||
);
|
||||
println!();
|
||||
|
||||
// Header
|
||||
println!(
|
||||
"{:<30} {:>6} {:>7} {:>8} {:>10} {:>6} {:>8}",
|
||||
"Task ID", "Score", "Label", "Tokens", "Cost", "Turns", "Time"
|
||||
);
|
||||
println!("{}", "-".repeat(80));
|
||||
|
||||
for task in tasks {
|
||||
let total_tokens = task.trace.input_tokens + task.trace.output_tokens;
|
||||
let task_id_display = if task.task_id.len() > 28 {
|
||||
format!("{}...", &task.task_id[..25])
|
||||
} else {
|
||||
task.task_id.clone()
|
||||
};
|
||||
println!(
|
||||
"{:<30} {:>6.3} {:>7} {:>8} {:>10.4} {:>6} {:>7.1}s",
|
||||
task_id_display,
|
||||
task.score.value,
|
||||
task.score.label,
|
||||
total_tokens,
|
||||
task.trace.estimated_cost_usd,
|
||||
task.trace.turns,
|
||||
task.trace.wall_time_ms as f64 / 1000.0,
|
||||
);
|
||||
}
|
||||
println!();
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_run_result_from_tasks() {
|
||||
let tasks = vec![
|
||||
TaskResult {
|
||||
task_id: "t1".to_string(),
|
||||
suite_id: "custom".to_string(),
|
||||
score: BenchScore {
|
||||
value: 1.0,
|
||||
label: "pass".to_string(),
|
||||
details: None,
|
||||
},
|
||||
trace: Trace {
|
||||
wall_time_ms: 1000,
|
||||
llm_calls: 2,
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
estimated_cost_usd: 0.01,
|
||||
tool_calls: vec![],
|
||||
turns: 1,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: "answer".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "default".to_string(),
|
||||
error: None,
|
||||
},
|
||||
TaskResult {
|
||||
task_id: "t2".to_string(),
|
||||
suite_id: "custom".to_string(),
|
||||
score: BenchScore {
|
||||
value: 0.0,
|
||||
label: "fail".to_string(),
|
||||
details: Some("wrong".to_string()),
|
||||
},
|
||||
trace: Trace {
|
||||
wall_time_ms: 2000,
|
||||
llm_calls: 3,
|
||||
input_tokens: 200,
|
||||
output_tokens: 100,
|
||||
estimated_cost_usd: 0.02,
|
||||
tool_calls: vec![],
|
||||
turns: 2,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: "wrong answer".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "default".to_string(),
|
||||
error: None,
|
||||
},
|
||||
];
|
||||
|
||||
let run = RunResult::from_tasks(
|
||||
Uuid::new_v4(),
|
||||
"custom",
|
||||
"default",
|
||||
"test-model",
|
||||
2,
|
||||
&tasks,
|
||||
Utc::now(),
|
||||
);
|
||||
|
||||
assert_eq!(run.pass_rate, 0.5);
|
||||
assert_eq!(run.avg_score, 0.5);
|
||||
assert_eq!(run.total_tasks, 2);
|
||||
assert_eq!(run.completed_tasks, 2);
|
||||
assert!((run.total_cost_usd - 0.03).abs() < f64::EPSILON);
|
||||
assert_eq!(run.total_wall_time_ms, 3000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_jsonl_roundtrip() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
|
||||
let result = TaskResult {
|
||||
task_id: "round-trip-test".to_string(),
|
||||
suite_id: "custom".to_string(),
|
||||
score: BenchScore::pass(),
|
||||
trace: Trace {
|
||||
wall_time_ms: 500,
|
||||
llm_calls: 1,
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
estimated_cost_usd: 0.001,
|
||||
tool_calls: vec![],
|
||||
turns: 1,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: "hello".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "test".to_string(),
|
||||
error: None,
|
||||
};
|
||||
|
||||
append_task_result(&path, &result).expect("append");
|
||||
append_task_result(&path, &result).expect("append");
|
||||
|
||||
let loaded = read_task_results(&path).expect("read");
|
||||
assert_eq!(loaded.len(), 2);
|
||||
assert_eq!(loaded[0].task_id, "round-trip-test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_completed_task_ids() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
|
||||
let result = TaskResult {
|
||||
task_id: "unique-id-1".to_string(),
|
||||
suite_id: "custom".to_string(),
|
||||
score: BenchScore::pass(),
|
||||
trace: Trace {
|
||||
wall_time_ms: 100,
|
||||
llm_calls: 1,
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
estimated_cost_usd: 0.0,
|
||||
tool_calls: vec![],
|
||||
turns: 1,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: "x".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "test".to_string(),
|
||||
error: None,
|
||||
};
|
||||
append_task_result(&path, &result).expect("append");
|
||||
|
||||
let ids = completed_task_ids(&path).expect("ids");
|
||||
assert!(ids.contains("unique-id-1"));
|
||||
assert!(!ids.contains("unique-id-2"));
|
||||
}
|
||||
}
|
||||
@@ -1,460 +0,0 @@
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use chrono::Utc;
|
||||
use tokio::sync::Mutex;
|
||||
use uuid::Uuid;
|
||||
|
||||
use ironclaw::agent::{Agent, AgentDeps};
|
||||
use ironclaw::channels::{ChannelManager, IncomingMessage};
|
||||
use ironclaw::config::AgentConfig;
|
||||
use ironclaw::llm::LlmProvider;
|
||||
use ironclaw::safety::SafetyLayer;
|
||||
use ironclaw::tools::ToolRegistry;
|
||||
|
||||
use crate::channel::BenchChannel;
|
||||
use crate::config::{BenchConfig, MatrixEntry};
|
||||
use crate::error::BenchError;
|
||||
use crate::instrumented_llm::InstrumentedLlm;
|
||||
use crate::results::{
|
||||
RunResult, TaskResult, Trace, append_task_result, completed_task_ids, run_dir, run_json_path,
|
||||
tasks_jsonl_path, write_run_result,
|
||||
};
|
||||
use crate::suite::{BenchSuite, BenchTask, ConversationTurn, TaskSubmission, TurnRole};
|
||||
|
||||
/// Parameters for running a single task in isolation.
|
||||
struct TaskRunParams<'a> {
|
||||
task: &'a BenchTask,
|
||||
suite_id: &'a str,
|
||||
config_label: &'a str,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
timeout: std::time::Duration,
|
||||
additional_tools: &'a [Arc<dyn ironclaw::tools::Tool>],
|
||||
}
|
||||
|
||||
/// Orchestrates benchmark execution: loads tasks, runs agent per task,
|
||||
/// scores results, writes JSONL output.
|
||||
pub struct BenchRunner {
|
||||
suite: Box<dyn BenchSuite>,
|
||||
config: BenchConfig,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
}
|
||||
|
||||
impl BenchRunner {
|
||||
pub fn new(
|
||||
suite: Box<dyn BenchSuite>,
|
||||
config: BenchConfig,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
) -> Self {
|
||||
Self {
|
||||
suite,
|
||||
config,
|
||||
llm,
|
||||
safety,
|
||||
}
|
||||
}
|
||||
|
||||
/// Run the benchmark for one matrix entry.
|
||||
///
|
||||
/// Returns the run_id for result retrieval.
|
||||
pub async fn run(
|
||||
&self,
|
||||
matrix: &MatrixEntry,
|
||||
sample: Option<usize>,
|
||||
task_filter: Option<&[String]>,
|
||||
tag_filter: Option<&[String]>,
|
||||
resume_run_id: Option<Uuid>,
|
||||
) -> Result<Uuid, BenchError> {
|
||||
let run_id = resume_run_id.unwrap_or_else(Uuid::new_v4);
|
||||
let results_base = &self.config.results_dir;
|
||||
let dir = run_dir(results_base, run_id);
|
||||
std::fs::create_dir_all(&dir)?;
|
||||
|
||||
let jsonl_path = tasks_jsonl_path(results_base, run_id);
|
||||
let json_path = run_json_path(results_base, run_id);
|
||||
|
||||
// Load completed task IDs for resume support
|
||||
let completed: HashSet<String> = if resume_run_id.is_some() {
|
||||
completed_task_ids(&jsonl_path)?
|
||||
} else {
|
||||
HashSet::new()
|
||||
};
|
||||
|
||||
if !completed.is_empty() {
|
||||
tracing::info!(
|
||||
"Resuming run {}: {} tasks already completed",
|
||||
run_id,
|
||||
completed.len()
|
||||
);
|
||||
}
|
||||
|
||||
// Load and filter tasks
|
||||
let mut tasks = self.suite.load_tasks().await?;
|
||||
|
||||
if let Some(ids) = task_filter {
|
||||
let id_set: HashSet<&str> = ids.iter().map(|s| s.as_str()).collect();
|
||||
tasks.retain(|t| id_set.contains(t.id.as_str()));
|
||||
}
|
||||
|
||||
if let Some(tags) = tag_filter {
|
||||
let tag_set: HashSet<&str> = tags.iter().map(|s| s.as_str()).collect();
|
||||
tasks.retain(|t| t.tags.iter().any(|tag| tag_set.contains(tag.as_str())));
|
||||
}
|
||||
|
||||
// Filter out already-completed tasks
|
||||
tasks.retain(|t| !completed.contains(&t.id));
|
||||
|
||||
// Sample if requested
|
||||
if let Some(n) = sample {
|
||||
tasks.truncate(n);
|
||||
}
|
||||
|
||||
let total_tasks = tasks.len() + completed.len();
|
||||
tracing::info!(
|
||||
"Running {} tasks for suite '{}' (config: {}, run: {})",
|
||||
tasks.len(),
|
||||
self.suite.id(),
|
||||
matrix.label,
|
||||
run_id
|
||||
);
|
||||
|
||||
let started_at = Utc::now();
|
||||
let all_results: Arc<Mutex<Vec<TaskResult>>> =
|
||||
Arc::new(Mutex::new(Vec::with_capacity(tasks.len())));
|
||||
|
||||
if self.config.parallelism <= 1 {
|
||||
// Sequential execution
|
||||
let additional_tools = self.suite.additional_tools();
|
||||
for (i, task) in tasks.iter().enumerate() {
|
||||
tracing::info!(
|
||||
"[{}/{}] Running task: {}",
|
||||
i + 1 + completed.len(),
|
||||
total_tasks,
|
||||
task.id
|
||||
);
|
||||
let params = TaskRunParams {
|
||||
task,
|
||||
suite_id: self.suite.id(),
|
||||
config_label: &matrix.label,
|
||||
llm: Arc::clone(&self.llm),
|
||||
safety: Arc::clone(&self.safety),
|
||||
timeout: task.timeout.unwrap_or(self.config.task_timeout),
|
||||
additional_tools: &additional_tools,
|
||||
};
|
||||
let result = run_task_isolated(params).await;
|
||||
append_task_result(&jsonl_path, &result)?;
|
||||
all_results.lock().await.push(result);
|
||||
}
|
||||
} else {
|
||||
// Parallel execution with bounded concurrency
|
||||
let semaphore = Arc::new(tokio::sync::Semaphore::new(self.config.parallelism));
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for (i, task) in tasks.into_iter().enumerate() {
|
||||
let sem = Arc::clone(&semaphore);
|
||||
let suite_id = self.suite.id().to_string();
|
||||
let config_label = matrix.label.clone();
|
||||
let llm = Arc::clone(&self.llm);
|
||||
let safety = Arc::clone(&self.safety);
|
||||
let timeout = task.timeout.unwrap_or(self.config.task_timeout);
|
||||
let results_ref = Arc::clone(&all_results);
|
||||
let jsonl = jsonl_path.clone();
|
||||
let completed_count = completed.len();
|
||||
let total = total_tasks;
|
||||
let additional_tools = self.suite.additional_tools();
|
||||
|
||||
handles.push(tokio::spawn(async move {
|
||||
let _permit = sem.acquire().await.expect("semaphore closed");
|
||||
tracing::info!(
|
||||
"[{}/{}] Running task: {}",
|
||||
i + 1 + completed_count,
|
||||
total,
|
||||
task.id
|
||||
);
|
||||
let params = TaskRunParams {
|
||||
task: &task,
|
||||
suite_id: &suite_id,
|
||||
config_label: &config_label,
|
||||
llm,
|
||||
safety,
|
||||
timeout,
|
||||
additional_tools: &additional_tools,
|
||||
};
|
||||
let result = run_task_isolated(params).await;
|
||||
if let Err(e) = append_task_result(&jsonl, &result) {
|
||||
tracing::error!("Failed to write result for {}: {}", task.id, e);
|
||||
}
|
||||
results_ref.lock().await.push(result);
|
||||
}));
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
if let Err(e) = handle.await {
|
||||
tracing::error!("Task panicked: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Score all results
|
||||
let results = all_results.lock().await;
|
||||
let mut scored: Vec<TaskResult> = Vec::with_capacity(results.len());
|
||||
for result in results.iter() {
|
||||
let task_opt = self
|
||||
.suite
|
||||
.load_tasks()
|
||||
.await?
|
||||
.into_iter()
|
||||
.find(|t| t.id == result.task_id);
|
||||
|
||||
if let Some(task) = task_opt {
|
||||
let submission = TaskSubmission {
|
||||
response: result.response.clone(),
|
||||
conversation: vec![],
|
||||
tool_calls: result
|
||||
.trace
|
||||
.tool_calls
|
||||
.iter()
|
||||
.map(|tc| tc.name.clone())
|
||||
.collect(),
|
||||
};
|
||||
match self.suite.score(&task, &submission).await {
|
||||
Ok(score) => {
|
||||
let mut scored_result = result.clone();
|
||||
scored_result.score = score;
|
||||
scored.push(scored_result);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Scoring failed for {}: {}", result.task_id, e);
|
||||
scored.push(result.clone());
|
||||
}
|
||||
}
|
||||
} else {
|
||||
scored.push(result.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// Combine with any previously completed results for the aggregate
|
||||
let mut all_for_aggregate = crate::results::read_task_results(&jsonl_path)?;
|
||||
// De-duplicate (prefer the newer scored versions)
|
||||
let scored_ids: HashSet<String> = scored.iter().map(|r| r.task_id.clone()).collect();
|
||||
all_for_aggregate.retain(|r| !scored_ids.contains(&r.task_id));
|
||||
all_for_aggregate.extend(scored);
|
||||
|
||||
let model_name = matrix.model.as_deref().unwrap_or(self.llm.model_name());
|
||||
|
||||
let run_result = RunResult::from_tasks(
|
||||
run_id,
|
||||
self.suite.id(),
|
||||
&matrix.label,
|
||||
model_name,
|
||||
total_tasks,
|
||||
&all_for_aggregate,
|
||||
started_at,
|
||||
);
|
||||
|
||||
write_run_result(&json_path, &run_result)?;
|
||||
|
||||
tracing::info!(
|
||||
"Run {} complete: {:.1}% pass rate, {:.3} avg score, ${:.4} cost",
|
||||
run_id,
|
||||
run_result.pass_rate * 100.0,
|
||||
run_result.avg_score,
|
||||
run_result.total_cost_usd,
|
||||
);
|
||||
|
||||
Ok(run_id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Run a single benchmark task in complete isolation.
|
||||
///
|
||||
/// Creates a fresh Agent + BenchChannel + InstrumentedLlm for the task,
|
||||
/// injects the prompt, waits for the response, and returns the result.
|
||||
async fn run_task_isolated(params: TaskRunParams<'_>) -> TaskResult {
|
||||
let TaskRunParams {
|
||||
task,
|
||||
suite_id,
|
||||
config_label,
|
||||
llm,
|
||||
safety,
|
||||
timeout,
|
||||
additional_tools,
|
||||
} = params;
|
||||
|
||||
let started_at = Utc::now();
|
||||
let start = Instant::now();
|
||||
|
||||
// Wrap LLM with instrumentation
|
||||
let instrumented = Arc::new(InstrumentedLlm::new(llm));
|
||||
|
||||
// Create bench channel
|
||||
let (bench_channel, msg_tx) = BenchChannel::new();
|
||||
let capture = bench_channel.capture();
|
||||
|
||||
// Build tool registry
|
||||
let tools = Arc::new(ToolRegistry::new());
|
||||
tools.register_builtin_tools();
|
||||
|
||||
// Register additional suite-specific tools
|
||||
for tool in additional_tools {
|
||||
tools.register(Arc::clone(tool)).await;
|
||||
}
|
||||
|
||||
// Build agent config (minimal, headless)
|
||||
let agent_config = AgentConfig {
|
||||
name: format!("bench-{}", task.id),
|
||||
max_parallel_jobs: 1,
|
||||
job_timeout: timeout,
|
||||
stuck_threshold: timeout,
|
||||
repair_check_interval: timeout + std::time::Duration::from_secs(999),
|
||||
max_repair_attempts: 0,
|
||||
use_planning: false,
|
||||
session_idle_timeout: timeout,
|
||||
};
|
||||
|
||||
let deps = AgentDeps {
|
||||
store: None,
|
||||
llm: instrumented.clone() as Arc<dyn LlmProvider>,
|
||||
safety,
|
||||
tools,
|
||||
workspace: None,
|
||||
extension_manager: None,
|
||||
};
|
||||
|
||||
let mut channels = ChannelManager::new();
|
||||
channels.add(Box::new(bench_channel));
|
||||
|
||||
let agent = Agent::new(agent_config, deps, channels, None, None, None);
|
||||
|
||||
// Build the full prompt with context
|
||||
let full_prompt = if let Some(ref ctx) = task.context {
|
||||
format!("{}\n\nContext:\n{}", task.prompt, ctx)
|
||||
} else {
|
||||
task.prompt.clone()
|
||||
};
|
||||
|
||||
// Inject the task prompt
|
||||
let incoming = IncomingMessage::new("bench", "bench-user", &full_prompt);
|
||||
if msg_tx.send(incoming).await.is_err() {
|
||||
return make_error_result(
|
||||
task,
|
||||
suite_id,
|
||||
config_label,
|
||||
started_at,
|
||||
"failed to send prompt",
|
||||
);
|
||||
}
|
||||
|
||||
// Record prompt in conversation
|
||||
{
|
||||
let mut cap = capture.lock().await;
|
||||
cap.conversation.push(ConversationTurn {
|
||||
role: TurnRole::User,
|
||||
content: full_prompt,
|
||||
});
|
||||
}
|
||||
|
||||
// Run agent with timeout.
|
||||
// After the first response, send /quit to end the session.
|
||||
let quit_tx = msg_tx.clone();
|
||||
let capture_for_quit = Arc::clone(&capture);
|
||||
let quit_handle = tokio::spawn(async move {
|
||||
// Poll for first response
|
||||
loop {
|
||||
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
|
||||
let cap = capture_for_quit.lock().await;
|
||||
if !cap.responses.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
// Give a small grace period for any final status events
|
||||
tokio::time::sleep(std::time::Duration::from_millis(200)).await;
|
||||
let quit = IncomingMessage::new("bench", "bench-user", "/quit");
|
||||
let _ = quit_tx.send(quit).await;
|
||||
});
|
||||
|
||||
let agent_result = tokio::time::timeout(timeout, agent.run()).await;
|
||||
|
||||
quit_handle.abort();
|
||||
|
||||
let wall_time = start.elapsed();
|
||||
let hit_timeout = agent_result.is_err();
|
||||
|
||||
if let Ok(Err(e)) = &agent_result {
|
||||
tracing::warn!("Agent error for task {}: {}", task.id, e);
|
||||
}
|
||||
|
||||
// Extract results from capture
|
||||
let cap = capture.lock().await;
|
||||
let response = cap.responses.last().cloned().unwrap_or_default();
|
||||
|
||||
let trace = Trace {
|
||||
wall_time_ms: wall_time.as_millis() as u64,
|
||||
llm_calls: instrumented.call_count(),
|
||||
input_tokens: instrumented.total_input_tokens(),
|
||||
output_tokens: instrumented.total_output_tokens(),
|
||||
estimated_cost_usd: instrumented.estimated_cost(),
|
||||
tool_calls: cap.tool_calls.clone(),
|
||||
turns: cap.responses.len() as u32,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout,
|
||||
};
|
||||
|
||||
let error = if hit_timeout {
|
||||
Some(format!("timeout after {}s", timeout.as_secs()))
|
||||
} else if let Ok(Err(e)) = &agent_result {
|
||||
Some(e.to_string())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
TaskResult {
|
||||
task_id: task.id.clone(),
|
||||
suite_id: suite_id.to_string(),
|
||||
score: crate::suite::BenchScore {
|
||||
value: 0.0,
|
||||
label: "pending".to_string(),
|
||||
details: None,
|
||||
},
|
||||
trace,
|
||||
response,
|
||||
started_at,
|
||||
finished_at: Utc::now(),
|
||||
config_label: config_label.to_string(),
|
||||
error,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_error_result(
|
||||
task: &BenchTask,
|
||||
suite_id: &str,
|
||||
config_label: &str,
|
||||
started_at: chrono::DateTime<Utc>,
|
||||
reason: &str,
|
||||
) -> TaskResult {
|
||||
TaskResult {
|
||||
task_id: task.id.clone(),
|
||||
suite_id: suite_id.to_string(),
|
||||
score: crate::suite::BenchScore::fail(reason),
|
||||
trace: Trace {
|
||||
wall_time_ms: 0,
|
||||
llm_calls: 0,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
estimated_cost_usd: 0.0,
|
||||
tool_calls: vec![],
|
||||
turns: 0,
|
||||
hit_iteration_limit: false,
|
||||
hit_timeout: false,
|
||||
},
|
||||
response: String::new(),
|
||||
started_at,
|
||||
finished_at: Utc::now(),
|
||||
config_label: config_label.to_string(),
|
||||
error: Some(reason.to_string()),
|
||||
}
|
||||
}
|
||||
@@ -1,113 +0,0 @@
|
||||
use regex::Regex;
|
||||
|
||||
use crate::suite::BenchScore;
|
||||
|
||||
/// Normalize an answer string for comparison: lowercase, trim whitespace,
|
||||
/// strip trailing punctuation, collapse internal whitespace.
|
||||
pub fn normalize_answer(s: &str) -> String {
|
||||
let trimmed = s.trim().to_lowercase();
|
||||
let collapsed: String = trimmed.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
collapsed.trim_end_matches(['.', ',', ';', '!']).to_string()
|
||||
}
|
||||
|
||||
/// Exact match after normalization.
|
||||
pub fn exact_match(expected: &str, actual: &str) -> BenchScore {
|
||||
let norm_expected = normalize_answer(expected);
|
||||
let norm_actual = normalize_answer(actual);
|
||||
if norm_expected == norm_actual {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!(
|
||||
"expected \"{norm_expected}\", got \"{norm_actual}\""
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the actual answer contains the expected substring (normalized).
|
||||
pub fn contains_match(expected_substring: &str, actual: &str) -> BenchScore {
|
||||
let norm_expected = normalize_answer(expected_substring);
|
||||
let norm_actual = normalize_answer(actual);
|
||||
if norm_actual.contains(&norm_expected) {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!("response does not contain \"{norm_expected}\""))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the actual answer matches a regex pattern.
|
||||
pub fn regex_match(pattern: &str, actual: &str) -> BenchScore {
|
||||
match Regex::new(pattern) {
|
||||
Ok(re) => {
|
||||
if re.is_match(actual) {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!("response does not match pattern /{pattern}/"))
|
||||
}
|
||||
}
|
||||
Err(e) => BenchScore::fail(format!("invalid regex pattern: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_normalize_answer() {
|
||||
assert_eq!(normalize_answer(" Hello World. "), "hello world");
|
||||
assert_eq!(normalize_answer("Yes!"), "yes");
|
||||
assert_eq!(normalize_answer("42"), "42");
|
||||
assert_eq!(normalize_answer(" "), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exact_match_pass() {
|
||||
let score = exact_match("Hello World", " hello world. ");
|
||||
assert_eq!(score.value, 1.0);
|
||||
assert_eq!(score.label, "pass");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exact_match_fail() {
|
||||
let score = exact_match("hello", "world");
|
||||
assert_eq!(score.value, 0.0);
|
||||
assert_eq!(score.label, "fail");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_contains_match_pass() {
|
||||
let score = contains_match("world", "Hello World!");
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_contains_match_fail() {
|
||||
let score = contains_match("xyz", "Hello World!");
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_pass() {
|
||||
let score = regex_match(r"\d{4}", "The year is 2024.");
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_fail() {
|
||||
let score = regex_match(r"\d{4}", "No numbers here.");
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_invalid_pattern() {
|
||||
let score = regex_match(r"[invalid", "anything");
|
||||
assert_eq!(score.value, 0.0);
|
||||
assert!(
|
||||
score
|
||||
.details
|
||||
.as_deref()
|
||||
.unwrap_or("")
|
||||
.contains("invalid regex")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,156 +0,0 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::error::BenchError;
|
||||
|
||||
/// A single task in a benchmark suite.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BenchTask {
|
||||
pub id: String,
|
||||
pub prompt: String,
|
||||
#[serde(default)]
|
||||
pub context: Option<String>,
|
||||
#[serde(default)]
|
||||
pub resources: Vec<TaskResource>,
|
||||
#[serde(default)]
|
||||
pub tags: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub expected_turns: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub timeout: Option<Duration>,
|
||||
#[serde(default)]
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
/// A resource attached to a benchmark task (file, URL, etc.).
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct TaskResource {
|
||||
pub name: String,
|
||||
pub path: String,
|
||||
#[serde(default)]
|
||||
pub resource_type: ResourceType,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ResourceType {
|
||||
#[default]
|
||||
File,
|
||||
Url,
|
||||
Directory,
|
||||
}
|
||||
|
||||
/// What the agent produced for scoring.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TaskSubmission {
|
||||
pub response: String,
|
||||
pub conversation: Vec<ConversationTurn>,
|
||||
pub tool_calls: Vec<String>,
|
||||
}
|
||||
|
||||
/// A single turn in a multi-turn conversation.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct ConversationTurn {
|
||||
pub role: TurnRole,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TurnRole {
|
||||
User,
|
||||
Assistant,
|
||||
System,
|
||||
}
|
||||
|
||||
/// Score for a single task.
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct BenchScore {
|
||||
/// 0.0 to 1.0 (1.0 = perfect).
|
||||
pub value: f64,
|
||||
/// "pass" / "fail" / "partial".
|
||||
pub label: String,
|
||||
#[serde(default)]
|
||||
pub details: Option<String>,
|
||||
}
|
||||
|
||||
impl BenchScore {
|
||||
pub fn pass() -> Self {
|
||||
Self {
|
||||
value: 1.0,
|
||||
label: "pass".to_string(),
|
||||
details: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn fail(details: impl Into<String>) -> Self {
|
||||
Self {
|
||||
value: 0.0,
|
||||
label: "fail".to_string(),
|
||||
details: Some(details.into()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn partial(value: f64, details: impl Into<String>) -> Self {
|
||||
Self {
|
||||
value: value.clamp(0.0, 1.0),
|
||||
label: "partial".to_string(),
|
||||
details: Some(details.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for benchmark suite adapters.
|
||||
///
|
||||
/// Each suite (GAIA, Tau-bench, custom, etc.) implements this trait
|
||||
/// to provide task loading, scoring, and optional lifecycle hooks.
|
||||
#[async_trait]
|
||||
pub trait BenchSuite: Send + Sync {
|
||||
/// Human-readable name (e.g., "GAIA Validation").
|
||||
fn name(&self) -> &str;
|
||||
|
||||
/// Machine ID (e.g., "gaia").
|
||||
fn id(&self) -> &str;
|
||||
|
||||
/// Load all tasks from the suite's data source.
|
||||
async fn load_tasks(&self) -> Result<Vec<BenchTask>, BenchError>;
|
||||
|
||||
/// Score the agent's submission against the expected answer.
|
||||
async fn score(
|
||||
&self,
|
||||
task: &BenchTask,
|
||||
submission: &TaskSubmission,
|
||||
) -> Result<BenchScore, BenchError>;
|
||||
|
||||
/// Optional: set up environment before running a task (clone repo, init DB, etc.).
|
||||
async fn setup_task(&self, _task: &BenchTask) -> Result<(), BenchError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Optional: tear down environment after a task completes.
|
||||
async fn teardown_task(&self, _task: &BenchTask) -> Result<(), BenchError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Optional: additional tools to register for this suite's tasks.
|
||||
fn additional_tools(&self) -> Vec<Arc<dyn ironclaw::tools::Tool>> {
|
||||
vec![]
|
||||
}
|
||||
|
||||
/// Optional: restrict which tools the agent can use (allowlist).
|
||||
fn tool_whitelist(&self) -> Option<Vec<String>> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Multi-turn: generate next simulated user message based on conversation so far.
|
||||
/// Return `None` to end the conversation.
|
||||
async fn next_user_message(
|
||||
&self,
|
||||
_task: &BenchTask,
|
||||
_conversation: &[ConversationTurn],
|
||||
) -> Result<Option<String>, BenchError> {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
@@ -21,6 +21,7 @@ use crate::context::JobContext;
|
||||
use crate::error::Error;
|
||||
use crate::extensions::ExtensionManager;
|
||||
use crate::history::Store;
|
||||
use crate::keys::KeyManager;
|
||||
use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::tools::ToolRegistry;
|
||||
@@ -64,6 +65,7 @@ pub struct AgentDeps {
|
||||
pub tools: Arc<ToolRegistry>,
|
||||
pub workspace: Option<Arc<Workspace>>,
|
||||
pub extension_manager: Option<Arc<ExtensionManager>>,
|
||||
pub key_manager: Option<Arc<KeyManager>>,
|
||||
}
|
||||
|
||||
/// The main agent that coordinates all components.
|
||||
|
||||
@@ -314,11 +314,11 @@ impl WsServerMessage {
|
||||
SseEvent::Thinking { .. } => "thinking",
|
||||
SseEvent::ToolStarted { .. } => "tool_started",
|
||||
SseEvent::ToolCompleted { .. } => "tool_completed",
|
||||
SseEvent::ToolResult { .. } => "tool_result",
|
||||
SseEvent::StreamChunk { .. } => "stream_chunk",
|
||||
SseEvent::Status { .. } => "status",
|
||||
SseEvent::ApprovalNeeded { .. } => "approval_needed",
|
||||
SseEvent::Error { .. } => "error",
|
||||
SseEvent::ToolResult { .. } => "tool_result",
|
||||
SseEvent::Heartbeat => "heartbeat",
|
||||
};
|
||||
let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
|
||||
|
||||
+745
@@ -0,0 +1,745 @@
|
||||
//! NEAR key management CLI commands.
|
||||
|
||||
use std::io::Write;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::Subcommand;
|
||||
use tokio::fs;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::history::Store;
|
||||
use crate::keys::KeyManager;
|
||||
use crate::keys::policy::{ChainSigRule, FunctionCallRule, PolicyConfig, SignatureDomain};
|
||||
use crate::keys::types::{
|
||||
AccessKeyPermission, NearAccountId, NearNetwork, format_yocto, parse_near_amount,
|
||||
};
|
||||
use crate::secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore};
|
||||
|
||||
/// Default policy config path.
|
||||
fn default_policy_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.map(|h| h.join(".ironclaw").join("key_policy.json"))
|
||||
.unwrap_or_else(|| PathBuf::from(".ironclaw/key_policy.json"))
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
pub enum KeyCommand {
|
||||
/// Generate a new ed25519 keypair
|
||||
Generate {
|
||||
/// Label for the key (used to reference it later)
|
||||
label: String,
|
||||
|
||||
/// NEAR account ID this key belongs to
|
||||
#[arg(long)]
|
||||
account: String,
|
||||
|
||||
/// Permission level: "full-access" or "function-call"
|
||||
#[arg(long, default_value = "function-call")]
|
||||
permission: String,
|
||||
|
||||
/// Contract to scope function-call keys to
|
||||
#[arg(long)]
|
||||
receiver: Option<String>,
|
||||
|
||||
/// Comma-separated method names (empty = all methods on contract)
|
||||
#[arg(long)]
|
||||
methods: Option<String>,
|
||||
|
||||
/// Allowance in NEAR (e.g., "1.5")
|
||||
#[arg(long)]
|
||||
allowance: Option<String>,
|
||||
|
||||
/// Network: mainnet, testnet, or RPC URL
|
||||
#[arg(long, default_value = "testnet")]
|
||||
network: String,
|
||||
},
|
||||
|
||||
/// Import an existing secret key
|
||||
Import {
|
||||
/// Label for the key
|
||||
label: String,
|
||||
|
||||
/// NEAR account ID
|
||||
#[arg(long)]
|
||||
account: String,
|
||||
|
||||
/// Permission level
|
||||
#[arg(long, default_value = "function-call")]
|
||||
permission: String,
|
||||
|
||||
/// Contract to scope function-call keys to
|
||||
#[arg(long)]
|
||||
receiver: Option<String>,
|
||||
|
||||
/// Comma-separated method names
|
||||
#[arg(long)]
|
||||
methods: Option<String>,
|
||||
|
||||
/// Allowance in NEAR
|
||||
#[arg(long)]
|
||||
allowance: Option<String>,
|
||||
|
||||
/// Network
|
||||
#[arg(long, default_value = "testnet")]
|
||||
network: String,
|
||||
},
|
||||
|
||||
/// List all stored keys
|
||||
List {
|
||||
/// Show verbose details
|
||||
#[arg(short, long)]
|
||||
verbose: bool,
|
||||
},
|
||||
|
||||
/// Show information about a key
|
||||
Info {
|
||||
/// Key label
|
||||
label: String,
|
||||
},
|
||||
|
||||
/// Remove a key
|
||||
Remove {
|
||||
/// Key label
|
||||
label: String,
|
||||
},
|
||||
|
||||
/// Export public key (NEVER exports private key)
|
||||
Export {
|
||||
/// Key label
|
||||
label: String,
|
||||
},
|
||||
|
||||
/// Manage transaction approval policy
|
||||
#[command(subcommand)]
|
||||
Policy(PolicyCommand),
|
||||
|
||||
/// Create encrypted backup of all keys
|
||||
Backup {
|
||||
/// Output file path
|
||||
#[arg(long)]
|
||||
output: PathBuf,
|
||||
|
||||
/// List keys in a backup without restoring (still needs passphrase)
|
||||
#[arg(long)]
|
||||
list: bool,
|
||||
},
|
||||
|
||||
/// Restore keys from encrypted backup
|
||||
Restore {
|
||||
/// Backup file path
|
||||
path: PathBuf,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug, Clone)]
|
||||
pub enum PolicyCommand {
|
||||
/// Show current policy configuration
|
||||
Show,
|
||||
|
||||
/// Set auto-approve transfer limit
|
||||
SetTransferLimit {
|
||||
/// Max NEAR amount for auto-approved transfers (e.g., "1.5")
|
||||
amount: String,
|
||||
},
|
||||
|
||||
/// Whitelist an account for transfers
|
||||
WhitelistAccount {
|
||||
/// Account ID to whitelist
|
||||
account: String,
|
||||
|
||||
/// Max transfer amount in NEAR
|
||||
#[arg(long)]
|
||||
max_transfer: Option<String>,
|
||||
},
|
||||
|
||||
/// Whitelist a validator for staking
|
||||
WhitelistValidator {
|
||||
/// Validator account ID
|
||||
validator: String,
|
||||
|
||||
/// Max stake amount in NEAR
|
||||
#[arg(long)]
|
||||
max_stake: Option<String>,
|
||||
},
|
||||
|
||||
/// Add a function call rule for a contract
|
||||
AddContractRule {
|
||||
/// Contract account ID
|
||||
contract: String,
|
||||
|
||||
/// Comma-separated method names (empty = all)
|
||||
#[arg(long)]
|
||||
methods: Option<String>,
|
||||
|
||||
/// Max deposit in NEAR
|
||||
#[arg(long, default_value = "0")]
|
||||
max_deposit: String,
|
||||
|
||||
/// Auto-approve matching calls
|
||||
#[arg(long)]
|
||||
auto_approve: bool,
|
||||
},
|
||||
|
||||
/// Add a chain signature rule
|
||||
AddChainSigRule {
|
||||
/// Derivation path pattern (supports * glob)
|
||||
path_pattern: String,
|
||||
|
||||
/// Signature domain: secp256k1 or ed25519
|
||||
#[arg(long, default_value = "secp256k1")]
|
||||
domain: String,
|
||||
|
||||
/// Max payload size in bytes
|
||||
#[arg(long, default_value = "4096")]
|
||||
max_payload: usize,
|
||||
|
||||
/// Auto-approve matching requests
|
||||
#[arg(long)]
|
||||
auto_approve: bool,
|
||||
},
|
||||
|
||||
/// Set daily cumulative spend limit
|
||||
SetDailyLimit {
|
||||
/// Max NEAR amount per day
|
||||
amount: String,
|
||||
},
|
||||
|
||||
/// Set per-transaction auto-approve limit
|
||||
SetTxLimit {
|
||||
/// Max NEAR amount per transaction
|
||||
amount: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Run a key management command.
|
||||
pub async fn run_key_command(cmd: KeyCommand) -> anyhow::Result<()> {
|
||||
match cmd {
|
||||
KeyCommand::Generate {
|
||||
label,
|
||||
account,
|
||||
permission,
|
||||
receiver,
|
||||
methods,
|
||||
allowance,
|
||||
network,
|
||||
} => {
|
||||
let manager = create_key_manager().await?;
|
||||
let account_id = NearAccountId::new(&account)?;
|
||||
let network: NearNetwork = network.parse()?;
|
||||
let perm = parse_permission(&permission, receiver, methods, allowance)?;
|
||||
|
||||
let metadata = manager
|
||||
.generate_key(&label, &account_id, perm.clone(), network)
|
||||
.await?;
|
||||
|
||||
println!("Key generated successfully:");
|
||||
println!(" Label: {}", metadata.label);
|
||||
println!(" Account: {}", metadata.account_id);
|
||||
println!(" Public key: {}", metadata.public_key);
|
||||
println!(" Permission: {}", perm);
|
||||
println!(" Network: {}", metadata.network);
|
||||
|
||||
if matches!(perm, AccessKeyPermission::FullAccess) {
|
||||
println!();
|
||||
println!(
|
||||
" WARNING: This is a FULL ACCESS key for {}.",
|
||||
metadata.account_id
|
||||
);
|
||||
println!(" If this is the ONLY full-access key for this account and you lose it,");
|
||||
println!(" the account becomes permanently inaccessible.");
|
||||
println!();
|
||||
println!(" Create a backup: ironclaw key backup --output <file>");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
KeyCommand::Import {
|
||||
label,
|
||||
account,
|
||||
permission,
|
||||
receiver,
|
||||
methods,
|
||||
allowance,
|
||||
network,
|
||||
} => {
|
||||
let manager = create_key_manager().await?;
|
||||
let account_id = NearAccountId::new(&account)?;
|
||||
let network: NearNetwork = network.parse()?;
|
||||
let perm = parse_permission(&permission, receiver, methods, allowance)?;
|
||||
|
||||
// Read secret key from stdin (hidden)
|
||||
print!("Paste secret key (ed25519:...): ");
|
||||
std::io::stdout().flush()?;
|
||||
let secret_key = read_hidden_line()?;
|
||||
println!();
|
||||
|
||||
if secret_key.is_empty() {
|
||||
anyhow::bail!("No secret key provided");
|
||||
}
|
||||
|
||||
let metadata = manager
|
||||
.import_key(&label, &account_id, &secret_key, perm.clone(), network)
|
||||
.await?;
|
||||
|
||||
println!("Key imported successfully:");
|
||||
println!(" Label: {}", metadata.label);
|
||||
println!(" Account: {}", metadata.account_id);
|
||||
println!(" Public key: {}", metadata.public_key);
|
||||
println!(" Permission: {}", perm);
|
||||
|
||||
if matches!(perm, AccessKeyPermission::FullAccess) {
|
||||
println!();
|
||||
println!(" WARNING: Full-access key imported. Back it up!");
|
||||
println!(" ironclaw key backup --output <file>");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
KeyCommand::List { verbose } => {
|
||||
let manager = create_key_manager().await?;
|
||||
let keys = manager.list_keys().await?;
|
||||
|
||||
if keys.is_empty() {
|
||||
println!("No keys stored.");
|
||||
println!("Generate one: ironclaw key generate <label> --account <id>");
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
println!("Stored keys:");
|
||||
println!();
|
||||
for key in keys {
|
||||
if verbose {
|
||||
println!(" {} ({})", key.label, key.network);
|
||||
println!(" Account: {}", key.account_id);
|
||||
println!(" Public key: {}", key.public_key);
|
||||
println!(" Permission: {}", key.permission);
|
||||
println!(
|
||||
" Created: {}",
|
||||
key.created_at.format("%Y-%m-%d %H:%M UTC")
|
||||
);
|
||||
println!();
|
||||
} else {
|
||||
println!(
|
||||
" {} | {} | {} | {}",
|
||||
key.label, key.account_id, key.permission, key.network
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
KeyCommand::Info { label } => {
|
||||
let manager = create_key_manager().await?;
|
||||
let key = manager.get_key(&label).await?;
|
||||
|
||||
println!("Key: {}", key.label);
|
||||
println!(" Account: {}", key.account_id);
|
||||
println!(" Public key: {}", key.public_key);
|
||||
println!(" Permission: {}", key.permission);
|
||||
println!(" Network: {}", key.network);
|
||||
println!(
|
||||
" Created: {}",
|
||||
key.created_at.format("%Y-%m-%d %H:%M UTC")
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
KeyCommand::Remove { label } => {
|
||||
let manager = create_key_manager().await?;
|
||||
manager.remove_key(&label).await?;
|
||||
println!("Key '{}' removed.", label);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
KeyCommand::Export { label } => {
|
||||
let manager = create_key_manager().await?;
|
||||
let pubkey = manager.export_public_key(&label).await?;
|
||||
println!("{}", pubkey.to_near_format());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
KeyCommand::Policy(policy_cmd) => run_policy_command(policy_cmd).await,
|
||||
|
||||
KeyCommand::Backup { output, list } => {
|
||||
if list {
|
||||
// List keys in backup
|
||||
let data = fs::read(&output).await?;
|
||||
|
||||
print!("Backup passphrase: ");
|
||||
std::io::stdout().flush()?;
|
||||
let passphrase = read_hidden_line()?;
|
||||
println!();
|
||||
|
||||
// We need to decrypt to list, so restore to a temp manager
|
||||
// and just display, not actually import
|
||||
let plaintext = crate::keys::decrypt_backup(&passphrase, &data)?;
|
||||
let backup: serde_json::Value = serde_json::from_slice(&plaintext)?;
|
||||
|
||||
if let Some(keys) = backup.get("keys").and_then(|k| k.as_array()) {
|
||||
println!("Keys in backup ({}):", output.display());
|
||||
for key in keys {
|
||||
let label = key.get("label").and_then(|l| l.as_str()).unwrap_or("?");
|
||||
let account = key
|
||||
.get("account_id")
|
||||
.and_then(|a| a.as_str())
|
||||
.unwrap_or("?");
|
||||
println!(" {} ({})", label, account);
|
||||
}
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let manager = create_key_manager().await?;
|
||||
|
||||
print!("Backup passphrase: ");
|
||||
std::io::stdout().flush()?;
|
||||
let passphrase = read_hidden_line()?;
|
||||
println!();
|
||||
|
||||
print!("Confirm passphrase: ");
|
||||
std::io::stdout().flush()?;
|
||||
let confirm = read_hidden_line()?;
|
||||
println!();
|
||||
|
||||
if passphrase != confirm {
|
||||
anyhow::bail!("Passphrases do not match");
|
||||
}
|
||||
|
||||
if passphrase.len() < 8 {
|
||||
anyhow::bail!("Passphrase must be at least 8 characters");
|
||||
}
|
||||
|
||||
let backup_data = manager.create_backup(&passphrase).await?;
|
||||
fs::write(&output, &backup_data).await?;
|
||||
|
||||
println!(
|
||||
"Backup created: {} ({} bytes)",
|
||||
output.display(),
|
||||
backup_data.len()
|
||||
);
|
||||
println!("Store this file securely. You'll need the passphrase to restore.");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
KeyCommand::Restore { path } => {
|
||||
let manager = create_key_manager().await?;
|
||||
|
||||
let data = fs::read(&path).await?;
|
||||
|
||||
print!("Backup passphrase: ");
|
||||
std::io::stdout().flush()?;
|
||||
let passphrase = read_hidden_line()?;
|
||||
println!();
|
||||
|
||||
let restored = manager.restore_backup(&data, &passphrase).await?;
|
||||
|
||||
if restored.is_empty() {
|
||||
println!("No new keys to restore (all already exist).");
|
||||
} else {
|
||||
println!("Restored {} keys:", restored.len());
|
||||
for label in &restored {
|
||||
println!(" {}", label);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_policy_command(cmd: PolicyCommand) -> anyhow::Result<()> {
|
||||
let policy_path = default_policy_path();
|
||||
|
||||
match cmd {
|
||||
PolicyCommand::Show => {
|
||||
let policy = load_policy(&policy_path).await?;
|
||||
let json = serde_json::to_string_pretty(&policy)?;
|
||||
println!("{}", json);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
PolicyCommand::SetTransferLimit { amount } => {
|
||||
let yocto = parse_near_amount(&amount)?;
|
||||
let mut policy = load_policy(&policy_path).await?;
|
||||
policy.transfer_auto_approve_max_yocto = yocto;
|
||||
save_policy(&policy_path, &policy).await?;
|
||||
println!("Transfer auto-approve limit set to {}", format_yocto(yocto));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
PolicyCommand::WhitelistAccount {
|
||||
account,
|
||||
max_transfer,
|
||||
} => {
|
||||
let mut policy = load_policy(&policy_path).await?;
|
||||
if !policy.transfer_whitelist.contains(&account) {
|
||||
policy.transfer_whitelist.push(account.clone());
|
||||
}
|
||||
if let Some(max) = max_transfer {
|
||||
policy.transfer_whitelist_max_yocto = parse_near_amount(&max)?;
|
||||
}
|
||||
save_policy(&policy_path, &policy).await?;
|
||||
println!("Account '{}' added to transfer whitelist", account);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
PolicyCommand::WhitelistValidator {
|
||||
validator,
|
||||
max_stake,
|
||||
} => {
|
||||
let mut policy = load_policy(&policy_path).await?;
|
||||
if !policy.stake_validator_whitelist.contains(&validator) {
|
||||
policy.stake_validator_whitelist.push(validator.clone());
|
||||
}
|
||||
if let Some(max) = max_stake {
|
||||
policy.stake_auto_approve_max_yocto = parse_near_amount(&max)?;
|
||||
}
|
||||
save_policy(&policy_path, &policy).await?;
|
||||
println!("Validator '{}' added to staking whitelist", validator);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
PolicyCommand::AddContractRule {
|
||||
contract,
|
||||
methods,
|
||||
max_deposit,
|
||||
auto_approve,
|
||||
} => {
|
||||
let mut policy = load_policy(&policy_path).await?;
|
||||
let deposit = parse_near_amount(&max_deposit)?;
|
||||
let method_list = methods
|
||||
.map(|m| m.split(',').map(|s| s.trim().to_string()).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
policy.function_call_rules.push(FunctionCallRule {
|
||||
receiver_id: contract.clone(),
|
||||
allowed_methods: method_list,
|
||||
max_deposit_yocto: deposit,
|
||||
max_gas: None,
|
||||
auto_approve,
|
||||
});
|
||||
save_policy(&policy_path, &policy).await?;
|
||||
println!(
|
||||
"Contract rule added for '{}' (auto_approve={})",
|
||||
contract, auto_approve
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
PolicyCommand::AddChainSigRule {
|
||||
path_pattern,
|
||||
domain,
|
||||
max_payload,
|
||||
auto_approve,
|
||||
} => {
|
||||
let domain = match domain.to_lowercase().as_str() {
|
||||
"secp256k1" => SignatureDomain::Secp256k1,
|
||||
"ed25519" => SignatureDomain::Ed25519,
|
||||
other => anyhow::bail!("Unknown domain '{}', expected secp256k1 or ed25519", other),
|
||||
};
|
||||
|
||||
let mut policy = load_policy(&policy_path).await?;
|
||||
policy.chain_sig_rules.push(ChainSigRule {
|
||||
allowed_paths: vec![path_pattern.clone()],
|
||||
allowed_domains: vec![domain],
|
||||
max_payload_bytes: max_payload,
|
||||
auto_approve,
|
||||
});
|
||||
save_policy(&policy_path, &policy).await?;
|
||||
println!(
|
||||
"Chain signature rule added for '{}' (auto_approve={})",
|
||||
path_pattern, auto_approve
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
PolicyCommand::SetDailyLimit { amount } => {
|
||||
let yocto = parse_near_amount(&amount)?;
|
||||
let mut policy = load_policy(&policy_path).await?;
|
||||
policy.daily_spend_limit_yocto = Some(yocto);
|
||||
save_policy(&policy_path, &policy).await?;
|
||||
println!("Daily spend limit set to {}", format_yocto(yocto));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
PolicyCommand::SetTxLimit { amount } => {
|
||||
let yocto = parse_near_amount(&amount)?;
|
||||
let mut policy = load_policy(&policy_path).await?;
|
||||
policy.per_tx_auto_approve_max_yocto = yocto;
|
||||
save_policy(&policy_path, &policy).await?;
|
||||
println!(
|
||||
"Per-transaction auto-approve limit set to {}",
|
||||
format_yocto(yocto)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn load_policy(path: &PathBuf) -> anyhow::Result<PolicyConfig> {
|
||||
if path.exists() {
|
||||
let content = fs::read_to_string(path).await?;
|
||||
Ok(serde_json::from_str(&content)?)
|
||||
} else {
|
||||
Ok(PolicyConfig::default())
|
||||
}
|
||||
}
|
||||
|
||||
async fn save_policy(path: &PathBuf, policy: &PolicyConfig) -> anyhow::Result<()> {
|
||||
if let Some(parent) = path.parent() {
|
||||
fs::create_dir_all(parent).await?;
|
||||
}
|
||||
let content = serde_json::to_string_pretty(policy)?;
|
||||
fs::write(path, content).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_permission(
|
||||
permission: &str,
|
||||
receiver: Option<String>,
|
||||
methods: Option<String>,
|
||||
allowance: Option<String>,
|
||||
) -> anyhow::Result<AccessKeyPermission> {
|
||||
match permission {
|
||||
"full-access" | "FullAccess" => Ok(AccessKeyPermission::FullAccess),
|
||||
"function-call" | "FunctionCall" => {
|
||||
let receiver_id = receiver
|
||||
.ok_or_else(|| anyhow::anyhow!("--receiver required for function-call keys"))?;
|
||||
|
||||
let method_names = methods
|
||||
.map(|m| m.split(',').map(|s| s.trim().to_string()).collect())
|
||||
.unwrap_or_default();
|
||||
|
||||
let allowance_yocto = allowance
|
||||
.map(|a| parse_near_amount(&a))
|
||||
.transpose()
|
||||
.map_err(|e| anyhow::anyhow!("invalid allowance: {}", e))?;
|
||||
|
||||
Ok(AccessKeyPermission::FunctionCall {
|
||||
allowance: allowance_yocto,
|
||||
receiver_id,
|
||||
method_names,
|
||||
})
|
||||
}
|
||||
other => Err(anyhow::anyhow!(
|
||||
"unknown permission '{}', expected full-access or function-call",
|
||||
other
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a KeyManager with the default secrets store.
|
||||
async fn create_key_manager() -> anyhow::Result<KeyManager> {
|
||||
let config = Config::from_env()?;
|
||||
let master_key = config.secrets.master_key().ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"SECRETS_MASTER_KEY not set. Run 'ironclaw onboard' first or set it in .env"
|
||||
)
|
||||
})?;
|
||||
|
||||
let store = Store::new(&config.database).await?;
|
||||
store.run_migrations().await?;
|
||||
|
||||
let crypto = SecretsCrypto::new(master_key.clone())?;
|
||||
let secrets_store: Arc<dyn SecretsStore + Send + Sync> =
|
||||
Arc::new(PostgresSecretsStore::new(store.pool(), Arc::new(crypto)));
|
||||
|
||||
let manager = KeyManager::new(secrets_store, "default".to_string());
|
||||
|
||||
// Load policy if it exists
|
||||
let policy_path = default_policy_path();
|
||||
if policy_path.exists() {
|
||||
let content = fs::read_to_string(&policy_path).await?;
|
||||
let policy: PolicyConfig = serde_json::from_str(&content)?;
|
||||
Ok(manager.with_policy(policy))
|
||||
} else {
|
||||
Ok(manager)
|
||||
}
|
||||
}
|
||||
|
||||
/// Read a line of input with hidden characters.
|
||||
fn read_hidden_line() -> anyhow::Result<String> {
|
||||
use crossterm::{
|
||||
event::{self, Event, KeyCode, KeyModifiers},
|
||||
terminal,
|
||||
};
|
||||
|
||||
let mut input = String::new();
|
||||
terminal::enable_raw_mode()?;
|
||||
|
||||
loop {
|
||||
if let Event::Key(key_event) = event::read()? {
|
||||
match key_event.code {
|
||||
KeyCode::Enter => break,
|
||||
KeyCode::Backspace => {
|
||||
if !input.is_empty() {
|
||||
input.pop();
|
||||
print!("\x08 \x08");
|
||||
std::io::stdout().flush()?;
|
||||
}
|
||||
}
|
||||
KeyCode::Char('c') if key_event.modifiers.contains(KeyModifiers::CONTROL) => {
|
||||
terminal::disable_raw_mode()?;
|
||||
return Err(anyhow::anyhow!("Interrupted"));
|
||||
}
|
||||
KeyCode::Char(c) => {
|
||||
input.push(c);
|
||||
print!("*");
|
||||
std::io::stdout().flush()?;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
terminal::disable_raw_mode()?;
|
||||
Ok(input)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::cli::key::parse_permission;
|
||||
use crate::keys::types::AccessKeyPermission;
|
||||
|
||||
#[test]
|
||||
fn test_parse_full_access() {
|
||||
let perm = parse_permission("full-access", None, None, None).unwrap();
|
||||
assert!(matches!(perm, AccessKeyPermission::FullAccess));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_function_call() {
|
||||
let perm = parse_permission(
|
||||
"function-call",
|
||||
Some("contract.near".to_string()),
|
||||
Some("deposit,withdraw".to_string()),
|
||||
Some("1.5".to_string()),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
match perm {
|
||||
AccessKeyPermission::FunctionCall {
|
||||
receiver_id,
|
||||
method_names,
|
||||
allowance,
|
||||
} => {
|
||||
assert_eq!(receiver_id, "contract.near");
|
||||
assert_eq!(method_names, vec!["deposit", "withdraw"]);
|
||||
assert!(allowance.is_some());
|
||||
}
|
||||
_ => panic!("expected FunctionCall"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_function_call_missing_receiver() {
|
||||
let result = parse_permission("function-call", None, None, None);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
}
|
||||
@@ -10,12 +10,14 @@
|
||||
//! - Checking system health (`status`)
|
||||
|
||||
mod config;
|
||||
pub mod key;
|
||||
mod mcp;
|
||||
pub mod memory;
|
||||
pub mod status;
|
||||
mod tool;
|
||||
|
||||
pub use config::{ConfigCommand, run_config_command};
|
||||
pub use key::{KeyCommand, run_key_command};
|
||||
pub use mcp::{McpCommand, run_mcp_command};
|
||||
pub use memory::{MemoryCommand, run_memory_command};
|
||||
pub use status::run_status_command;
|
||||
@@ -78,6 +80,10 @@ pub enum Command {
|
||||
#[command(subcommand)]
|
||||
Tool(ToolCommand),
|
||||
|
||||
/// Manage NEAR blockchain keys
|
||||
#[command(subcommand)]
|
||||
Key(KeyCommand),
|
||||
|
||||
/// Manage MCP servers (hosted tool providers)
|
||||
#[command(subcommand)]
|
||||
Mcp(McpCommand),
|
||||
|
||||
@@ -39,6 +39,9 @@ pub enum Error {
|
||||
|
||||
#[error("Workspace error: {0}")]
|
||||
Workspace(#[from] WorkspaceError),
|
||||
|
||||
#[error("Key management error: {0}")]
|
||||
Key(#[from] crate::keys::KeyError),
|
||||
}
|
||||
|
||||
/// Configuration-related errors.
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
//! Cross-chain signing via v1.signer MPC contract.
|
||||
//!
|
||||
//! Enables signing payloads for other chains (Ethereum, Bitcoin, etc.)
|
||||
//! using NEAR's chain signatures MPC network.
|
||||
|
||||
use crate::keys::KeyError;
|
||||
use crate::keys::policy::SignatureDomain;
|
||||
use crate::keys::transaction::{Action, FunctionCall, MAX_GAS, ONE_YOCTO};
|
||||
|
||||
/// The chain signatures MPC contract on mainnet.
|
||||
pub const CHAIN_SIGNATURES_CONTRACT_MAINNET: &str = "v1.signer";
|
||||
|
||||
/// The chain signatures MPC contract on testnet.
|
||||
pub const CHAIN_SIGNATURES_CONTRACT_TESTNET: &str = "v1.signer-prod.testnet";
|
||||
|
||||
/// Build a FunctionCall action for requesting a chain signature.
|
||||
pub fn build_chain_signature_action(
|
||||
payload: &[u8],
|
||||
derivation_path: &str,
|
||||
_domain: SignatureDomain,
|
||||
) -> Result<Action, KeyError> {
|
||||
let args = serde_json::json!({
|
||||
"request": {
|
||||
"payload": payload.iter().map(|b| *b as u32).collect::<Vec<u32>>(),
|
||||
"path": derivation_path,
|
||||
"key_version": 0,
|
||||
},
|
||||
});
|
||||
|
||||
let args_bytes = serde_json::to_vec(&args).map_err(|e| KeyError::ChainSignatureError {
|
||||
reason: format!("failed to serialize chain sig args: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(Action::FunctionCall(FunctionCall {
|
||||
method_name: "sign".to_string(),
|
||||
args: args_bytes,
|
||||
gas: MAX_GAS,
|
||||
deposit: ONE_YOCTO,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Parse the result of a chain signature request from the transaction outcome.
|
||||
pub fn parse_chain_signature_result(
|
||||
outcome: &serde_json::Value,
|
||||
) -> Result<ChainSignatureResult, KeyError> {
|
||||
// The result is in the SuccessValue field, base64-encoded
|
||||
let success_value = outcome
|
||||
.get("SuccessValue")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| KeyError::ChainSignatureError {
|
||||
reason: "no SuccessValue in chain signature outcome".to_string(),
|
||||
})?;
|
||||
|
||||
let decoded = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, success_value)
|
||||
.map_err(|e| KeyError::ChainSignatureError {
|
||||
reason: format!("failed to decode chain sig result: {}", e),
|
||||
})?;
|
||||
|
||||
let result_str = String::from_utf8(decoded).map_err(|e| KeyError::ChainSignatureError {
|
||||
reason: format!("chain sig result is not UTF-8: {}", e),
|
||||
})?;
|
||||
|
||||
let result_json: serde_json::Value =
|
||||
serde_json::from_str(&result_str).map_err(|e| KeyError::ChainSignatureError {
|
||||
reason: format!("failed to parse chain sig result JSON: {}", e),
|
||||
})?;
|
||||
|
||||
// Extract big_r and s components
|
||||
let big_r = result_json
|
||||
.get("big_r")
|
||||
.and_then(|v| v.get("affine_point"))
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| KeyError::ChainSignatureError {
|
||||
reason: "missing big_r.affine_point in chain sig result".to_string(),
|
||||
})?
|
||||
.to_string();
|
||||
|
||||
let s = result_json
|
||||
.get("s")
|
||||
.and_then(|v| v.get("scalar"))
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| KeyError::ChainSignatureError {
|
||||
reason: "missing s.scalar in chain sig result".to_string(),
|
||||
})?
|
||||
.to_string();
|
||||
|
||||
let recovery_id = result_json
|
||||
.get("recovery_id")
|
||||
.and_then(|v| v.as_u64())
|
||||
.map(|v| v as u8);
|
||||
|
||||
Ok(ChainSignatureResult {
|
||||
big_r,
|
||||
s,
|
||||
recovery_id,
|
||||
})
|
||||
}
|
||||
|
||||
/// Result from a chain signature request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChainSignatureResult {
|
||||
/// The R component (affine point, hex-encoded).
|
||||
pub big_r: String,
|
||||
/// The s component (scalar, hex-encoded).
|
||||
pub s: String,
|
||||
/// Recovery ID for ECDSA (relevant for Ethereum).
|
||||
pub recovery_id: Option<u8>,
|
||||
}
|
||||
|
||||
/// Get the chain signatures contract address for a network.
|
||||
pub fn chain_sig_contract(network: &crate::keys::types::NearNetwork) -> &str {
|
||||
match network {
|
||||
crate::keys::types::NearNetwork::Mainnet => CHAIN_SIGNATURES_CONTRACT_MAINNET,
|
||||
crate::keys::types::NearNetwork::Testnet => CHAIN_SIGNATURES_CONTRACT_TESTNET,
|
||||
crate::keys::types::NearNetwork::Custom(_) => CHAIN_SIGNATURES_CONTRACT_TESTNET,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::keys::chain_signatures::{
|
||||
build_chain_signature_action, chain_sig_contract, parse_chain_signature_result,
|
||||
};
|
||||
use crate::keys::policy::SignatureDomain;
|
||||
use crate::keys::transaction::{Action, MAX_GAS, ONE_YOCTO};
|
||||
use crate::keys::types::NearNetwork;
|
||||
|
||||
#[test]
|
||||
fn test_build_chain_signature_action() {
|
||||
let payload = vec![0u8; 32];
|
||||
let action =
|
||||
build_chain_signature_action(&payload, "ethereum-1", SignatureDomain::Secp256k1)
|
||||
.unwrap();
|
||||
|
||||
match action {
|
||||
Action::FunctionCall(fc) => {
|
||||
assert_eq!(fc.method_name, "sign");
|
||||
assert_eq!(fc.gas, MAX_GAS);
|
||||
assert_eq!(fc.deposit, ONE_YOCTO);
|
||||
|
||||
// Verify args parse correctly
|
||||
let args: serde_json::Value = serde_json::from_slice(&fc.args).unwrap();
|
||||
assert!(args.get("request").is_some());
|
||||
let path = args["request"]["path"].as_str().unwrap();
|
||||
assert_eq!(path, "ethereum-1");
|
||||
}
|
||||
_ => panic!("expected FunctionCall action"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_chain_signature_result() {
|
||||
let result_json = serde_json::json!({
|
||||
"big_r": {"affine_point": "02abc123"},
|
||||
"s": {"scalar": "def456"},
|
||||
"recovery_id": 0
|
||||
});
|
||||
|
||||
let result_str = serde_json::to_string(&result_json).unwrap();
|
||||
let encoded = base64::Engine::encode(
|
||||
&base64::engine::general_purpose::STANDARD,
|
||||
result_str.as_bytes(),
|
||||
);
|
||||
|
||||
let outcome = serde_json::json!({"SuccessValue": encoded});
|
||||
let result = parse_chain_signature_result(&outcome).unwrap();
|
||||
|
||||
assert_eq!(result.big_r, "02abc123");
|
||||
assert_eq!(result.s, "def456");
|
||||
assert_eq!(result.recovery_id, Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chain_sig_contract_addresses() {
|
||||
assert_eq!(chain_sig_contract(&NearNetwork::Mainnet), "v1.signer");
|
||||
assert_eq!(
|
||||
chain_sig_contract(&NearNetwork::Testnet),
|
||||
"v1.signer-prod.testnet"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
//! Error types for NEAR key management.
|
||||
|
||||
use crate::secrets::SecretError;
|
||||
|
||||
/// Errors from NEAR key operations.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum KeyError {
|
||||
#[error("Key not found: {label}")]
|
||||
NotFound { label: String },
|
||||
|
||||
#[error("Key already exists: {label}")]
|
||||
AlreadyExists { label: String },
|
||||
|
||||
#[error("Invalid key format: {reason}")]
|
||||
InvalidKeyFormat { reason: String },
|
||||
|
||||
#[error("Invalid account ID: {reason}")]
|
||||
InvalidAccountId { reason: String },
|
||||
|
||||
#[error("Signing failed: {reason}")]
|
||||
SigningFailed { reason: String },
|
||||
|
||||
#[error("Serialization failed: {0}")]
|
||||
SerializationFailed(String),
|
||||
|
||||
#[error("Approval required: {operation}")]
|
||||
ApprovalRequired { operation: String },
|
||||
|
||||
#[error("Policy denied: {reason}")]
|
||||
PolicyDenied { reason: String },
|
||||
|
||||
#[error("RPC error: {reason}")]
|
||||
RpcError { reason: String },
|
||||
|
||||
#[error("Stale nonce: cached {cached}, chain {chain}")]
|
||||
StaleNonce { cached: u64, chain: u64 },
|
||||
|
||||
#[error("Insufficient allowance: needed {needed}, available {available}")]
|
||||
InsufficientAllowance { needed: u128, available: u128 },
|
||||
|
||||
#[error("Permission denied: {reason}")]
|
||||
PermissionDenied { reason: String },
|
||||
|
||||
#[error("Chain signature error: {reason}")]
|
||||
ChainSignatureError { reason: String },
|
||||
|
||||
#[error("Backup error: {reason}")]
|
||||
BackupError { reason: String },
|
||||
|
||||
#[error("Secret store error: {0}")]
|
||||
SecretStore(#[from] SecretError),
|
||||
|
||||
#[error("Network error: {0}")]
|
||||
Network(#[from] reqwest::Error),
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
}
|
||||
@@ -0,0 +1,183 @@
|
||||
//! NEP-413 intent construction and signing.
|
||||
//!
|
||||
//! Provides types and signing for NEAR intents following the NEP-413 standard.
|
||||
//! Intents are signed messages that authorize actions on a verifying contract.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::keys::KeyError;
|
||||
use crate::keys::signer::sign_hash;
|
||||
use crate::keys::types::NearPublicKey;
|
||||
use crate::secrets::SecretsStore;
|
||||
|
||||
/// NEP-413 intent message to be signed.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct IntentMessage {
|
||||
/// Account signing the intent.
|
||||
pub signer_id: String,
|
||||
/// Contract that will verify the signature.
|
||||
pub verifying_contract: String,
|
||||
/// Deadline (block height or timestamp) after which the intent expires.
|
||||
pub deadline: String,
|
||||
/// Unique nonce to prevent replay.
|
||||
pub nonce: String,
|
||||
/// List of intent actions.
|
||||
pub intents: Vec<IntentAction>,
|
||||
}
|
||||
|
||||
/// An action within an intent.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum IntentAction {
|
||||
/// Token difference (swap, deposit, etc.)
|
||||
TokenDiff { token: String, amount: String },
|
||||
/// Add a public key to the account.
|
||||
AddPublicKey { public_key: String },
|
||||
/// Custom action with arbitrary data.
|
||||
Custom {
|
||||
action_type: String,
|
||||
data: serde_json::Value,
|
||||
},
|
||||
}
|
||||
|
||||
/// A signed NEP-413 intent.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SignedIntent {
|
||||
pub standard: String,
|
||||
pub payload: IntentMessage,
|
||||
pub public_key: String,
|
||||
pub signature: String,
|
||||
}
|
||||
|
||||
/// Construct the NEP-413 signing payload.
|
||||
///
|
||||
/// The payload is: SHA-256(tag + message_json + nonce + recipient)
|
||||
/// where tag is the NEP-413 tag prefix.
|
||||
pub fn nep413_signing_payload(message: &IntentMessage) -> Result<[u8; 32], KeyError> {
|
||||
let message_json = serde_json::to_string(message).map_err(|e| {
|
||||
KeyError::SerializationFailed(format!("failed to serialize intent message: {}", e))
|
||||
})?;
|
||||
|
||||
// NEP-413 tag
|
||||
const NEP413_TAG: u32 = 2147484061; // (1 << 31) + 413
|
||||
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(NEP413_TAG.to_le_bytes());
|
||||
hasher.update(message_json.as_bytes());
|
||||
|
||||
Ok(hasher.finalize().into())
|
||||
}
|
||||
|
||||
/// Sign an intent message using a key from the secrets store.
|
||||
pub async fn sign_intent(
|
||||
secrets_store: &dyn SecretsStore,
|
||||
user_id: &str,
|
||||
label: &str,
|
||||
public_key: &NearPublicKey,
|
||||
intent: IntentMessage,
|
||||
) -> Result<SignedIntent, KeyError> {
|
||||
let hash = nep413_signing_payload(&intent)?;
|
||||
let signature_bytes = sign_hash(secrets_store, user_id, label, &hash).await?;
|
||||
|
||||
// Base64-encode the signature
|
||||
let signature =
|
||||
base64::Engine::encode(&base64::engine::general_purpose::STANDARD, signature_bytes);
|
||||
|
||||
Ok(SignedIntent {
|
||||
standard: "nep413".to_string(),
|
||||
payload: intent,
|
||||
public_key: public_key.to_near_format(),
|
||||
signature,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use ed25519_dalek::SigningKey;
|
||||
use secrecy::SecretString;
|
||||
|
||||
use crate::keys::intents::{IntentAction, IntentMessage, nep413_signing_payload, sign_intent};
|
||||
use crate::keys::signer::public_key_from_secret;
|
||||
use crate::secrets::{CreateSecretParams, InMemorySecretsStore, SecretsCrypto, SecretsStore};
|
||||
|
||||
fn test_store() -> Arc<InMemorySecretsStore> {
|
||||
let key = "0123456789abcdef0123456789abcdef";
|
||||
let crypto = Arc::new(SecretsCrypto::new(SecretString::from(key.to_string())).unwrap());
|
||||
Arc::new(InMemorySecretsStore::new(crypto))
|
||||
}
|
||||
|
||||
fn test_intent() -> IntentMessage {
|
||||
IntentMessage {
|
||||
signer_id: "alice.near".to_string(),
|
||||
verifying_contract: "intents.near".to_string(),
|
||||
deadline: "100000000".to_string(),
|
||||
nonce: "unique-nonce-123".to_string(),
|
||||
intents: vec![IntentAction::TokenDiff {
|
||||
token: "wrap.near".to_string(),
|
||||
amount: "1000000".to_string(),
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nep413_payload_deterministic() {
|
||||
let intent = test_intent();
|
||||
let hash1 = nep413_signing_payload(&intent).unwrap();
|
||||
let hash2 = nep413_signing_payload(&intent).unwrap();
|
||||
assert_eq!(hash1, hash2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_nep413_payload_different_nonces() {
|
||||
let mut intent1 = test_intent();
|
||||
let mut intent2 = test_intent();
|
||||
intent1.nonce = "nonce-1".to_string();
|
||||
intent2.nonce = "nonce-2".to_string();
|
||||
|
||||
let hash1 = nep413_signing_payload(&intent1).unwrap();
|
||||
let hash2 = nep413_signing_payload(&intent2).unwrap();
|
||||
assert_ne!(hash1, hash2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sign_intent_roundtrip() {
|
||||
let store = test_store();
|
||||
|
||||
// Generate a key
|
||||
let signing_key = SigningKey::generate(&mut rand::rngs::OsRng);
|
||||
let verifying_key = signing_key.verifying_key();
|
||||
let mut combined = Vec::with_capacity(64);
|
||||
combined.extend_from_slice(signing_key.as_bytes());
|
||||
combined.extend_from_slice(verifying_key.as_bytes());
|
||||
let secret = format!("ed25519:{}", bs58::encode(&combined).into_string());
|
||||
|
||||
store
|
||||
.create(
|
||||
"user1",
|
||||
CreateSecretParams::new("near_key:intent-signer", &secret)
|
||||
.with_provider("near_keys"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let public_key = public_key_from_secret(&secret).unwrap();
|
||||
let intent = test_intent();
|
||||
|
||||
let signed = sign_intent(
|
||||
store.as_ref(),
|
||||
"user1",
|
||||
"intent-signer",
|
||||
&public_key,
|
||||
intent,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(signed.standard, "nep413");
|
||||
assert_eq!(signed.public_key, public_key.to_near_format());
|
||||
assert!(!signed.signature.is_empty());
|
||||
}
|
||||
}
|
||||
+994
@@ -0,0 +1,994 @@
|
||||
//! NEAR key management for IronClaw.
|
||||
//!
|
||||
//! Manages NEAR Protocol blockchain keys so the agent can sign transactions,
|
||||
//! intents, and cross-chain signature requests.
|
||||
//!
|
||||
//! # Security Model
|
||||
//!
|
||||
//! Hybrid custody: the agent holds scoped function-call keys for routine
|
||||
//! operations. High-value operations require explicit user approval.
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌─────────────────────────────────────────────────────────────────────────┐
|
||||
//! │ Key Management │
|
||||
//! │ │
|
||||
//! │ KeyManager ──► SecretsStore (AES-256-GCM encrypted private keys) │
|
||||
//! │ │ │
|
||||
//! │ ├──► Signer (ed25519 sign, Zeroize on drop) │
|
||||
//! │ ├──► Policy (analyze transaction, evaluate rules, approve/deny) │
|
||||
//! │ ├──► SpendTracker (daily cumulative spend) │
|
||||
//! │ └──► RPC Client (nonce, submit, status) │
|
||||
//! │ │
|
||||
//! │ INVARIANT: Private keys NEVER reach the LLM or WASM boundary. │
|
||||
//! └─────────────────────────────────────────────────────────────────────────┘
|
||||
//! ```
|
||||
|
||||
pub mod chain_signatures;
|
||||
mod error;
|
||||
pub mod intents;
|
||||
pub mod policy;
|
||||
pub mod rpc;
|
||||
pub mod signer;
|
||||
pub mod spending;
|
||||
pub mod transaction;
|
||||
pub mod types;
|
||||
|
||||
pub use error::KeyError;
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use chrono::Utc;
|
||||
use ed25519_dalek::SigningKey;
|
||||
use tokio::fs;
|
||||
use zeroize::Zeroize;
|
||||
|
||||
use crate::keys::policy::{
|
||||
ChainSigAnalysis, PolicyConfig, PolicyDecision, SignatureDomain, analyze_transaction,
|
||||
infer_target_chain,
|
||||
};
|
||||
use crate::keys::rpc::NearRpcClient;
|
||||
use crate::keys::signer::{public_key_from_secret, sign_hash};
|
||||
use crate::keys::spending::SpendTracker;
|
||||
use crate::keys::transaction::{BlockHash, Signature, SignedTransaction, Transaction};
|
||||
use crate::keys::types::{
|
||||
AccessKeyPermission, KeyMetadata, KeyStore, KeyType, NearAccountId, NearNetwork, NearPublicKey,
|
||||
};
|
||||
use crate::secrets::{CreateSecretParams, SecretsStore};
|
||||
|
||||
/// Result of a signing operation.
|
||||
#[derive(Debug)]
|
||||
pub enum SignResult {
|
||||
/// Transaction was signed (policy auto-approved).
|
||||
Signed {
|
||||
transaction: SignedTransaction,
|
||||
analysis: policy::TransactionAnalysis,
|
||||
},
|
||||
/// User must approve before signing can proceed.
|
||||
ApprovalRequired {
|
||||
analysis: policy::TransactionAnalysis,
|
||||
reasons: Vec<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Central key management struct.
|
||||
pub struct KeyManager {
|
||||
secrets_store: Arc<dyn SecretsStore + Send + Sync>,
|
||||
metadata_path: PathBuf,
|
||||
policy: PolicyConfig,
|
||||
spend_tracker: SpendTracker,
|
||||
user_id: String,
|
||||
}
|
||||
|
||||
impl KeyManager {
|
||||
/// Create a new KeyManager.
|
||||
pub fn new(secrets_store: Arc<dyn SecretsStore + Send + Sync>, user_id: String) -> Self {
|
||||
Self {
|
||||
secrets_store,
|
||||
metadata_path: default_keys_path(),
|
||||
policy: PolicyConfig::default(),
|
||||
spend_tracker: SpendTracker::new(SpendTracker::default_path()),
|
||||
user_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set a custom metadata path (for testing).
|
||||
pub fn with_metadata_path(mut self, path: PathBuf) -> Self {
|
||||
self.metadata_path = path;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the policy config.
|
||||
pub fn with_policy(mut self, policy: PolicyConfig) -> Self {
|
||||
self.policy = policy;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set a custom spend tracker (for testing).
|
||||
pub fn with_spend_tracker(mut self, tracker: SpendTracker) -> Self {
|
||||
self.spend_tracker = tracker;
|
||||
self
|
||||
}
|
||||
|
||||
/// Get a reference to the current policy config.
|
||||
pub fn policy(&self) -> &PolicyConfig {
|
||||
&self.policy
|
||||
}
|
||||
|
||||
/// Get a mutable reference to the policy config.
|
||||
pub fn policy_mut(&mut self) -> &mut PolicyConfig {
|
||||
&mut self.policy
|
||||
}
|
||||
|
||||
// -- Key lifecycle --
|
||||
|
||||
/// Generate a new ed25519 keypair and store it.
|
||||
pub async fn generate_key(
|
||||
&self,
|
||||
label: &str,
|
||||
account_id: &NearAccountId,
|
||||
permission: AccessKeyPermission,
|
||||
network: NearNetwork,
|
||||
) -> Result<KeyMetadata, KeyError> {
|
||||
// Check for duplicates
|
||||
let store = self.load_store().await?;
|
||||
if store.keys.contains_key(label) {
|
||||
return Err(KeyError::AlreadyExists {
|
||||
label: label.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Generate keypair
|
||||
let signing_key = SigningKey::generate(&mut rand::rngs::OsRng);
|
||||
let verifying_key = signing_key.verifying_key();
|
||||
|
||||
// Build NEAR-format secret: ed25519:<base58(seed + pubkey)>
|
||||
let mut combined = Vec::with_capacity(64);
|
||||
combined.extend_from_slice(signing_key.as_bytes());
|
||||
combined.extend_from_slice(verifying_key.as_bytes());
|
||||
let secret_key = format!("ed25519:{}", bs58::encode(&combined).into_string());
|
||||
combined.zeroize();
|
||||
// signing_key drops here (Zeroize on drop)
|
||||
|
||||
let public_key = NearPublicKey {
|
||||
key_type: KeyType::Ed25519,
|
||||
data: verifying_key.to_bytes(),
|
||||
};
|
||||
|
||||
// Store private key in secrets store
|
||||
let secret_name = format!("near_key:{}", label);
|
||||
self.secrets_store
|
||||
.create(
|
||||
&self.user_id,
|
||||
CreateSecretParams::new(&secret_name, &secret_key).with_provider("near_keys"),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Build metadata
|
||||
let metadata = KeyMetadata {
|
||||
label: label.to_string(),
|
||||
account_id: account_id.to_string(),
|
||||
public_key: public_key.to_near_format(),
|
||||
permission,
|
||||
network,
|
||||
created_at: Utc::now(),
|
||||
cached_nonce: None,
|
||||
};
|
||||
|
||||
// Save metadata
|
||||
let mut store = self.load_store().await?;
|
||||
store.keys.insert(label.to_string(), metadata.clone());
|
||||
self.save_store(&store).await?;
|
||||
|
||||
Ok(metadata)
|
||||
}
|
||||
|
||||
/// Import an existing key from a NEAR-format secret key string.
|
||||
pub async fn import_key(
|
||||
&self,
|
||||
label: &str,
|
||||
account_id: &NearAccountId,
|
||||
secret_key: &str,
|
||||
permission: AccessKeyPermission,
|
||||
network: NearNetwork,
|
||||
) -> Result<KeyMetadata, KeyError> {
|
||||
// Check for duplicates
|
||||
let store = self.load_store().await?;
|
||||
if store.keys.contains_key(label) {
|
||||
return Err(KeyError::AlreadyExists {
|
||||
label: label.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Validate and derive public key
|
||||
let public_key = public_key_from_secret(secret_key)?;
|
||||
|
||||
// Store private key in secrets store
|
||||
let secret_name = format!("near_key:{}", label);
|
||||
self.secrets_store
|
||||
.create(
|
||||
&self.user_id,
|
||||
CreateSecretParams::new(&secret_name, secret_key).with_provider("near_keys"),
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Build metadata
|
||||
let metadata = KeyMetadata {
|
||||
label: label.to_string(),
|
||||
account_id: account_id.to_string(),
|
||||
public_key: public_key.to_near_format(),
|
||||
permission,
|
||||
network,
|
||||
created_at: Utc::now(),
|
||||
cached_nonce: None,
|
||||
};
|
||||
|
||||
// Save metadata
|
||||
let mut store = self.load_store().await?;
|
||||
store.keys.insert(label.to_string(), metadata.clone());
|
||||
self.save_store(&store).await?;
|
||||
|
||||
Ok(metadata)
|
||||
}
|
||||
|
||||
/// List all stored keys (metadata only).
|
||||
pub async fn list_keys(&self) -> Result<Vec<KeyMetadata>, KeyError> {
|
||||
let store = self.load_store().await?;
|
||||
let mut keys: Vec<KeyMetadata> = store.keys.values().cloned().collect();
|
||||
keys.sort_by(|a, b| a.label.cmp(&b.label));
|
||||
Ok(keys)
|
||||
}
|
||||
|
||||
/// Get metadata for a specific key.
|
||||
pub async fn get_key(&self, label: &str) -> Result<KeyMetadata, KeyError> {
|
||||
let store = self.load_store().await?;
|
||||
store
|
||||
.keys
|
||||
.get(label)
|
||||
.cloned()
|
||||
.ok_or_else(|| KeyError::NotFound {
|
||||
label: label.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Remove a key (deletes from secrets store and metadata).
|
||||
pub async fn remove_key(&self, label: &str) -> Result<(), KeyError> {
|
||||
let mut store = self.load_store().await?;
|
||||
if store.keys.remove(label).is_none() {
|
||||
return Err(KeyError::NotFound {
|
||||
label: label.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Delete from secrets store
|
||||
let secret_name = format!("near_key:{}", label);
|
||||
let _ = self.secrets_store.delete(&self.user_id, &secret_name).await;
|
||||
|
||||
self.save_store(&store).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Export the public key (NEVER the private key).
|
||||
pub async fn export_public_key(&self, label: &str) -> Result<NearPublicKey, KeyError> {
|
||||
let metadata = self.get_key(label).await?;
|
||||
NearPublicKey::from_near_format(&metadata.public_key)
|
||||
}
|
||||
|
||||
// -- Transaction signing --
|
||||
|
||||
/// Sign a transaction with policy enforcement.
|
||||
pub async fn sign_transaction(
|
||||
&self,
|
||||
label: &str,
|
||||
receiver_id: &NearAccountId,
|
||||
actions: Vec<transaction::Action>,
|
||||
) -> Result<SignResult, KeyError> {
|
||||
let metadata = self.get_key(label).await?;
|
||||
|
||||
// Analyze
|
||||
let analysis = analyze_transaction(
|
||||
receiver_id.as_str(),
|
||||
&actions,
|
||||
&metadata.permission,
|
||||
&self.policy,
|
||||
);
|
||||
|
||||
// Check spend
|
||||
let daily_spend = self.spend_tracker.get_daily_spend().await?;
|
||||
|
||||
// Evaluate policy
|
||||
let decision = self
|
||||
.policy
|
||||
.evaluate(&analysis, &metadata.permission, daily_spend);
|
||||
|
||||
match decision {
|
||||
PolicyDecision::Deny { reason } => Err(KeyError::PolicyDenied { reason }),
|
||||
PolicyDecision::RequireApproval { reasons } => {
|
||||
Ok(SignResult::ApprovalRequired { analysis, reasons })
|
||||
}
|
||||
PolicyDecision::AutoApprove => {
|
||||
let signed = self
|
||||
.build_and_sign(label, &metadata, receiver_id, actions)
|
||||
.await?;
|
||||
|
||||
// Record spend
|
||||
if analysis.total_value_yocto > 0 {
|
||||
let _ = self
|
||||
.spend_tracker
|
||||
.record_spend(analysis.total_value_yocto, analysis.summary.clone(), None)
|
||||
.await;
|
||||
}
|
||||
|
||||
Ok(SignResult::Signed {
|
||||
transaction: signed,
|
||||
analysis,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Request a chain signature via MPC.
|
||||
pub async fn request_chain_signature(
|
||||
&self,
|
||||
label: &str,
|
||||
payload: &[u8],
|
||||
derivation_path: &str,
|
||||
domain: SignatureDomain,
|
||||
) -> Result<SignResult, KeyError> {
|
||||
let metadata = self.get_key(label).await?;
|
||||
|
||||
// Build chain sig analysis
|
||||
let chain_sig = ChainSigAnalysis {
|
||||
derivation_path: derivation_path.to_string(),
|
||||
domain,
|
||||
target_chain: infer_target_chain(derivation_path),
|
||||
payload_size: payload.len(),
|
||||
risk_level: policy::RiskLevel::Medium,
|
||||
};
|
||||
|
||||
let daily_spend = self.spend_tracker.get_daily_spend().await?;
|
||||
let decision = self.policy.evaluate_chain_sig(&chain_sig, daily_spend);
|
||||
|
||||
// Build the function call action
|
||||
let action =
|
||||
chain_signatures::build_chain_signature_action(payload, derivation_path, domain)?;
|
||||
|
||||
let contract = chain_signatures::chain_sig_contract(&metadata.network);
|
||||
let contract_id = NearAccountId::new(contract)?;
|
||||
|
||||
// Analyze the underlying transaction too
|
||||
let analysis = analyze_transaction(
|
||||
contract,
|
||||
&[action.clone()],
|
||||
&metadata.permission,
|
||||
&self.policy,
|
||||
);
|
||||
|
||||
match decision {
|
||||
PolicyDecision::Deny { reason } => Err(KeyError::PolicyDenied { reason }),
|
||||
PolicyDecision::RequireApproval { reasons } => {
|
||||
Ok(SignResult::ApprovalRequired { analysis, reasons })
|
||||
}
|
||||
PolicyDecision::AutoApprove => {
|
||||
let signed = self
|
||||
.build_and_sign(label, &metadata, &contract_id, vec![action])
|
||||
.await?;
|
||||
Ok(SignResult::Signed {
|
||||
transaction: signed,
|
||||
analysis,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Build and sign a transaction (internal, after policy check passes).
|
||||
async fn build_and_sign(
|
||||
&self,
|
||||
label: &str,
|
||||
metadata: &KeyMetadata,
|
||||
receiver_id: &NearAccountId,
|
||||
actions: Vec<transaction::Action>,
|
||||
) -> Result<SignedTransaction, KeyError> {
|
||||
let public_key = NearPublicKey::from_near_format(&metadata.public_key)?;
|
||||
|
||||
// Get nonce and block hash from RPC
|
||||
let rpc = NearRpcClient::new(&metadata.network);
|
||||
let access_key = rpc
|
||||
.view_access_key(&metadata.account_id, &metadata.public_key)
|
||||
.await?;
|
||||
|
||||
let nonce = access_key.nonce + 1;
|
||||
let block_hash = BlockHash::from_base58(&access_key.block_hash)?;
|
||||
|
||||
let signer_id = NearAccountId::new(&metadata.account_id)?;
|
||||
|
||||
let tx = Transaction {
|
||||
signer_id,
|
||||
public_key,
|
||||
nonce,
|
||||
receiver_id: receiver_id.clone(),
|
||||
block_hash,
|
||||
actions,
|
||||
};
|
||||
|
||||
// Hash and sign
|
||||
let hash = tx.hash_for_signing()?;
|
||||
let sig_bytes = sign_hash(self.secrets_store.as_ref(), &self.user_id, label, &hash).await?;
|
||||
|
||||
Ok(SignedTransaction {
|
||||
transaction: tx,
|
||||
signature: Signature {
|
||||
key_type: KeyType::Ed25519,
|
||||
data: sig_bytes,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// -- Backup / Restore --
|
||||
|
||||
/// Create an encrypted backup of all keys.
|
||||
pub async fn create_backup(&self, passphrase: &str) -> Result<Vec<u8>, KeyError> {
|
||||
let store = self.load_store().await?;
|
||||
|
||||
let mut entries = Vec::new();
|
||||
for (label, metadata) in &store.keys {
|
||||
let secret_name = format!("near_key:{}", label);
|
||||
let decrypted = self
|
||||
.secrets_store
|
||||
.get_decrypted(&self.user_id, &secret_name)
|
||||
.await
|
||||
.map_err(|e| KeyError::BackupError {
|
||||
reason: format!("failed to decrypt key '{}': {}", label, e),
|
||||
})?;
|
||||
|
||||
entries.push(KeyBackupEntry {
|
||||
label: label.clone(),
|
||||
account_id: metadata.account_id.clone(),
|
||||
secret_key_near_format: decrypted.expose().to_string(),
|
||||
permission: metadata.permission.clone(),
|
||||
network: metadata.network.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
let backup = KeyBackup {
|
||||
version: 1,
|
||||
created_at: Utc::now(),
|
||||
keys: entries,
|
||||
};
|
||||
|
||||
let plaintext = serde_json::to_vec(&backup).map_err(|e| KeyError::BackupError {
|
||||
reason: format!("failed to serialize backup: {}", e),
|
||||
})?;
|
||||
|
||||
encrypt_backup(passphrase, &plaintext)
|
||||
}
|
||||
|
||||
/// Restore keys from an encrypted backup.
|
||||
pub async fn restore_backup(
|
||||
&self,
|
||||
backup_data: &[u8],
|
||||
passphrase: &str,
|
||||
) -> Result<Vec<String>, KeyError> {
|
||||
let plaintext = decrypt_backup(passphrase, backup_data)?;
|
||||
|
||||
let backup: KeyBackup =
|
||||
serde_json::from_slice(&plaintext).map_err(|e| KeyError::BackupError {
|
||||
reason: format!("failed to parse backup: {}", e),
|
||||
})?;
|
||||
|
||||
let mut restored = Vec::new();
|
||||
for entry in backup.keys {
|
||||
// Validate the key
|
||||
let _ = public_key_from_secret(&entry.secret_key_near_format)?;
|
||||
let account_id = NearAccountId::new(&entry.account_id)?;
|
||||
|
||||
// Import (skip if already exists)
|
||||
match self
|
||||
.import_key(
|
||||
&entry.label,
|
||||
&account_id,
|
||||
&entry.secret_key_near_format,
|
||||
entry.permission,
|
||||
entry.network,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => restored.push(entry.label),
|
||||
Err(KeyError::AlreadyExists { .. }) => {
|
||||
// Skip existing keys
|
||||
}
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
// Update backup timestamp
|
||||
let mut store = self.load_store().await?;
|
||||
store.last_backup_at = Some(Utc::now());
|
||||
self.save_store(&store).await?;
|
||||
|
||||
Ok(restored)
|
||||
}
|
||||
|
||||
// -- Internal helpers --
|
||||
|
||||
async fn load_store(&self) -> Result<KeyStore, KeyError> {
|
||||
if !self.metadata_path.exists() {
|
||||
return Ok(KeyStore::default());
|
||||
}
|
||||
|
||||
let content = fs::read_to_string(&self.metadata_path).await?;
|
||||
serde_json::from_str(&content).map_err(|e| {
|
||||
KeyError::Io(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
format!("corrupt keys.json: {}", e),
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
async fn save_store(&self, store: &KeyStore) -> Result<(), KeyError> {
|
||||
if let Some(parent) = self.metadata_path.parent() {
|
||||
fs::create_dir_all(parent).await?;
|
||||
}
|
||||
|
||||
let content = serde_json::to_string_pretty(store).map_err(|e| {
|
||||
KeyError::SerializationFailed(format!("failed to serialize key store: {}", e))
|
||||
})?;
|
||||
|
||||
fs::write(&self.metadata_path, content).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Default path for keys metadata.
|
||||
fn default_keys_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.map(|h| h.join(".ironclaw").join("keys.json"))
|
||||
.unwrap_or_else(|| PathBuf::from(".ironclaw/keys.json"))
|
||||
}
|
||||
|
||||
// -- Backup encryption --
|
||||
|
||||
/// Backup file magic bytes.
|
||||
const BACKUP_MAGIC: &[u8; 4] = b"ICLK";
|
||||
const BACKUP_VERSION: u32 = 1;
|
||||
const ARGON2_SALT_LEN: usize = 32;
|
||||
const AES_NONCE_LEN: usize = 12;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct KeyBackup {
|
||||
version: u32,
|
||||
created_at: chrono::DateTime<Utc>,
|
||||
keys: Vec<KeyBackupEntry>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct KeyBackupEntry {
|
||||
label: String,
|
||||
account_id: String,
|
||||
secret_key_near_format: String,
|
||||
permission: AccessKeyPermission,
|
||||
network: NearNetwork,
|
||||
}
|
||||
|
||||
fn encrypt_backup(passphrase: &str, plaintext: &[u8]) -> Result<Vec<u8>, KeyError> {
|
||||
use aes_gcm::{Aes256Gcm, KeyInit, Nonce, aead::Aead};
|
||||
use argon2::Argon2;
|
||||
|
||||
// Generate salt
|
||||
let mut salt = [0u8; ARGON2_SALT_LEN];
|
||||
rand::RngCore::fill_bytes(&mut rand::rngs::OsRng, &mut salt);
|
||||
|
||||
// Derive key with Argon2id
|
||||
let mut derived_key = [0u8; 32];
|
||||
Argon2::default()
|
||||
.hash_password_into(passphrase.as_bytes(), &salt, &mut derived_key)
|
||||
.map_err(|e| KeyError::BackupError {
|
||||
reason: format!("Argon2 key derivation failed: {}", e),
|
||||
})?;
|
||||
|
||||
// Encrypt with AES-256-GCM
|
||||
let cipher = Aes256Gcm::new_from_slice(&derived_key).map_err(|e| KeyError::BackupError {
|
||||
reason: format!("failed to create cipher: {}", e),
|
||||
})?;
|
||||
|
||||
let mut nonce_bytes = [0u8; AES_NONCE_LEN];
|
||||
rand::RngCore::fill_bytes(&mut rand::rngs::OsRng, &mut nonce_bytes);
|
||||
let nonce = Nonce::from_slice(&nonce_bytes);
|
||||
|
||||
let ciphertext = cipher
|
||||
.encrypt(nonce, plaintext)
|
||||
.map_err(|e| KeyError::BackupError {
|
||||
reason: format!("encryption failed: {}", e),
|
||||
})?;
|
||||
|
||||
// Assemble: magic + version + salt + nonce + ciphertext
|
||||
let mut output = Vec::new();
|
||||
output.extend_from_slice(BACKUP_MAGIC);
|
||||
output.extend_from_slice(&BACKUP_VERSION.to_le_bytes());
|
||||
output.extend_from_slice(&salt);
|
||||
output.extend_from_slice(&nonce_bytes);
|
||||
output.extend_from_slice(&ciphertext);
|
||||
|
||||
derived_key.zeroize();
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
pub(crate) fn decrypt_backup(passphrase: &str, data: &[u8]) -> Result<Vec<u8>, KeyError> {
|
||||
use aes_gcm::{Aes256Gcm, KeyInit, Nonce, aead::Aead};
|
||||
use argon2::Argon2;
|
||||
|
||||
let header_len = 4 + 4 + ARGON2_SALT_LEN + AES_NONCE_LEN;
|
||||
if data.len() < header_len {
|
||||
return Err(KeyError::BackupError {
|
||||
reason: "backup file too short".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check magic
|
||||
if &data[..4] != BACKUP_MAGIC {
|
||||
return Err(KeyError::BackupError {
|
||||
reason: "not a valid IronClaw backup file".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check version
|
||||
let version = u32::from_le_bytes(data[4..8].try_into().unwrap());
|
||||
if version != BACKUP_VERSION {
|
||||
return Err(KeyError::BackupError {
|
||||
reason: format!("unsupported backup version: {}", version),
|
||||
});
|
||||
}
|
||||
|
||||
let salt = &data[8..8 + ARGON2_SALT_LEN];
|
||||
let nonce_bytes = &data[8 + ARGON2_SALT_LEN..header_len];
|
||||
let ciphertext = &data[header_len..];
|
||||
|
||||
// Derive key
|
||||
let mut derived_key = [0u8; 32];
|
||||
Argon2::default()
|
||||
.hash_password_into(passphrase.as_bytes(), salt, &mut derived_key)
|
||||
.map_err(|e| KeyError::BackupError {
|
||||
reason: format!("Argon2 key derivation failed: {}", e),
|
||||
})?;
|
||||
|
||||
// Decrypt
|
||||
let cipher = Aes256Gcm::new_from_slice(&derived_key).map_err(|e| KeyError::BackupError {
|
||||
reason: format!("failed to create cipher: {}", e),
|
||||
})?;
|
||||
|
||||
let nonce = Nonce::from_slice(nonce_bytes);
|
||||
let plaintext = cipher
|
||||
.decrypt(nonce, ciphertext)
|
||||
.map_err(|_| KeyError::BackupError {
|
||||
reason: "decryption failed (wrong passphrase?)".to_string(),
|
||||
})?;
|
||||
|
||||
derived_key.zeroize();
|
||||
|
||||
Ok(plaintext)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use secrecy::SecretString;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use crate::keys::spending::SpendTracker;
|
||||
use crate::keys::transaction::{Action, ONE_NEAR, Transfer};
|
||||
use crate::keys::types::{AccessKeyPermission, NearAccountId, NearNetwork};
|
||||
use crate::keys::{KeyManager, SignResult};
|
||||
use crate::secrets::{InMemorySecretsStore, SecretsCrypto};
|
||||
|
||||
fn test_manager(dir: &TempDir) -> KeyManager {
|
||||
let key = "0123456789abcdef0123456789abcdef";
|
||||
let crypto = Arc::new(SecretsCrypto::new(SecretString::from(key.to_string())).unwrap());
|
||||
let store: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
|
||||
Arc::new(InMemorySecretsStore::new(crypto));
|
||||
|
||||
KeyManager::new(store, "test_user".to_string())
|
||||
.with_metadata_path(dir.path().join("keys.json"))
|
||||
.with_spend_tracker(SpendTracker::new(dir.path().join("spend.json")))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_generate_key() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
let metadata = manager
|
||||
.generate_key(
|
||||
"test-key",
|
||||
&account,
|
||||
AccessKeyPermission::FunctionCall {
|
||||
allowance: None,
|
||||
receiver_id: "intents.near".to_string(),
|
||||
method_names: vec![],
|
||||
},
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(metadata.label, "test-key");
|
||||
assert_eq!(metadata.account_id, "alice.testnet");
|
||||
assert!(metadata.public_key.starts_with("ed25519:"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_generate_duplicate_key_fails() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
|
||||
manager
|
||||
.generate_key(
|
||||
"dup",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = manager
|
||||
.generate_key(
|
||||
"dup",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(crate::keys::KeyError::AlreadyExists { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_list_keys() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
|
||||
assert_eq!(manager.list_keys().await.unwrap().len(), 0);
|
||||
|
||||
manager
|
||||
.generate_key(
|
||||
"key-1",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
manager
|
||||
.generate_key(
|
||||
"key-2",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let keys = manager.list_keys().await.unwrap();
|
||||
assert_eq!(keys.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_remove_key() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
|
||||
manager
|
||||
.generate_key(
|
||||
"to-remove",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
manager.remove_key("to-remove").await.unwrap();
|
||||
assert!(manager.get_key("to-remove").await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_export_public_key() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
|
||||
let metadata = manager
|
||||
.generate_key(
|
||||
"export-test",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let pubkey = manager.export_public_key("export-test").await.unwrap();
|
||||
assert_eq!(pubkey.to_near_format(), metadata.public_key);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_import_key() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
let account = NearAccountId::new("bob.testnet").unwrap();
|
||||
|
||||
// Generate a test secret key
|
||||
let signing_key = ed25519_dalek::SigningKey::generate(&mut rand::rngs::OsRng);
|
||||
let verifying_key = signing_key.verifying_key();
|
||||
let mut combined = Vec::with_capacity(64);
|
||||
combined.extend_from_slice(signing_key.as_bytes());
|
||||
combined.extend_from_slice(verifying_key.as_bytes());
|
||||
let secret = format!("ed25519:{}", bs58::encode(&combined).into_string());
|
||||
|
||||
let metadata = manager
|
||||
.import_key(
|
||||
"imported",
|
||||
&account,
|
||||
&secret,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(metadata.label, "imported");
|
||||
assert!(metadata.public_key.starts_with("ed25519:"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_backup_and_restore_roundtrip() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
|
||||
// Generate a key
|
||||
manager
|
||||
.generate_key(
|
||||
"backup-test",
|
||||
&account,
|
||||
AccessKeyPermission::FunctionCall {
|
||||
allowance: None,
|
||||
receiver_id: "contract.near".to_string(),
|
||||
method_names: vec!["deposit".to_string()],
|
||||
},
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Create backup
|
||||
let backup_data = manager.create_backup("test-passphrase").await.unwrap();
|
||||
assert!(!backup_data.is_empty());
|
||||
|
||||
// Restore into a fresh manager
|
||||
let dir2 = TempDir::new().unwrap();
|
||||
let manager2 = test_manager(&dir2);
|
||||
|
||||
let restored = manager2
|
||||
.restore_backup(&backup_data, "test-passphrase")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(restored, vec!["backup-test"]);
|
||||
|
||||
// Verify the restored key
|
||||
let keys = manager2.list_keys().await.unwrap();
|
||||
assert_eq!(keys.len(), 1);
|
||||
assert_eq!(keys[0].label, "backup-test");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_backup_wrong_passphrase() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
|
||||
manager
|
||||
.generate_key(
|
||||
"test",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let backup_data = manager.create_backup("correct").await.unwrap();
|
||||
|
||||
let dir2 = TempDir::new().unwrap();
|
||||
let manager2 = test_manager(&dir2);
|
||||
|
||||
let result = manager2.restore_backup(&backup_data, "wrong").await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sign_transaction_policy_deny() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let key = "0123456789abcdef0123456789abcdef";
|
||||
let crypto = Arc::new(SecretsCrypto::new(SecretString::from(key.to_string())).unwrap());
|
||||
let store: Arc<dyn crate::secrets::SecretsStore + Send + Sync> =
|
||||
Arc::new(InMemorySecretsStore::new(crypto));
|
||||
|
||||
let mut manager = KeyManager::new(store, "test_user".to_string())
|
||||
.with_metadata_path(dir.path().join("keys.json"))
|
||||
.with_spend_tracker(SpendTracker::new(dir.path().join("spend.json")));
|
||||
|
||||
// Deny full access operations
|
||||
manager.policy_mut().deny_full_access_operations = true;
|
||||
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
manager
|
||||
.generate_key(
|
||||
"denied",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let receiver = NearAccountId::new("bob.testnet").unwrap();
|
||||
let result = manager
|
||||
.sign_transaction(
|
||||
"denied",
|
||||
&receiver,
|
||||
vec![Action::Transfer(Transfer { deposit: 0 })],
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(crate::keys::KeyError::PolicyDenied { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sign_transaction_requires_approval() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let manager = test_manager(&dir);
|
||||
|
||||
let account = NearAccountId::new("alice.testnet").unwrap();
|
||||
manager
|
||||
.generate_key(
|
||||
"signer",
|
||||
&account,
|
||||
AccessKeyPermission::FullAccess,
|
||||
NearNetwork::Testnet,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let receiver = NearAccountId::new("unknown.testnet").unwrap();
|
||||
let result = manager
|
||||
.sign_transaction(
|
||||
"signer",
|
||||
&receiver,
|
||||
vec![Action::Transfer(Transfer {
|
||||
deposit: 100 * ONE_NEAR,
|
||||
})],
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Default policy requires approval for any transfer
|
||||
assert!(matches!(result, SignResult::ApprovalRequired { .. }));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,917 @@
|
||||
//! Transaction analysis and policy engine for NEAR key operations.
|
||||
//!
|
||||
//! Every transaction is decomposed into a `TransactionAnalysis` before any
|
||||
//! signing happens. The policy engine then evaluates the analysis against
|
||||
//! a configurable ruleset. Most restrictive rule always wins.
|
||||
//!
|
||||
//! # Pipeline
|
||||
//!
|
||||
//! ```text
|
||||
//! Transaction -> analyze_transaction() -> TransactionAnalysis
|
||||
//! |
|
||||
//! PolicyConfig.evaluate() <-------+
|
||||
//! |
|
||||
//! PolicyDecision { AutoApprove | RequireApproval | Deny }
|
||||
//! ```
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::keys::transaction::{Action, ONE_NEAR};
|
||||
use crate::keys::types::{AccessKeyPermission, format_yocto};
|
||||
|
||||
/// Risk level for a single action within a transaction.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
|
||||
pub enum RiskLevel {
|
||||
Low,
|
||||
Medium,
|
||||
High,
|
||||
Critical,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for RiskLevel {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
RiskLevel::Low => write!(f, "LOW"),
|
||||
RiskLevel::Medium => write!(f, "MEDIUM"),
|
||||
RiskLevel::High => write!(f, "HIGH"),
|
||||
RiskLevel::Critical => write!(f, "CRITICAL"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Category of a transaction action for policy evaluation.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum ActionCategory {
|
||||
Transfer,
|
||||
FunctionCall,
|
||||
Stake,
|
||||
AddKey { is_full_access: bool },
|
||||
DeleteKey,
|
||||
DeployContract,
|
||||
CreateAccount,
|
||||
DeleteAccount,
|
||||
}
|
||||
|
||||
/// Analysis of a single action within a transaction.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ActionAnalysis {
|
||||
pub category: ActionCategory,
|
||||
pub value_yocto: u128,
|
||||
pub receiver: String,
|
||||
pub method: Option<String>,
|
||||
pub description: String,
|
||||
pub risk_level: RiskLevel,
|
||||
}
|
||||
|
||||
/// Complete analysis of a transaction.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TransactionAnalysis {
|
||||
pub actions: Vec<ActionAnalysis>,
|
||||
pub total_value_yocto: u128,
|
||||
pub receivers: Vec<String>,
|
||||
pub uses_full_access_key: bool,
|
||||
pub summary: String,
|
||||
}
|
||||
|
||||
/// Analyze a transaction's actions for policy evaluation.
|
||||
pub fn analyze_transaction(
|
||||
receiver_id: &str,
|
||||
actions: &[Action],
|
||||
key_permission: &AccessKeyPermission,
|
||||
policy: &PolicyConfig,
|
||||
) -> TransactionAnalysis {
|
||||
let uses_full_access_key = matches!(key_permission, AccessKeyPermission::FullAccess);
|
||||
let mut action_analyses = Vec::new();
|
||||
let mut total_value = 0u128;
|
||||
|
||||
for action in actions {
|
||||
let analysis = analyze_action(action, receiver_id, policy);
|
||||
total_value = total_value.saturating_add(analysis.value_yocto);
|
||||
action_analyses.push(analysis);
|
||||
}
|
||||
|
||||
let receivers: Vec<String> = action_analyses
|
||||
.iter()
|
||||
.map(|a| a.receiver.clone())
|
||||
.collect::<std::collections::HashSet<_>>()
|
||||
.into_iter()
|
||||
.collect();
|
||||
|
||||
let summary = build_summary(&action_analyses, total_value);
|
||||
|
||||
TransactionAnalysis {
|
||||
actions: action_analyses,
|
||||
total_value_yocto: total_value,
|
||||
receivers,
|
||||
uses_full_access_key,
|
||||
summary,
|
||||
}
|
||||
}
|
||||
|
||||
fn analyze_action(action: &Action, receiver_id: &str, policy: &PolicyConfig) -> ActionAnalysis {
|
||||
match action {
|
||||
Action::Transfer(t) => {
|
||||
let is_whitelisted = policy.transfer_whitelist.contains(&receiver_id.to_string());
|
||||
let risk = if t.deposit == 0 || (t.deposit < ONE_NEAR && is_whitelisted) {
|
||||
RiskLevel::Low
|
||||
} else if t.deposit < policy.transfer_whitelist_max_yocto && is_whitelisted {
|
||||
RiskLevel::Medium
|
||||
} else {
|
||||
RiskLevel::High
|
||||
};
|
||||
|
||||
ActionAnalysis {
|
||||
category: ActionCategory::Transfer,
|
||||
value_yocto: t.deposit,
|
||||
receiver: receiver_id.to_string(),
|
||||
method: None,
|
||||
description: format!("Transfer {} to {}", format_yocto(t.deposit), receiver_id),
|
||||
risk_level: risk,
|
||||
}
|
||||
}
|
||||
|
||||
Action::FunctionCall(fc) => {
|
||||
let has_matching_rule = policy
|
||||
.function_call_rules
|
||||
.iter()
|
||||
.any(|r| r.receiver_id == receiver_id && fc.deposit <= r.max_deposit_yocto);
|
||||
|
||||
let risk = if fc.deposit == 0 && has_matching_rule {
|
||||
RiskLevel::Low
|
||||
} else if fc.deposit == 0 || has_matching_rule {
|
||||
RiskLevel::Medium
|
||||
} else {
|
||||
RiskLevel::High
|
||||
};
|
||||
|
||||
ActionAnalysis {
|
||||
category: ActionCategory::FunctionCall,
|
||||
value_yocto: fc.deposit,
|
||||
receiver: receiver_id.to_string(),
|
||||
method: Some(fc.method_name.clone()),
|
||||
description: format!(
|
||||
"FunctionCall {}::{}{}",
|
||||
receiver_id,
|
||||
fc.method_name,
|
||||
if fc.deposit > 0 {
|
||||
format!(" ({})", format_yocto(fc.deposit))
|
||||
} else {
|
||||
String::new()
|
||||
}
|
||||
),
|
||||
risk_level: risk,
|
||||
}
|
||||
}
|
||||
|
||||
Action::Stake(s) => {
|
||||
let risk = if policy
|
||||
.stake_validator_whitelist
|
||||
.contains(&receiver_id.to_string())
|
||||
&& s.stake <= policy.stake_auto_approve_max_yocto
|
||||
{
|
||||
RiskLevel::Medium
|
||||
} else {
|
||||
RiskLevel::High
|
||||
};
|
||||
|
||||
ActionAnalysis {
|
||||
category: ActionCategory::Stake,
|
||||
value_yocto: s.stake,
|
||||
receiver: receiver_id.to_string(),
|
||||
method: None,
|
||||
description: format!("Stake {} with {}", format_yocto(s.stake), receiver_id),
|
||||
risk_level: risk,
|
||||
}
|
||||
}
|
||||
|
||||
Action::AddKey(ak) => {
|
||||
let is_full_access = borsh_permission_is_full_access(&ak.access_key.permission);
|
||||
ActionAnalysis {
|
||||
category: ActionCategory::AddKey { is_full_access },
|
||||
value_yocto: 0,
|
||||
receiver: receiver_id.to_string(),
|
||||
method: None,
|
||||
description: if is_full_access {
|
||||
format!("AddKey (FullAccess) to {}", receiver_id)
|
||||
} else {
|
||||
format!("AddKey (FunctionCall) to {}", receiver_id)
|
||||
},
|
||||
risk_level: if is_full_access {
|
||||
RiskLevel::Critical
|
||||
} else {
|
||||
RiskLevel::High
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
Action::DeleteKey(_) => ActionAnalysis {
|
||||
category: ActionCategory::DeleteKey,
|
||||
value_yocto: 0,
|
||||
receiver: receiver_id.to_string(),
|
||||
method: None,
|
||||
description: format!("DeleteKey on {}", receiver_id),
|
||||
risk_level: RiskLevel::High,
|
||||
},
|
||||
|
||||
Action::DeployContract(_) => ActionAnalysis {
|
||||
category: ActionCategory::DeployContract,
|
||||
value_yocto: 0,
|
||||
receiver: receiver_id.to_string(),
|
||||
method: None,
|
||||
description: format!("DeployContract to {}", receiver_id),
|
||||
risk_level: RiskLevel::Critical,
|
||||
},
|
||||
|
||||
Action::CreateAccount => ActionAnalysis {
|
||||
category: ActionCategory::CreateAccount,
|
||||
value_yocto: 0,
|
||||
receiver: receiver_id.to_string(),
|
||||
method: None,
|
||||
description: format!("CreateAccount {}", receiver_id),
|
||||
risk_level: RiskLevel::Medium,
|
||||
},
|
||||
|
||||
Action::DeleteAccount(_) => ActionAnalysis {
|
||||
category: ActionCategory::DeleteAccount,
|
||||
value_yocto: 0,
|
||||
receiver: receiver_id.to_string(),
|
||||
method: None,
|
||||
description: format!("DeleteAccount {}", receiver_id),
|
||||
risk_level: RiskLevel::Critical,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn borsh_permission_is_full_access(
|
||||
perm: &crate::keys::transaction::AccessKeyPermissionBorsh,
|
||||
) -> bool {
|
||||
matches!(
|
||||
perm,
|
||||
crate::keys::transaction::AccessKeyPermissionBorsh::FullAccess
|
||||
)
|
||||
}
|
||||
|
||||
fn build_summary(actions: &[ActionAnalysis], total_value: u128) -> String {
|
||||
let mut lines = Vec::new();
|
||||
for (i, a) in actions.iter().enumerate() {
|
||||
lines.push(format!(" {}. {} [{}]", i + 1, a.description, a.risk_level));
|
||||
}
|
||||
if total_value > 0 {
|
||||
lines.push(format!(" Total value: {}", format_yocto(total_value)));
|
||||
}
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
/// Policy decision after evaluating a transaction.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PolicyDecision {
|
||||
/// Transaction can proceed without user interaction.
|
||||
AutoApprove,
|
||||
/// User must approve before signing.
|
||||
RequireApproval { reasons: Vec<String> },
|
||||
/// Transaction is denied by policy (not even user can override).
|
||||
Deny { reason: String },
|
||||
}
|
||||
|
||||
/// Configurable policy rules for transaction approval.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PolicyConfig {
|
||||
// Transfer rules
|
||||
pub transfer_auto_approve_max_yocto: u128,
|
||||
pub transfer_whitelist_max_yocto: u128,
|
||||
pub transfer_whitelist: Vec<String>,
|
||||
|
||||
// Function call rules
|
||||
pub function_call_rules: Vec<FunctionCallRule>,
|
||||
|
||||
// Staking rules
|
||||
pub stake_validator_whitelist: Vec<String>,
|
||||
pub stake_auto_approve_max_yocto: u128,
|
||||
|
||||
// Key management rules
|
||||
pub allow_add_scoped_keys_to: Vec<String>,
|
||||
|
||||
// Chain signature rules
|
||||
pub chain_sig_rules: Vec<ChainSigRule>,
|
||||
|
||||
// Global limits
|
||||
pub daily_spend_limit_yocto: Option<u128>,
|
||||
pub per_tx_auto_approve_max_yocto: u128,
|
||||
|
||||
// Blanket denials
|
||||
pub deny_full_access_operations: bool,
|
||||
pub deny_delete_account: bool,
|
||||
}
|
||||
|
||||
impl Default for PolicyConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
transfer_auto_approve_max_yocto: 0,
|
||||
transfer_whitelist_max_yocto: ONE_NEAR,
|
||||
transfer_whitelist: Vec::new(),
|
||||
function_call_rules: Vec::new(),
|
||||
stake_validator_whitelist: Vec::new(),
|
||||
stake_auto_approve_max_yocto: 0,
|
||||
allow_add_scoped_keys_to: Vec::new(),
|
||||
chain_sig_rules: Vec::new(),
|
||||
daily_spend_limit_yocto: None,
|
||||
per_tx_auto_approve_max_yocto: 0,
|
||||
deny_full_access_operations: false,
|
||||
deny_delete_account: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A function call rule for policy evaluation.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct FunctionCallRule {
|
||||
pub receiver_id: String,
|
||||
/// Empty = all methods on this contract.
|
||||
pub allowed_methods: Vec<String>,
|
||||
pub max_deposit_yocto: u128,
|
||||
pub max_gas: Option<u64>,
|
||||
pub auto_approve: bool,
|
||||
}
|
||||
|
||||
/// Signature domain for chain signatures.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum SignatureDomain {
|
||||
Secp256k1 = 0,
|
||||
Ed25519 = 1,
|
||||
}
|
||||
|
||||
/// A chain signature rule for policy evaluation.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ChainSigRule {
|
||||
pub allowed_paths: Vec<String>,
|
||||
pub allowed_domains: Vec<SignatureDomain>,
|
||||
pub max_payload_bytes: usize,
|
||||
pub auto_approve: bool,
|
||||
}
|
||||
|
||||
/// Analysis specific to chain signature requests.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ChainSigAnalysis {
|
||||
pub derivation_path: String,
|
||||
pub domain: SignatureDomain,
|
||||
pub target_chain: Option<String>,
|
||||
pub payload_size: usize,
|
||||
pub risk_level: RiskLevel,
|
||||
}
|
||||
|
||||
impl PolicyConfig {
|
||||
/// Evaluate a transaction analysis against this policy.
|
||||
///
|
||||
/// Returns the most restrictive decision across all actions.
|
||||
pub fn evaluate(
|
||||
&self,
|
||||
analysis: &TransactionAnalysis,
|
||||
key_permission: &AccessKeyPermission,
|
||||
daily_spend: u128,
|
||||
) -> PolicyDecision {
|
||||
let mut reasons = Vec::new();
|
||||
|
||||
// Blanket denials first
|
||||
if self.deny_full_access_operations && analysis.uses_full_access_key {
|
||||
return PolicyDecision::Deny {
|
||||
reason: "full-access key operations are denied by policy".to_string(),
|
||||
};
|
||||
}
|
||||
|
||||
for action in &analysis.actions {
|
||||
if self.deny_delete_account && matches!(action.category, ActionCategory::DeleteAccount)
|
||||
{
|
||||
return PolicyDecision::Deny {
|
||||
reason: "account deletion is denied by policy".to_string(),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Daily spend limit
|
||||
if let Some(limit) = self.daily_spend_limit_yocto {
|
||||
if daily_spend.saturating_add(analysis.total_value_yocto) > limit {
|
||||
reasons.push(format!(
|
||||
"daily spend limit exceeded: {} + {} > {}",
|
||||
format_yocto(daily_spend),
|
||||
format_yocto(analysis.total_value_yocto),
|
||||
format_yocto(limit)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// Per-transaction limit
|
||||
if analysis.total_value_yocto > self.per_tx_auto_approve_max_yocto
|
||||
&& self.per_tx_auto_approve_max_yocto > 0
|
||||
{
|
||||
reasons.push(format!(
|
||||
"transaction value {} exceeds per-tx auto-approve limit {}",
|
||||
format_yocto(analysis.total_value_yocto),
|
||||
format_yocto(self.per_tx_auto_approve_max_yocto)
|
||||
));
|
||||
}
|
||||
|
||||
// Per-action evaluation
|
||||
for action in &analysis.actions {
|
||||
if let Some(reason) = self.evaluate_action(action, key_permission) {
|
||||
reasons.push(reason);
|
||||
}
|
||||
}
|
||||
|
||||
if reasons.is_empty() {
|
||||
PolicyDecision::AutoApprove
|
||||
} else {
|
||||
PolicyDecision::RequireApproval { reasons }
|
||||
}
|
||||
}
|
||||
|
||||
/// Evaluate a chain signature request.
|
||||
pub fn evaluate_chain_sig(
|
||||
&self,
|
||||
chain_sig: &ChainSigAnalysis,
|
||||
daily_spend: u128,
|
||||
) -> PolicyDecision {
|
||||
let mut reasons = Vec::new();
|
||||
|
||||
// Check daily limit (chain sigs don't have a value, but check anyway)
|
||||
if let Some(limit) = self.daily_spend_limit_yocto {
|
||||
if daily_spend > limit {
|
||||
reasons.push("daily spend limit exceeded".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
// Find matching chain sig rule
|
||||
let matching_rule = self.chain_sig_rules.iter().find(|rule| {
|
||||
rule.allowed_domains.contains(&chain_sig.domain)
|
||||
&& chain_sig.payload_size <= rule.max_payload_bytes
|
||||
&& rule
|
||||
.allowed_paths
|
||||
.iter()
|
||||
.any(|pattern| glob_matches(pattern, &chain_sig.derivation_path))
|
||||
});
|
||||
|
||||
match matching_rule {
|
||||
Some(rule) if rule.auto_approve => PolicyDecision::AutoApprove,
|
||||
Some(_) => {
|
||||
reasons.push(format!(
|
||||
"chain signature for path '{}' requires approval",
|
||||
chain_sig.derivation_path
|
||||
));
|
||||
PolicyDecision::RequireApproval { reasons }
|
||||
}
|
||||
None => {
|
||||
reasons.push(format!(
|
||||
"no matching chain signature rule for path '{}'",
|
||||
chain_sig.derivation_path
|
||||
));
|
||||
PolicyDecision::RequireApproval { reasons }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn evaluate_action(
|
||||
&self,
|
||||
action: &ActionAnalysis,
|
||||
key_permission: &AccessKeyPermission,
|
||||
) -> Option<String> {
|
||||
match &action.category {
|
||||
ActionCategory::Transfer => {
|
||||
// Auto-approve to whitelisted accounts under threshold
|
||||
if self.transfer_whitelist.contains(&action.receiver)
|
||||
&& action.value_yocto <= self.transfer_whitelist_max_yocto
|
||||
{
|
||||
return None;
|
||||
}
|
||||
// Auto-approve small transfers to anyone
|
||||
if action.value_yocto <= self.transfer_auto_approve_max_yocto {
|
||||
return None;
|
||||
}
|
||||
Some(format!(
|
||||
"transfer {} to {} exceeds auto-approve threshold",
|
||||
format_yocto(action.value_yocto),
|
||||
action.receiver
|
||||
))
|
||||
}
|
||||
|
||||
ActionCategory::FunctionCall => {
|
||||
// Check if key is already scoped to this receiver with zero deposit
|
||||
if let AccessKeyPermission::FunctionCall {
|
||||
receiver_id,
|
||||
method_names,
|
||||
..
|
||||
} = key_permission
|
||||
{
|
||||
if receiver_id == &action.receiver
|
||||
&& action.value_yocto == 0
|
||||
&& (method_names.is_empty()
|
||||
|| action
|
||||
.method
|
||||
.as_ref()
|
||||
.map(|m| method_names.contains(m))
|
||||
.unwrap_or(false))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
}
|
||||
|
||||
// Check function call rules
|
||||
if let Some(method) = &action.method {
|
||||
for rule in &self.function_call_rules {
|
||||
if rule.receiver_id == action.receiver
|
||||
&& (rule.allowed_methods.is_empty()
|
||||
|| rule.allowed_methods.contains(method))
|
||||
&& action.value_yocto <= rule.max_deposit_yocto
|
||||
&& rule.auto_approve
|
||||
{
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some(format!(
|
||||
"function call {} requires approval",
|
||||
action.description
|
||||
))
|
||||
}
|
||||
|
||||
ActionCategory::Stake => {
|
||||
if self.stake_validator_whitelist.contains(&action.receiver)
|
||||
&& action.value_yocto <= self.stake_auto_approve_max_yocto
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(format!("stake {} requires approval", action.description))
|
||||
}
|
||||
|
||||
ActionCategory::AddKey { is_full_access } => {
|
||||
if *is_full_access {
|
||||
Some("adding full-access key requires approval".to_string())
|
||||
} else {
|
||||
Some("adding function-call key requires approval".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
ActionCategory::DeleteKey
|
||||
| ActionCategory::DeployContract
|
||||
| ActionCategory::CreateAccount
|
||||
| ActionCategory::DeleteAccount => {
|
||||
Some(format!("{} requires approval", action.description))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple glob matching: supports `*` as wildcard for any suffix.
|
||||
fn glob_matches(pattern: &str, value: &str) -> bool {
|
||||
if let Some(prefix) = pattern.strip_suffix('*') {
|
||||
value.starts_with(prefix)
|
||||
} else {
|
||||
pattern == value
|
||||
}
|
||||
}
|
||||
|
||||
/// Infer target chain from a derivation path.
|
||||
pub fn infer_target_chain(derivation_path: &str) -> Option<String> {
|
||||
let lower = derivation_path.to_lowercase();
|
||||
if lower.starts_with("ethereum") || lower.starts_with("eth") {
|
||||
Some("Ethereum".to_string())
|
||||
} else if lower.starts_with("bitcoin") || lower.starts_with("btc") {
|
||||
Some("Bitcoin".to_string())
|
||||
} else if lower.starts_with("near") {
|
||||
Some("NEAR".to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::keys::policy::{
|
||||
ChainSigAnalysis, ChainSigRule, FunctionCallRule, PolicyConfig, PolicyDecision, RiskLevel,
|
||||
SignatureDomain, analyze_transaction, glob_matches, infer_target_chain,
|
||||
};
|
||||
use crate::keys::transaction::{Action, FunctionCall, ONE_NEAR, TGAS, Transfer};
|
||||
use crate::keys::types::AccessKeyPermission;
|
||||
|
||||
fn default_policy() -> PolicyConfig {
|
||||
PolicyConfig::default()
|
||||
}
|
||||
|
||||
fn permissive_policy() -> PolicyConfig {
|
||||
PolicyConfig {
|
||||
transfer_auto_approve_max_yocto: ONE_NEAR,
|
||||
transfer_whitelist_max_yocto: 10 * ONE_NEAR,
|
||||
transfer_whitelist: vec!["bob.near".to_string()],
|
||||
function_call_rules: vec![FunctionCallRule {
|
||||
receiver_id: "intents.near".to_string(),
|
||||
allowed_methods: vec!["execute_intents".to_string()],
|
||||
max_deposit_yocto: 0,
|
||||
max_gas: None,
|
||||
auto_approve: true,
|
||||
}],
|
||||
per_tx_auto_approve_max_yocto: 5 * ONE_NEAR,
|
||||
daily_spend_limit_yocto: Some(50 * ONE_NEAR),
|
||||
..default_policy()
|
||||
}
|
||||
}
|
||||
|
||||
// -- Transfer tests --
|
||||
|
||||
#[test]
|
||||
fn test_transfer_below_auto_approve() {
|
||||
let policy = permissive_policy();
|
||||
let actions = vec![Action::Transfer(Transfer {
|
||||
deposit: ONE_NEAR / 2,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("someone.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
assert_eq!(decision, PolicyDecision::AutoApprove);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transfer_above_threshold_requires_approval() {
|
||||
let policy = permissive_policy();
|
||||
let actions = vec![Action::Transfer(Transfer {
|
||||
deposit: 2 * ONE_NEAR,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("unknown.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transfer_to_whitelisted_account() {
|
||||
let policy = permissive_policy();
|
||||
let actions = vec![Action::Transfer(Transfer {
|
||||
deposit: 5 * ONE_NEAR,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("bob.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
assert_eq!(decision, PolicyDecision::AutoApprove);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transfer_to_whitelisted_above_whitelist_limit() {
|
||||
let policy = permissive_policy();
|
||||
let actions = vec![Action::Transfer(Transfer {
|
||||
deposit: 15 * ONE_NEAR,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("bob.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
// 15 NEAR > whitelist max (10 NEAR), and > per_tx limit (5 NEAR)
|
||||
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
|
||||
}
|
||||
|
||||
// -- Function call tests --
|
||||
|
||||
#[test]
|
||||
fn test_function_call_matching_rule_auto_approve() {
|
||||
let policy = permissive_policy();
|
||||
let actions = vec![Action::FunctionCall(FunctionCall {
|
||||
method_name: "execute_intents".to_string(),
|
||||
args: vec![],
|
||||
gas: 30 * TGAS,
|
||||
deposit: 0,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("intents.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
assert_eq!(decision, PolicyDecision::AutoApprove);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_function_call_scoped_key_auto_approve() {
|
||||
let policy = default_policy();
|
||||
let actions = vec![Action::FunctionCall(FunctionCall {
|
||||
method_name: "deposit".to_string(),
|
||||
args: vec![],
|
||||
gas: 30 * TGAS,
|
||||
deposit: 0,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FunctionCall {
|
||||
allowance: None,
|
||||
receiver_id: "contract.near".to_string(),
|
||||
method_names: vec!["deposit".to_string()],
|
||||
};
|
||||
let analysis = analyze_transaction("contract.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
assert_eq!(decision, PolicyDecision::AutoApprove);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_function_call_no_rule_requires_approval() {
|
||||
let policy = default_policy();
|
||||
let actions = vec![Action::FunctionCall(FunctionCall {
|
||||
method_name: "dangerous_method".to_string(),
|
||||
args: vec![],
|
||||
gas: 30 * TGAS,
|
||||
deposit: ONE_NEAR,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("unknown.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
|
||||
}
|
||||
|
||||
// -- Blanket denial tests --
|
||||
|
||||
#[test]
|
||||
fn test_deny_full_access_operations() {
|
||||
let policy = PolicyConfig {
|
||||
deny_full_access_operations: true,
|
||||
..default_policy()
|
||||
};
|
||||
let actions = vec![Action::Transfer(Transfer { deposit: 0 })];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("bob.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
assert!(matches!(decision, PolicyDecision::Deny { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deny_delete_account() {
|
||||
let policy = PolicyConfig {
|
||||
deny_delete_account: true,
|
||||
..default_policy()
|
||||
};
|
||||
let actions = vec![Action::DeleteAccount(
|
||||
crate::keys::transaction::DeleteAccount {
|
||||
beneficiary_id: crate::keys::types::NearAccountId::new("bob.near").unwrap(),
|
||||
},
|
||||
)];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("alice.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
assert!(matches!(decision, PolicyDecision::Deny { .. }));
|
||||
}
|
||||
|
||||
// -- Daily spend limit tests --
|
||||
|
||||
#[test]
|
||||
fn test_daily_spend_limit_under() {
|
||||
let policy = permissive_policy();
|
||||
let actions = vec![Action::Transfer(Transfer {
|
||||
deposit: ONE_NEAR / 2,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("bob.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 10 * ONE_NEAR);
|
||||
assert_eq!(decision, PolicyDecision::AutoApprove);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_daily_spend_limit_exceeded() {
|
||||
let policy = permissive_policy();
|
||||
let actions = vec![Action::Transfer(Transfer {
|
||||
deposit: ONE_NEAR / 2,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("bob.near", &actions, &perm, &policy);
|
||||
// Current daily spend is 50 NEAR (at limit), adding 0.5 NEAR puts us over
|
||||
let decision = policy.evaluate(&analysis, &perm, 50 * ONE_NEAR);
|
||||
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
|
||||
}
|
||||
|
||||
// -- Per-transaction limit tests --
|
||||
|
||||
#[test]
|
||||
fn test_per_tx_limit() {
|
||||
let policy = permissive_policy();
|
||||
let actions = vec![Action::Transfer(Transfer {
|
||||
deposit: 6 * ONE_NEAR,
|
||||
})];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("bob.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
// 6 NEAR > per_tx_auto_approve_max (5 NEAR)
|
||||
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
|
||||
}
|
||||
|
||||
// -- Most restrictive wins --
|
||||
|
||||
#[test]
|
||||
fn test_mixed_actions_most_restrictive_wins() {
|
||||
let policy = permissive_policy();
|
||||
// One auto-approvable + one that requires approval
|
||||
let actions = vec![
|
||||
Action::FunctionCall(FunctionCall {
|
||||
method_name: "execute_intents".to_string(),
|
||||
args: vec![],
|
||||
gas: 30 * TGAS,
|
||||
deposit: 0,
|
||||
}),
|
||||
Action::Transfer(Transfer {
|
||||
deposit: 100 * ONE_NEAR,
|
||||
}),
|
||||
];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("intents.near", &actions, &perm, &policy);
|
||||
let decision = policy.evaluate(&analysis, &perm, 0);
|
||||
// Transfer is too large, so the whole tx requires approval
|
||||
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
|
||||
}
|
||||
|
||||
// -- Transaction analysis tests --
|
||||
|
||||
#[test]
|
||||
fn test_analysis_total_value() {
|
||||
let policy = default_policy();
|
||||
let actions = vec![
|
||||
Action::Transfer(Transfer {
|
||||
deposit: 2 * ONE_NEAR,
|
||||
}),
|
||||
Action::FunctionCall(FunctionCall {
|
||||
method_name: "deposit".to_string(),
|
||||
args: vec![],
|
||||
gas: 30 * TGAS,
|
||||
deposit: ONE_NEAR,
|
||||
}),
|
||||
];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("bob.near", &actions, &perm, &policy);
|
||||
assert_eq!(analysis.total_value_yocto, 3 * ONE_NEAR);
|
||||
assert_eq!(analysis.actions.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_analysis_risk_levels() {
|
||||
let policy = default_policy();
|
||||
let actions = vec![
|
||||
Action::Transfer(Transfer { deposit: 0 }),
|
||||
Action::DeleteAccount(crate::keys::transaction::DeleteAccount {
|
||||
beneficiary_id: crate::keys::types::NearAccountId::new("bob.near").unwrap(),
|
||||
}),
|
||||
];
|
||||
let perm = AccessKeyPermission::FullAccess;
|
||||
let analysis = analyze_transaction("alice.near", &actions, &perm, &policy);
|
||||
assert_eq!(analysis.actions[0].risk_level, RiskLevel::Low);
|
||||
assert_eq!(analysis.actions[1].risk_level, RiskLevel::Critical);
|
||||
}
|
||||
|
||||
// -- Chain signature tests --
|
||||
|
||||
#[test]
|
||||
fn test_chain_sig_no_rule_requires_approval() {
|
||||
let policy = default_policy();
|
||||
let chain_sig = ChainSigAnalysis {
|
||||
derivation_path: "ethereum-1".to_string(),
|
||||
domain: SignatureDomain::Secp256k1,
|
||||
target_chain: Some("Ethereum".to_string()),
|
||||
payload_size: 256,
|
||||
risk_level: RiskLevel::Medium,
|
||||
};
|
||||
let decision = policy.evaluate_chain_sig(&chain_sig, 0);
|
||||
assert!(matches!(decision, PolicyDecision::RequireApproval { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_chain_sig_matching_rule_auto_approve() {
|
||||
let policy = PolicyConfig {
|
||||
chain_sig_rules: vec![ChainSigRule {
|
||||
allowed_paths: vec!["ethereum-*".to_string()],
|
||||
allowed_domains: vec![SignatureDomain::Secp256k1],
|
||||
max_payload_bytes: 1024,
|
||||
auto_approve: true,
|
||||
}],
|
||||
..default_policy()
|
||||
};
|
||||
let chain_sig = ChainSigAnalysis {
|
||||
derivation_path: "ethereum-1".to_string(),
|
||||
domain: SignatureDomain::Secp256k1,
|
||||
target_chain: Some("Ethereum".to_string()),
|
||||
payload_size: 256,
|
||||
risk_level: RiskLevel::Medium,
|
||||
};
|
||||
let decision = policy.evaluate_chain_sig(&chain_sig, 0);
|
||||
assert_eq!(decision, PolicyDecision::AutoApprove);
|
||||
}
|
||||
|
||||
// -- Glob matching tests --
|
||||
|
||||
#[test]
|
||||
fn test_glob_matches() {
|
||||
assert!(glob_matches("ethereum-*", "ethereum-1"));
|
||||
assert!(glob_matches("ethereum-*", "ethereum-mainnet"));
|
||||
assert!(!glob_matches("ethereum-*", "bitcoin-0"));
|
||||
assert!(glob_matches("exact-match", "exact-match"));
|
||||
assert!(!glob_matches("exact-match", "other"));
|
||||
}
|
||||
|
||||
// -- Infer target chain --
|
||||
|
||||
#[test]
|
||||
fn test_infer_target_chain() {
|
||||
assert_eq!(
|
||||
infer_target_chain("ethereum-1"),
|
||||
Some("Ethereum".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
infer_target_chain("bitcoin/0/0"),
|
||||
Some("Bitcoin".to_string())
|
||||
);
|
||||
assert_eq!(infer_target_chain("unknown-path"), None);
|
||||
}
|
||||
}
|
||||
+297
@@ -0,0 +1,297 @@
|
||||
//! Lightweight NEAR JSON-RPC client.
|
||||
//!
|
||||
//! Thin reqwest wrapper for the subset of NEAR RPC we need:
|
||||
//! - view_access_key (nonce + block_hash for transaction building)
|
||||
//! - send_transaction (submit signed transaction)
|
||||
//! - tx_status (poll for result)
|
||||
//! - view_account (check balance)
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::keys::KeyError;
|
||||
use crate::keys::types::NearNetwork;
|
||||
|
||||
/// NEAR RPC client.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NearRpcClient {
|
||||
client: reqwest::Client,
|
||||
rpc_url: String,
|
||||
}
|
||||
|
||||
impl NearRpcClient {
|
||||
pub fn new(network: &NearNetwork) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
rpc_url: network.rpc_url().to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_url(url: &str) -> Self {
|
||||
Self {
|
||||
client: reqwest::Client::new(),
|
||||
rpc_url: url.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch access key info (nonce + block hash) for signing a transaction.
|
||||
pub async fn view_access_key(
|
||||
&self,
|
||||
account_id: &str,
|
||||
public_key: &str,
|
||||
) -> Result<AccessKeyView, KeyError> {
|
||||
let response: RpcResponse<AccessKeyView> = self
|
||||
.call(
|
||||
"query",
|
||||
serde_json::json!({
|
||||
"request_type": "view_access_key",
|
||||
"finality": "final",
|
||||
"account_id": account_id,
|
||||
"public_key": public_key,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(response.result)
|
||||
}
|
||||
|
||||
/// Submit a signed transaction (fire and forget, returns tx hash).
|
||||
pub async fn send_transaction_async(&self, signed_tx_base64: &str) -> Result<String, KeyError> {
|
||||
let response: RpcResponse<serde_json::Value> = self
|
||||
.call("broadcast_tx_async", serde_json::json!([signed_tx_base64]))
|
||||
.await?;
|
||||
|
||||
response
|
||||
.result
|
||||
.as_str()
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| KeyError::RpcError {
|
||||
reason: "unexpected response from broadcast_tx_async".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Submit a signed transaction and wait for result.
|
||||
pub async fn send_transaction(&self, signed_tx_base64: &str) -> Result<TxOutcome, KeyError> {
|
||||
let response: RpcResponse<TxOutcome> = self
|
||||
.call("broadcast_tx_commit", serde_json::json!([signed_tx_base64]))
|
||||
.await?;
|
||||
|
||||
Ok(response.result)
|
||||
}
|
||||
|
||||
/// Check transaction status.
|
||||
pub async fn tx_status(&self, tx_hash: &str, sender_id: &str) -> Result<TxOutcome, KeyError> {
|
||||
let response: RpcResponse<TxOutcome> = self
|
||||
.call("tx", serde_json::json!([tx_hash, sender_id]))
|
||||
.await?;
|
||||
|
||||
Ok(response.result)
|
||||
}
|
||||
|
||||
/// View account information.
|
||||
pub async fn view_account(&self, account_id: &str) -> Result<AccountView, KeyError> {
|
||||
let response: RpcResponse<AccountView> = self
|
||||
.call(
|
||||
"query",
|
||||
serde_json::json!({
|
||||
"request_type": "view_account",
|
||||
"finality": "final",
|
||||
"account_id": account_id,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(response.result)
|
||||
}
|
||||
|
||||
/// Make a JSON-RPC 2.0 call.
|
||||
async fn call<T: for<'de> Deserialize<'de>>(
|
||||
&self,
|
||||
method: &str,
|
||||
params: serde_json::Value,
|
||||
) -> Result<RpcResponse<T>, KeyError> {
|
||||
let request = RpcRequest {
|
||||
jsonrpc: "2.0",
|
||||
id: "ironclaw",
|
||||
method,
|
||||
params,
|
||||
};
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.post(&self.rpc_url)
|
||||
.json(&request)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(KeyError::RpcError {
|
||||
reason: format!("HTTP {}: {}", status, truncate(&body, 200)),
|
||||
});
|
||||
}
|
||||
|
||||
let body = response.text().await?;
|
||||
let parsed: serde_json::Value =
|
||||
serde_json::from_str(&body).map_err(|e| KeyError::RpcError {
|
||||
reason: format!("invalid JSON response: {}", e),
|
||||
})?;
|
||||
|
||||
// Check for JSON-RPC error
|
||||
if let Some(error) = parsed.get("error") {
|
||||
let cause = error
|
||||
.get("cause")
|
||||
.and_then(|c| c.get("name"))
|
||||
.and_then(|n| n.as_str())
|
||||
.unwrap_or("unknown");
|
||||
let message = error
|
||||
.get("message")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or("unknown error");
|
||||
return Err(KeyError::RpcError {
|
||||
reason: format!("{}: {}", cause, message),
|
||||
});
|
||||
}
|
||||
|
||||
serde_json::from_value(parsed).map_err(|e| KeyError::RpcError {
|
||||
reason: format!("failed to parse RPC response: {}", e),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate(s: &str, max: usize) -> String {
|
||||
if s.len() <= max {
|
||||
s.to_string()
|
||||
} else {
|
||||
format!("{}...", &s[..max])
|
||||
}
|
||||
}
|
||||
|
||||
/// JSON-RPC 2.0 request.
|
||||
#[derive(Serialize)]
|
||||
struct RpcRequest<'a> {
|
||||
jsonrpc: &'a str,
|
||||
id: &'a str,
|
||||
method: &'a str,
|
||||
params: serde_json::Value,
|
||||
}
|
||||
|
||||
/// JSON-RPC 2.0 response.
|
||||
#[derive(Deserialize)]
|
||||
struct RpcResponse<T> {
|
||||
result: T,
|
||||
}
|
||||
|
||||
/// Access key view from RPC.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct AccessKeyView {
|
||||
pub nonce: u64,
|
||||
pub block_hash: String,
|
||||
pub permission: serde_json::Value,
|
||||
}
|
||||
|
||||
/// Transaction outcome from RPC.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct TxOutcome {
|
||||
pub status: serde_json::Value,
|
||||
pub transaction: Option<serde_json::Value>,
|
||||
pub transaction_outcome: Option<serde_json::Value>,
|
||||
pub receipts_outcome: Option<Vec<serde_json::Value>>,
|
||||
}
|
||||
|
||||
impl TxOutcome {
|
||||
/// Check if the transaction succeeded.
|
||||
pub fn is_success(&self) -> bool {
|
||||
if let Some(obj) = self.status.as_object() {
|
||||
obj.contains_key("SuccessValue") || obj.contains_key("SuccessReceiptId")
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the failure reason if the transaction failed.
|
||||
pub fn failure_reason(&self) -> Option<String> {
|
||||
if let Some(obj) = self.status.as_object() {
|
||||
if let Some(failure) = obj.get("Failure") {
|
||||
return Some(format!("{}", failure));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Account view from RPC.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct AccountView {
|
||||
pub amount: String,
|
||||
pub locked: String,
|
||||
pub storage_usage: u64,
|
||||
pub code_hash: String,
|
||||
pub block_height: u64,
|
||||
pub block_hash: String,
|
||||
}
|
||||
|
||||
impl AccountView {
|
||||
/// Parse the balance as u128 (yoctoNEAR).
|
||||
pub fn balance_yocto(&self) -> Result<u128, KeyError> {
|
||||
self.amount.parse::<u128>().map_err(|e| KeyError::RpcError {
|
||||
reason: format!("failed to parse account balance '{}': {}", self.amount, e),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::keys::rpc::{AccessKeyView, AccountView, TxOutcome};
|
||||
|
||||
#[test]
|
||||
fn test_tx_outcome_success() {
|
||||
let outcome = TxOutcome {
|
||||
status: serde_json::json!({"SuccessValue": ""}),
|
||||
transaction: None,
|
||||
transaction_outcome: None,
|
||||
receipts_outcome: None,
|
||||
};
|
||||
assert!(outcome.is_success());
|
||||
assert!(outcome.failure_reason().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tx_outcome_failure() {
|
||||
let outcome = TxOutcome {
|
||||
status: serde_json::json!({"Failure": {"ActionError": "..."}}),
|
||||
transaction: None,
|
||||
transaction_outcome: None,
|
||||
receipts_outcome: None,
|
||||
};
|
||||
assert!(!outcome.is_success());
|
||||
assert!(outcome.failure_reason().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_access_key_view_deserialize() {
|
||||
let json = serde_json::json!({
|
||||
"nonce": 42,
|
||||
"block_hash": "11111111111111111111111111111111",
|
||||
"permission": "FullAccess"
|
||||
});
|
||||
let view: AccessKeyView = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(view.nonce, 42);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_account_view_balance() {
|
||||
let view = AccountView {
|
||||
amount: "1000000000000000000000000".to_string(), // 1 NEAR
|
||||
locked: "0".to_string(),
|
||||
storage_usage: 100,
|
||||
code_hash: "11111111111111111111111111111111".to_string(),
|
||||
block_height: 1000,
|
||||
block_hash: "11111111111111111111111111111111".to_string(),
|
||||
};
|
||||
assert_eq!(
|
||||
view.balance_yocto().unwrap(),
|
||||
1_000_000_000_000_000_000_000_000
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
//! Ed25519 signing for NEAR transactions.
|
||||
//!
|
||||
//! SECURITY: Private keys are held in memory for the absolute minimum time.
|
||||
//! The flow is: decrypt -> construct SigningKey -> sign -> drop (Zeroize).
|
||||
//! The `ed25519_dalek::SigningKey` implements Zeroize, so memory is zeroed on drop.
|
||||
|
||||
use ed25519_dalek::Signer;
|
||||
use sha2::{Digest, Sha256};
|
||||
use zeroize::Zeroize;
|
||||
|
||||
use crate::keys::KeyError;
|
||||
use crate::keys::types::NearPublicKey;
|
||||
use crate::secrets::SecretsStore;
|
||||
|
||||
/// Parse a NEAR-format secret key and extract the 32-byte ed25519 seed.
|
||||
///
|
||||
/// NEAR secret keys are formatted as `ed25519:<base58-encoded-64-bytes>`.
|
||||
/// The 64 bytes are the seed (32) + public key (32) concatenated.
|
||||
/// Some wallets store only the 32-byte seed with the same prefix.
|
||||
fn parse_near_secret_key(near_format: &str) -> Result<[u8; 32], KeyError> {
|
||||
let data_str =
|
||||
near_format
|
||||
.strip_prefix("ed25519:")
|
||||
.ok_or_else(|| KeyError::InvalidKeyFormat {
|
||||
reason: "secret key must start with 'ed25519:'".to_string(),
|
||||
})?;
|
||||
|
||||
let mut bytes = bs58::decode(data_str)
|
||||
.into_vec()
|
||||
.map_err(|e| KeyError::InvalidKeyFormat {
|
||||
reason: format!("invalid base58 in secret key: {}", e),
|
||||
})?;
|
||||
|
||||
let seed = match bytes.len() {
|
||||
64 => {
|
||||
// Standard NEAR format: seed (32) + public key (32)
|
||||
let mut seed = [0u8; 32];
|
||||
seed.copy_from_slice(&bytes[..32]);
|
||||
bytes.zeroize();
|
||||
seed
|
||||
}
|
||||
32 => {
|
||||
// Some wallets export just the seed
|
||||
let mut seed = [0u8; 32];
|
||||
seed.copy_from_slice(&bytes);
|
||||
bytes.zeroize();
|
||||
seed
|
||||
}
|
||||
other => {
|
||||
bytes.zeroize();
|
||||
return Err(KeyError::InvalidKeyFormat {
|
||||
reason: format!("ed25519 secret key must be 32 or 64 bytes, got {}", other),
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
Ok(seed)
|
||||
}
|
||||
|
||||
/// Derive the public key from a NEAR-format secret key string.
|
||||
pub fn public_key_from_secret(near_format_secret: &str) -> Result<NearPublicKey, KeyError> {
|
||||
let seed = parse_near_secret_key(near_format_secret)?;
|
||||
let signing_key = ed25519_dalek::SigningKey::from_bytes(&seed);
|
||||
let verifying_key = signing_key.verifying_key();
|
||||
// signing_key implements Zeroize on drop
|
||||
|
||||
Ok(NearPublicKey {
|
||||
key_type: crate::keys::types::KeyType::Ed25519,
|
||||
data: verifying_key.to_bytes(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Sign a 32-byte SHA-256 hash using a key from the secrets store.
|
||||
///
|
||||
/// This is the core signing function. It:
|
||||
/// 1. Decrypts the private key from the secrets store
|
||||
/// 2. Parses the NEAR-format key to extract the ed25519 seed
|
||||
/// 3. Constructs a SigningKey (implements Zeroize on drop)
|
||||
/// 4. Signs the hash
|
||||
/// 5. Drops the SigningKey (memory zeroed)
|
||||
///
|
||||
/// The plaintext key exists in memory for microseconds.
|
||||
pub async fn sign_hash(
|
||||
secrets_store: &dyn SecretsStore,
|
||||
user_id: &str,
|
||||
label: &str,
|
||||
hash: &[u8; 32],
|
||||
) -> Result<[u8; 64], KeyError> {
|
||||
let secret_name = format!("near_key:{}", label);
|
||||
let decrypted = secrets_store
|
||||
.get_decrypted(user_id, &secret_name)
|
||||
.await
|
||||
.map_err(|e| KeyError::SigningFailed {
|
||||
reason: format!("failed to decrypt key '{}': {}", label, e),
|
||||
})?;
|
||||
|
||||
let mut seed = parse_near_secret_key(decrypted.expose())?;
|
||||
let signing_key = ed25519_dalek::SigningKey::from_bytes(&seed);
|
||||
seed.zeroize();
|
||||
|
||||
let signature = signing_key.sign(hash);
|
||||
// signing_key drops here, Zeroize zeroes the key material
|
||||
|
||||
Ok(signature.to_bytes())
|
||||
}
|
||||
|
||||
/// SHA-256 hash of data (used for transaction signing).
|
||||
pub fn sha256_hash(data: &[u8]) -> [u8; 32] {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(data);
|
||||
hasher.finalize().into()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use ed25519_dalek::{Signer, SigningKey, Verifier, VerifyingKey};
|
||||
use secrecy::SecretString;
|
||||
|
||||
use crate::keys::signer::{
|
||||
parse_near_secret_key, public_key_from_secret, sha256_hash, sign_hash,
|
||||
};
|
||||
use crate::keys::types::KeyType;
|
||||
use crate::secrets::{CreateSecretParams, InMemorySecretsStore, SecretsCrypto, SecretsStore};
|
||||
|
||||
fn test_store() -> Arc<InMemorySecretsStore> {
|
||||
let key = "0123456789abcdef0123456789abcdef";
|
||||
let crypto = Arc::new(SecretsCrypto::new(SecretString::from(key.to_string())).unwrap());
|
||||
Arc::new(InMemorySecretsStore::new(crypto))
|
||||
}
|
||||
|
||||
/// Generate a test keypair and return (near_format_secret, near_format_public).
|
||||
fn generate_test_keypair() -> (String, String) {
|
||||
let signing_key = SigningKey::generate(&mut rand::rngs::OsRng);
|
||||
let verifying_key = signing_key.verifying_key();
|
||||
|
||||
// NEAR format: ed25519:<base58(seed + pubkey)>
|
||||
let mut combined = Vec::with_capacity(64);
|
||||
combined.extend_from_slice(signing_key.as_bytes());
|
||||
combined.extend_from_slice(verifying_key.as_bytes());
|
||||
|
||||
let secret = format!("ed25519:{}", bs58::encode(&combined).into_string());
|
||||
let public = format!(
|
||||
"ed25519:{}",
|
||||
bs58::encode(verifying_key.as_bytes()).into_string()
|
||||
);
|
||||
|
||||
(secret, public)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_near_secret_key_64_bytes() {
|
||||
let (secret, _) = generate_test_keypair();
|
||||
let seed = parse_near_secret_key(&secret).unwrap();
|
||||
assert_eq!(seed.len(), 32);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_near_secret_key_32_bytes() {
|
||||
// Some wallets export just the 32-byte seed
|
||||
let signing_key = SigningKey::generate(&mut rand::rngs::OsRng);
|
||||
let secret = format!(
|
||||
"ed25519:{}",
|
||||
bs58::encode(signing_key.as_bytes()).into_string()
|
||||
);
|
||||
let seed = parse_near_secret_key(&secret).unwrap();
|
||||
assert_eq!(seed, *signing_key.as_bytes());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_invalid_prefix() {
|
||||
assert!(parse_near_secret_key("secp256k1:abc").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_public_key_from_secret() {
|
||||
let (secret, expected_public) = generate_test_keypair();
|
||||
let pubkey = public_key_from_secret(&secret).unwrap();
|
||||
assert_eq!(pubkey.key_type, KeyType::Ed25519);
|
||||
assert_eq!(pubkey.to_near_format(), expected_public);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sign_and_verify_roundtrip() {
|
||||
let signing_key = SigningKey::generate(&mut rand::rngs::OsRng);
|
||||
let verifying_key = signing_key.verifying_key();
|
||||
|
||||
let message = b"test message for signing";
|
||||
let hash = sha256_hash(message);
|
||||
|
||||
let signature = signing_key.sign(&hash);
|
||||
|
||||
// Verify
|
||||
assert!(verifying_key.verify(&hash, &signature).is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sign_hash_from_store() {
|
||||
let store = test_store();
|
||||
let (secret, _public) = generate_test_keypair();
|
||||
|
||||
// Store the key
|
||||
store
|
||||
.create(
|
||||
"user1",
|
||||
CreateSecretParams::new("near_key:test-signer", &secret).with_provider("near_keys"),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Sign
|
||||
let hash = sha256_hash(b"test transaction data");
|
||||
let sig_bytes = sign_hash(store.as_ref(), "user1", "test-signer", &hash)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Verify using the public key derived from the secret
|
||||
let pubkey = public_key_from_secret(&secret).unwrap();
|
||||
let verifying_key = VerifyingKey::from_bytes(pubkey.as_bytes()).unwrap();
|
||||
let signature = ed25519_dalek::Signature::from_bytes(&sig_bytes);
|
||||
assert!(verifying_key.verify(&hash, &signature).is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_sign_hash_key_not_found() {
|
||||
let store = test_store();
|
||||
let hash = [0u8; 32];
|
||||
let result = sign_hash(store.as_ref(), "user1", "nonexistent", &hash).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sha256_hash() {
|
||||
let hash = sha256_hash(b"hello");
|
||||
let hex: String = hash.iter().map(|b| format!("{:02x}", b)).collect();
|
||||
// Known SHA-256 of "hello"
|
||||
assert_eq!(
|
||||
hex,
|
||||
"2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
//! Daily spend tracking for rate-limiting value transfers.
|
||||
//!
|
||||
//! Tracks cumulative daily spend in yoctoNEAR to enforce `daily_spend_limit_yocto`.
|
||||
//! Persisted to `~/.ironclaw/spend_tracking.json`. Resets automatically at midnight UTC.
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
use chrono::{DateTime, NaiveDate, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::fs;
|
||||
|
||||
use crate::keys::KeyError;
|
||||
use crate::keys::types::format_yocto;
|
||||
|
||||
/// Tracks daily cumulative spend for policy enforcement.
|
||||
pub struct SpendTracker {
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl SpendTracker {
|
||||
pub fn new(path: PathBuf) -> Self {
|
||||
Self { path }
|
||||
}
|
||||
|
||||
/// Default location: `~/.ironclaw/spend_tracking.json`
|
||||
pub fn default_path() -> PathBuf {
|
||||
dirs::home_dir()
|
||||
.map(|h| h.join(".ironclaw").join("spend_tracking.json"))
|
||||
.unwrap_or_else(|| PathBuf::from(".ironclaw/spend_tracking.json"))
|
||||
}
|
||||
|
||||
/// Get today's cumulative spend in yoctoNEAR.
|
||||
pub async fn get_daily_spend(&self) -> Result<u128, KeyError> {
|
||||
let data = self.load().await?;
|
||||
let today = Utc::now().date_naive();
|
||||
|
||||
Ok(data
|
||||
.records
|
||||
.iter()
|
||||
.find(|r| r.date == today)
|
||||
.map(|r| r.total_spent_yocto)
|
||||
.unwrap_or(0))
|
||||
}
|
||||
|
||||
/// Record a spend after successful transaction submission.
|
||||
pub async fn record_spend(
|
||||
&self,
|
||||
value_yocto: u128,
|
||||
description: String,
|
||||
tx_hash: Option<String>,
|
||||
) -> Result<(), KeyError> {
|
||||
let mut data = self.load().await?;
|
||||
let today = Utc::now().date_naive();
|
||||
|
||||
let record = data.records.iter_mut().find(|r| r.date == today);
|
||||
|
||||
let entry = SpendEntry {
|
||||
timestamp: Utc::now(),
|
||||
tx_hash,
|
||||
value_yocto,
|
||||
description,
|
||||
};
|
||||
|
||||
if let Some(record) = record {
|
||||
record.total_spent_yocto = record.total_spent_yocto.saturating_add(value_yocto);
|
||||
record.transactions.push(entry);
|
||||
} else {
|
||||
data.records.push(SpendRecord {
|
||||
date: today,
|
||||
total_spent_yocto: value_yocto,
|
||||
transactions: vec![entry],
|
||||
});
|
||||
}
|
||||
|
||||
// Keep only last 30 days of records
|
||||
let cutoff = Utc::now().date_naive() - chrono::Duration::days(30);
|
||||
data.records.retain(|r| r.date >= cutoff);
|
||||
|
||||
self.save(&data).await
|
||||
}
|
||||
|
||||
/// Get spend history for the last N days.
|
||||
pub async fn get_history(&self, days: u32) -> Result<Vec<SpendRecord>, KeyError> {
|
||||
let data = self.load().await?;
|
||||
let cutoff = Utc::now().date_naive() - chrono::Duration::days(days as i64);
|
||||
|
||||
Ok(data
|
||||
.records
|
||||
.into_iter()
|
||||
.filter(|r| r.date >= cutoff)
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn load(&self) -> Result<SpendData, KeyError> {
|
||||
if !self.path.exists() {
|
||||
return Ok(SpendData::default());
|
||||
}
|
||||
|
||||
let content = fs::read_to_string(&self.path).await?;
|
||||
serde_json::from_str(&content).map_err(|e| {
|
||||
KeyError::Io(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
format!("corrupt spend tracking data: {}", e),
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
async fn save(&self, data: &SpendData) -> Result<(), KeyError> {
|
||||
if let Some(parent) = self.path.parent() {
|
||||
fs::create_dir_all(parent).await?;
|
||||
}
|
||||
|
||||
let content = serde_json::to_string_pretty(data).map_err(|e| {
|
||||
KeyError::SerializationFailed(format!("failed to serialize spend data: {}", e))
|
||||
})?;
|
||||
|
||||
fs::write(&self.path, content).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Serialize, Deserialize)]
|
||||
struct SpendData {
|
||||
records: Vec<SpendRecord>,
|
||||
}
|
||||
|
||||
/// A day's spend record with audit trail.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SpendRecord {
|
||||
pub date: NaiveDate,
|
||||
pub total_spent_yocto: u128,
|
||||
pub transactions: Vec<SpendEntry>,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for SpendRecord {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"{}: {} ({} txns)",
|
||||
self.date,
|
||||
format_yocto(self.total_spent_yocto),
|
||||
self.transactions.len()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// A single spend entry in the audit trail.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SpendEntry {
|
||||
pub timestamp: DateTime<Utc>,
|
||||
pub tx_hash: Option<String>,
|
||||
pub value_yocto: u128,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use tempfile::TempDir;
|
||||
|
||||
use crate::keys::spending::SpendTracker;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_empty_spend() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let tracker = SpendTracker::new(dir.path().join("spend.json"));
|
||||
assert_eq!(tracker.get_daily_spend().await.unwrap(), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_record_and_query_spend() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let tracker = SpendTracker::new(dir.path().join("spend.json"));
|
||||
|
||||
tracker
|
||||
.record_spend(
|
||||
1_000_000,
|
||||
"test transfer".to_string(),
|
||||
Some("hash1".to_string()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(tracker.get_daily_spend().await.unwrap(), 1_000_000);
|
||||
|
||||
tracker
|
||||
.record_spend(2_000_000, "another transfer".to_string(), None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(tracker.get_daily_spend().await.unwrap(), 3_000_000);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_history() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let tracker = SpendTracker::new(dir.path().join("spend.json"));
|
||||
|
||||
tracker
|
||||
.record_spend(100, "test".to_string(), None)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let history = tracker.get_history(7).await.unwrap();
|
||||
assert_eq!(history.len(), 1);
|
||||
assert_eq!(history[0].total_spent_yocto, 100);
|
||||
assert_eq!(history[0].transactions.len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,445 @@
|
||||
//! Minimal NEAR transaction types with borsh serialization.
|
||||
//!
|
||||
//! Hand-rolled types that produce byte-identical borsh output to near-primitives,
|
||||
//! without pulling in the massive nearcore dependency tree.
|
||||
//!
|
||||
//! # Serialization Format
|
||||
//!
|
||||
//! NEAR transactions are borsh-serialized, then SHA-256 hashed for signing.
|
||||
//! The signed transaction includes the original transaction + ed25519 signature.
|
||||
|
||||
use borsh::BorshSerialize;
|
||||
|
||||
use crate::keys::signer::sha256_hash;
|
||||
use crate::keys::types::{KeyType, NearAccountId, NearPublicKey};
|
||||
|
||||
/// A NEAR transaction ready for signing.
|
||||
#[derive(Debug, Clone, BorshSerialize)]
|
||||
pub struct Transaction {
|
||||
pub signer_id: NearAccountId,
|
||||
pub public_key: NearPublicKey,
|
||||
pub nonce: u64,
|
||||
pub receiver_id: NearAccountId,
|
||||
pub block_hash: BlockHash,
|
||||
pub actions: Vec<Action>,
|
||||
}
|
||||
|
||||
impl Transaction {
|
||||
/// Borsh-serialize and SHA-256 hash for signing.
|
||||
pub fn hash_for_signing(&self) -> Result<[u8; 32], crate::keys::KeyError> {
|
||||
let bytes = borsh::to_vec(self).map_err(|e| {
|
||||
crate::keys::KeyError::SerializationFailed(format!(
|
||||
"failed to serialize transaction: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
Ok(sha256_hash(&bytes))
|
||||
}
|
||||
}
|
||||
|
||||
/// A signed NEAR transaction with ed25519 signature.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SignedTransaction {
|
||||
pub transaction: Transaction,
|
||||
pub signature: Signature,
|
||||
}
|
||||
|
||||
impl SignedTransaction {
|
||||
/// Encode as base64 for RPC submission.
|
||||
pub fn to_base64(&self) -> Result<String, crate::keys::KeyError> {
|
||||
let bytes = self.to_borsh()?;
|
||||
Ok(base64::Engine::encode(
|
||||
&base64::engine::general_purpose::STANDARD,
|
||||
&bytes,
|
||||
))
|
||||
}
|
||||
|
||||
/// Borsh-serialize the signed transaction.
|
||||
pub fn to_borsh(&self) -> Result<Vec<u8>, crate::keys::KeyError> {
|
||||
let mut buf = Vec::new();
|
||||
borsh::BorshSerialize::serialize(&self.transaction, &mut buf).map_err(|e| {
|
||||
crate::keys::KeyError::SerializationFailed(format!(
|
||||
"failed to serialize signed transaction: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
borsh::BorshSerialize::serialize(&self.signature, &mut buf).map_err(|e| {
|
||||
crate::keys::KeyError::SerializationFailed(format!(
|
||||
"failed to serialize signature: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
/// Get the transaction hash (the hash that was signed).
|
||||
pub fn tx_hash(&self) -> Result<[u8; 32], crate::keys::KeyError> {
|
||||
self.transaction.hash_for_signing()
|
||||
}
|
||||
}
|
||||
|
||||
/// Block hash (32 bytes), used as recent block reference for transaction validity.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct BlockHash(pub [u8; 32]);
|
||||
|
||||
impl BorshSerialize for BlockHash {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
writer.write_all(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl BlockHash {
|
||||
pub fn from_base58(s: &str) -> Result<Self, crate::keys::KeyError> {
|
||||
let bytes =
|
||||
bs58::decode(s)
|
||||
.into_vec()
|
||||
.map_err(|e| crate::keys::KeyError::InvalidKeyFormat {
|
||||
reason: format!("invalid base58 block hash: {}", e),
|
||||
})?;
|
||||
if bytes.len() != 32 {
|
||||
return Err(crate::keys::KeyError::InvalidKeyFormat {
|
||||
reason: format!("block hash must be 32 bytes, got {}", bytes.len()),
|
||||
});
|
||||
}
|
||||
let mut hash = [0u8; 32];
|
||||
hash.copy_from_slice(&bytes);
|
||||
Ok(Self(hash))
|
||||
}
|
||||
}
|
||||
|
||||
/// Ed25519 signature (NEAR uses key_type prefix for borsh serialization).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Signature {
|
||||
pub key_type: KeyType,
|
||||
pub data: [u8; 64],
|
||||
}
|
||||
|
||||
impl BorshSerialize for Signature {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
BorshSerialize::serialize(&self.key_type, writer)?;
|
||||
writer.write_all(&self.data)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// NEAR transaction action variants.
|
||||
///
|
||||
/// Only includes the variants we actually need for key management operations.
|
||||
/// Borsh enum discriminants MUST match near-primitives exactly.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum Action {
|
||||
CreateAccount, // 0
|
||||
DeployContract(DeployContract), // 1
|
||||
FunctionCall(FunctionCall), // 2
|
||||
Transfer(Transfer), // 3
|
||||
Stake(Stake), // 4
|
||||
AddKey(AddKey), // 5
|
||||
DeleteKey(DeleteKey), // 6
|
||||
DeleteAccount(DeleteAccount), // 7
|
||||
}
|
||||
|
||||
impl BorshSerialize for Action {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
match self {
|
||||
Action::CreateAccount => {
|
||||
BorshSerialize::serialize(&0u8, writer)?;
|
||||
}
|
||||
Action::DeployContract(v) => {
|
||||
BorshSerialize::serialize(&1u8, writer)?;
|
||||
BorshSerialize::serialize(v, writer)?;
|
||||
}
|
||||
Action::FunctionCall(v) => {
|
||||
BorshSerialize::serialize(&2u8, writer)?;
|
||||
BorshSerialize::serialize(v, writer)?;
|
||||
}
|
||||
Action::Transfer(v) => {
|
||||
BorshSerialize::serialize(&3u8, writer)?;
|
||||
BorshSerialize::serialize(v, writer)?;
|
||||
}
|
||||
Action::Stake(v) => {
|
||||
BorshSerialize::serialize(&4u8, writer)?;
|
||||
BorshSerialize::serialize(v, writer)?;
|
||||
}
|
||||
Action::AddKey(v) => {
|
||||
BorshSerialize::serialize(&5u8, writer)?;
|
||||
BorshSerialize::serialize(v, writer)?;
|
||||
}
|
||||
Action::DeleteKey(v) => {
|
||||
BorshSerialize::serialize(&6u8, writer)?;
|
||||
BorshSerialize::serialize(v, writer)?;
|
||||
}
|
||||
Action::DeleteAccount(v) => {
|
||||
BorshSerialize::serialize(&7u8, writer)?;
|
||||
BorshSerialize::serialize(v, writer)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Deploy contract action.
|
||||
#[derive(Debug, Clone, BorshSerialize)]
|
||||
pub struct DeployContract {
|
||||
pub code: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Function call action.
|
||||
#[derive(Debug, Clone, BorshSerialize)]
|
||||
pub struct FunctionCall {
|
||||
pub method_name: String,
|
||||
pub args: Vec<u8>,
|
||||
pub gas: u64,
|
||||
pub deposit: u128,
|
||||
}
|
||||
|
||||
/// Transfer action.
|
||||
#[derive(Debug, Clone, BorshSerialize)]
|
||||
pub struct Transfer {
|
||||
pub deposit: u128,
|
||||
}
|
||||
|
||||
/// Stake action.
|
||||
#[derive(Debug, Clone, BorshSerialize)]
|
||||
pub struct Stake {
|
||||
pub stake: u128,
|
||||
pub public_key: NearPublicKey,
|
||||
}
|
||||
|
||||
/// Add key action.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AddKey {
|
||||
pub public_key: NearPublicKey,
|
||||
pub access_key: AccessKeyBorsh,
|
||||
}
|
||||
|
||||
impl BorshSerialize for AddKey {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
BorshSerialize::serialize(&self.public_key, writer)?;
|
||||
BorshSerialize::serialize(&self.access_key, writer)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Delete key action.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DeleteKey {
|
||||
pub public_key: NearPublicKey,
|
||||
}
|
||||
|
||||
impl BorshSerialize for DeleteKey {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
BorshSerialize::serialize(&self.public_key, writer)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Delete account action.
|
||||
#[derive(Debug, Clone, BorshSerialize)]
|
||||
pub struct DeleteAccount {
|
||||
pub beneficiary_id: NearAccountId,
|
||||
}
|
||||
|
||||
/// Borsh-serializable access key (for AddKey actions).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccessKeyBorsh {
|
||||
pub nonce: u64,
|
||||
pub permission: AccessKeyPermissionBorsh,
|
||||
}
|
||||
|
||||
impl BorshSerialize for AccessKeyBorsh {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
BorshSerialize::serialize(&self.nonce, writer)?;
|
||||
BorshSerialize::serialize(&self.permission, writer)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Borsh-serializable access key permission.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum AccessKeyPermissionBorsh {
|
||||
FunctionCall(FunctionCallPermissionBorsh),
|
||||
FullAccess,
|
||||
}
|
||||
|
||||
impl BorshSerialize for AccessKeyPermissionBorsh {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
match self {
|
||||
AccessKeyPermissionBorsh::FunctionCall(fc) => {
|
||||
BorshSerialize::serialize(&0u8, writer)?;
|
||||
BorshSerialize::serialize(fc, writer)?;
|
||||
}
|
||||
AccessKeyPermissionBorsh::FullAccess => {
|
||||
BorshSerialize::serialize(&1u8, writer)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Borsh-serializable function call permission.
|
||||
#[derive(Debug, Clone, BorshSerialize)]
|
||||
pub struct FunctionCallPermissionBorsh {
|
||||
/// Allowance in yoctoNEAR (None = unlimited within key scope).
|
||||
pub allowance: Option<u128>,
|
||||
pub receiver_id: String,
|
||||
pub method_names: Vec<String>,
|
||||
}
|
||||
|
||||
/// Standard gas amounts.
|
||||
pub const TGAS: u64 = 1_000_000_000_000;
|
||||
|
||||
/// 300 TGas, the maximum per transaction.
|
||||
pub const MAX_GAS: u64 = 300 * TGAS;
|
||||
|
||||
/// 1 yoctoNEAR, commonly used as a deposit to indicate "attached" value.
|
||||
pub const ONE_YOCTO: u128 = 1;
|
||||
|
||||
/// 1 NEAR in yoctoNEAR.
|
||||
pub const ONE_NEAR: u128 = 1_000_000_000_000_000_000_000_000;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::keys::transaction::{
|
||||
AccessKeyBorsh, AccessKeyPermissionBorsh, Action, BlockHash, FunctionCall,
|
||||
FunctionCallPermissionBorsh, MAX_GAS, ONE_NEAR, ONE_YOCTO, Signature, TGAS, Transaction,
|
||||
Transfer,
|
||||
};
|
||||
use crate::keys::types::{KeyType, NearAccountId, NearPublicKey};
|
||||
|
||||
fn test_public_key() -> NearPublicKey {
|
||||
NearPublicKey::from_near_format("ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp")
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transfer_action_borsh() {
|
||||
let action = Action::Transfer(Transfer { deposit: ONE_NEAR });
|
||||
let bytes = borsh::to_vec(&action).unwrap();
|
||||
// Discriminant (1 byte) + u128 (16 bytes)
|
||||
assert_eq!(bytes.len(), 1 + 16);
|
||||
assert_eq!(bytes[0], 3); // Transfer = discriminant 3
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_function_call_action_borsh() {
|
||||
let action = Action::FunctionCall(FunctionCall {
|
||||
method_name: "deposit".to_string(),
|
||||
args: b"{}".to_vec(),
|
||||
gas: 30 * TGAS,
|
||||
deposit: ONE_YOCTO,
|
||||
});
|
||||
let bytes = borsh::to_vec(&action).unwrap();
|
||||
assert_eq!(bytes[0], 2); // FunctionCall = discriminant 2
|
||||
// Verify it serializes without error
|
||||
assert!(bytes.len() > 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transaction_hash_for_signing() {
|
||||
let tx = Transaction {
|
||||
signer_id: NearAccountId::new("alice.near").unwrap(),
|
||||
public_key: test_public_key(),
|
||||
nonce: 1,
|
||||
receiver_id: NearAccountId::new("bob.near").unwrap(),
|
||||
block_hash: BlockHash([0u8; 32]),
|
||||
actions: vec![Action::Transfer(Transfer { deposit: ONE_NEAR })],
|
||||
};
|
||||
|
||||
let hash = tx.hash_for_signing().unwrap();
|
||||
assert_eq!(hash.len(), 32);
|
||||
|
||||
// Same transaction should produce same hash
|
||||
let hash2 = tx.hash_for_signing().unwrap();
|
||||
assert_eq!(hash, hash2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transaction_different_nonce_different_hash() {
|
||||
let tx1 = Transaction {
|
||||
signer_id: NearAccountId::new("alice.near").unwrap(),
|
||||
public_key: test_public_key(),
|
||||
nonce: 1,
|
||||
receiver_id: NearAccountId::new("bob.near").unwrap(),
|
||||
block_hash: BlockHash([0u8; 32]),
|
||||
actions: vec![Action::Transfer(Transfer { deposit: ONE_NEAR })],
|
||||
};
|
||||
|
||||
let tx2 = Transaction {
|
||||
nonce: 2,
|
||||
..tx1.clone()
|
||||
};
|
||||
|
||||
assert_ne!(
|
||||
tx1.hash_for_signing().unwrap(),
|
||||
tx2.hash_for_signing().unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_signed_transaction_to_base64() {
|
||||
let tx = Transaction {
|
||||
signer_id: NearAccountId::new("alice.near").unwrap(),
|
||||
public_key: test_public_key(),
|
||||
nonce: 1,
|
||||
receiver_id: NearAccountId::new("bob.near").unwrap(),
|
||||
block_hash: BlockHash([0u8; 32]),
|
||||
actions: vec![Action::Transfer(Transfer { deposit: ONE_NEAR })],
|
||||
};
|
||||
|
||||
let signed = crate::keys::transaction::SignedTransaction {
|
||||
transaction: tx,
|
||||
signature: Signature {
|
||||
key_type: KeyType::Ed25519,
|
||||
data: [0u8; 64],
|
||||
},
|
||||
};
|
||||
|
||||
let b64 = signed.to_base64().unwrap();
|
||||
assert!(!b64.is_empty());
|
||||
|
||||
// Should be valid base64
|
||||
let decoded =
|
||||
base64::Engine::decode(&base64::engine::general_purpose::STANDARD, &b64).unwrap();
|
||||
assert!(!decoded.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_block_hash_from_base58() {
|
||||
let hash_str = "11111111111111111111111111111111"; // 32 zero bytes in base58
|
||||
let hash = BlockHash::from_base58(hash_str).unwrap();
|
||||
assert_eq!(hash.0, [0u8; 32]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_access_key_borsh_full_access() {
|
||||
let ak = AccessKeyBorsh {
|
||||
nonce: 0,
|
||||
permission: AccessKeyPermissionBorsh::FullAccess,
|
||||
};
|
||||
let bytes = borsh::to_vec(&ak).unwrap();
|
||||
// u64 (8 bytes) + discriminant (1 byte)
|
||||
assert_eq!(bytes.len(), 9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_access_key_borsh_function_call() {
|
||||
let ak = AccessKeyBorsh {
|
||||
nonce: 0,
|
||||
permission: AccessKeyPermissionBorsh::FunctionCall(FunctionCallPermissionBorsh {
|
||||
allowance: Some(ONE_NEAR),
|
||||
receiver_id: "contract.near".to_string(),
|
||||
method_names: vec!["deposit".to_string()],
|
||||
}),
|
||||
};
|
||||
let bytes = borsh::to_vec(&ak).unwrap();
|
||||
assert!(!bytes.is_empty());
|
||||
// First 8 bytes = nonce, then discriminant 0 for FunctionCall
|
||||
assert_eq!(bytes[8], 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gas_constants() {
|
||||
assert_eq!(TGAS, 1_000_000_000_000);
|
||||
assert_eq!(MAX_GAS, 300_000_000_000_000);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,563 @@
|
||||
//! Core types for NEAR key management.
|
||||
//!
|
||||
//! Types for account IDs, public keys, access key permissions, network selection,
|
||||
//! and key metadata. All types validate on construction to prevent invalid states.
|
||||
//!
|
||||
//! SECURITY: Debug impls on key-related types MUST redact secret material.
|
||||
|
||||
use std::fmt;
|
||||
use std::str::FromStr;
|
||||
|
||||
use borsh::BorshSerialize;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::keys::KeyError;
|
||||
|
||||
/// NEAR account ID with validation.
|
||||
///
|
||||
/// Rules: 2-64 chars, lowercase alphanumeric + `.`, `-`, `_`.
|
||||
/// No leading/trailing separators, no consecutive separators.
|
||||
#[derive(Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub struct NearAccountId(String);
|
||||
|
||||
impl NearAccountId {
|
||||
pub fn new(id: &str) -> Result<Self, KeyError> {
|
||||
Self::validate(id)?;
|
||||
Ok(Self(id.to_string()))
|
||||
}
|
||||
|
||||
fn validate(id: &str) -> Result<(), KeyError> {
|
||||
if id.len() < 2 || id.len() > 64 {
|
||||
return Err(KeyError::InvalidAccountId {
|
||||
reason: format!("account ID must be 2-64 characters, got {}", id.len()),
|
||||
});
|
||||
}
|
||||
|
||||
let bytes = id.as_bytes();
|
||||
|
||||
// No leading/trailing separators
|
||||
if matches!(bytes[0], b'.' | b'-' | b'_') {
|
||||
return Err(KeyError::InvalidAccountId {
|
||||
reason: "account ID must not start with a separator".to_string(),
|
||||
});
|
||||
}
|
||||
if matches!(bytes[bytes.len() - 1], b'.' | b'-' | b'_') {
|
||||
return Err(KeyError::InvalidAccountId {
|
||||
reason: "account ID must not end with a separator".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
for ch in id.chars() {
|
||||
if !matches!(ch, 'a'..='z' | '0'..='9' | '.' | '-' | '_') {
|
||||
return Err(KeyError::InvalidAccountId {
|
||||
reason: format!("invalid character '{}' in account ID", ch),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for NearAccountId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}", self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for NearAccountId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "NearAccountId({})", self.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for NearAccountId {
|
||||
type Err = KeyError;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
Self::new(s)
|
||||
}
|
||||
}
|
||||
|
||||
impl BorshSerialize for NearAccountId {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
// NEAR protocol serializes account IDs as length-prefixed UTF-8 strings.
|
||||
BorshSerialize::serialize(&self.0, writer)
|
||||
}
|
||||
}
|
||||
|
||||
/// Key type discriminant for borsh serialization (matches NEAR protocol).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum KeyType {
|
||||
Ed25519 = 0,
|
||||
}
|
||||
|
||||
impl BorshSerialize for KeyType {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
BorshSerialize::serialize(&(*self as u8), writer)
|
||||
}
|
||||
}
|
||||
|
||||
/// NEAR public key with format parsing.
|
||||
///
|
||||
/// Parses the NEAR format: `ed25519:<base58-encoded-32-bytes>`
|
||||
#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct NearPublicKey {
|
||||
pub key_type: KeyType,
|
||||
pub data: [u8; 32],
|
||||
}
|
||||
|
||||
impl NearPublicKey {
|
||||
/// Parse from NEAR format string: `ed25519:<base58>`
|
||||
pub fn from_near_format(s: &str) -> Result<Self, KeyError> {
|
||||
let s = s.trim();
|
||||
let data_str = s
|
||||
.strip_prefix("ed25519:")
|
||||
.ok_or_else(|| KeyError::InvalidKeyFormat {
|
||||
reason: "public key must start with 'ed25519:'".to_string(),
|
||||
})?;
|
||||
|
||||
let bytes = bs58::decode(data_str)
|
||||
.into_vec()
|
||||
.map_err(|e| KeyError::InvalidKeyFormat {
|
||||
reason: format!("invalid base58 in public key: {}", e),
|
||||
})?;
|
||||
|
||||
if bytes.len() != 32 {
|
||||
return Err(KeyError::InvalidKeyFormat {
|
||||
reason: format!("ed25519 public key must be 32 bytes, got {}", bytes.len()),
|
||||
});
|
||||
}
|
||||
|
||||
let mut data = [0u8; 32];
|
||||
data.copy_from_slice(&bytes);
|
||||
|
||||
Ok(Self {
|
||||
key_type: KeyType::Ed25519,
|
||||
data,
|
||||
})
|
||||
}
|
||||
|
||||
/// Format as NEAR string: `ed25519:<base58>`
|
||||
pub fn to_near_format(&self) -> String {
|
||||
format!("ed25519:{}", bs58::encode(&self.data).into_string())
|
||||
}
|
||||
|
||||
/// Raw 32-byte key data.
|
||||
pub fn as_bytes(&self) -> &[u8; 32] {
|
||||
&self.data
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for NearPublicKey {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
write!(f, "{}", self.to_near_format())
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for NearPublicKey {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let encoded = bs58::encode(&self.data).into_string();
|
||||
let preview = if encoded.len() > 8 {
|
||||
&encoded[..8]
|
||||
} else {
|
||||
&encoded
|
||||
};
|
||||
write!(f, "NearPublicKey(ed25519:{}...)", preview)
|
||||
}
|
||||
}
|
||||
|
||||
impl BorshSerialize for NearPublicKey {
|
||||
fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
|
||||
// NEAR protocol: key_type byte + 32 bytes of key data
|
||||
BorshSerialize::serialize(&self.key_type, writer)?;
|
||||
writer.write_all(&self.data)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Access key permission level.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum AccessKeyPermission {
|
||||
FullAccess,
|
||||
FunctionCall {
|
||||
/// Max NEAR that can be spent (None = unlimited within key's scope).
|
||||
allowance: Option<u128>,
|
||||
/// Contract this key is scoped to.
|
||||
receiver_id: String,
|
||||
/// Allowed method names (empty = all methods on the contract).
|
||||
method_names: Vec<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl fmt::Display for AccessKeyPermission {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
AccessKeyPermission::FullAccess => write!(f, "FullAccess"),
|
||||
AccessKeyPermission::FunctionCall {
|
||||
receiver_id,
|
||||
method_names,
|
||||
allowance,
|
||||
} => {
|
||||
write!(f, "FunctionCall({}", receiver_id)?;
|
||||
if !method_names.is_empty() {
|
||||
write!(f, "::{}", method_names.join(","))?;
|
||||
}
|
||||
if let Some(a) = allowance {
|
||||
write!(f, ", allowance={})", format_yocto(*a))?;
|
||||
} else {
|
||||
write!(f, ")")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// NEAR network configuration.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum NearNetwork {
|
||||
Mainnet,
|
||||
Testnet,
|
||||
Custom(String),
|
||||
}
|
||||
|
||||
impl NearNetwork {
|
||||
pub fn rpc_url(&self) -> &str {
|
||||
match self {
|
||||
NearNetwork::Mainnet => "https://rpc.mainnet.near.org",
|
||||
NearNetwork::Testnet => "https://rpc.testnet.near.org",
|
||||
NearNetwork::Custom(url) => url.as_str(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for NearNetwork {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
NearNetwork::Mainnet => write!(f, "mainnet"),
|
||||
NearNetwork::Testnet => write!(f, "testnet"),
|
||||
NearNetwork::Custom(url) => write!(f, "custom({})", url),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl FromStr for NearNetwork {
|
||||
type Err = KeyError;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"mainnet" => Ok(NearNetwork::Mainnet),
|
||||
"testnet" => Ok(NearNetwork::Testnet),
|
||||
url if url.starts_with("http") => Ok(NearNetwork::Custom(url.to_string())),
|
||||
other => Err(KeyError::InvalidKeyFormat {
|
||||
reason: format!(
|
||||
"unknown network '{}', expected mainnet, testnet, or an RPC URL",
|
||||
other
|
||||
),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Metadata for a stored key (public info only, no secrets).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct KeyMetadata {
|
||||
pub label: String,
|
||||
pub account_id: String,
|
||||
pub public_key: String,
|
||||
pub permission: AccessKeyPermission,
|
||||
pub network: NearNetwork,
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// Cached nonce for transaction building (avoids extra RPC round-trip).
|
||||
pub cached_nonce: Option<u64>,
|
||||
}
|
||||
|
||||
/// Top-level structure for ~/.ironclaw/keys.json
|
||||
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
|
||||
pub struct KeyStore {
|
||||
pub keys: std::collections::HashMap<String, KeyMetadata>,
|
||||
pub last_backup_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// Format yoctoNEAR as human-readable NEAR amount.
|
||||
pub fn format_yocto(yocto: u128) -> String {
|
||||
const ONE_NEAR: u128 = 1_000_000_000_000_000_000_000_000;
|
||||
const ONE_MILLI_NEAR: u128 = ONE_NEAR / 1000;
|
||||
if yocto == 0 {
|
||||
return "0 NEAR".to_string();
|
||||
}
|
||||
if yocto >= ONE_MILLI_NEAR {
|
||||
let whole = yocto / ONE_NEAR;
|
||||
let frac = (yocto % ONE_NEAR) / ONE_MILLI_NEAR; // 3 decimal places
|
||||
if frac == 0 {
|
||||
format!("{} NEAR", whole)
|
||||
} else {
|
||||
format!("{}.{:03} NEAR", whole, frac)
|
||||
}
|
||||
} else {
|
||||
format!("{} yoctoNEAR", yocto)
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a NEAR amount string into yoctoNEAR.
|
||||
///
|
||||
/// Accepts: "1", "0.5", "1.5 NEAR", "100000 yoctoNEAR"
|
||||
pub fn parse_near_amount(s: &str) -> Result<u128, KeyError> {
|
||||
const ONE_NEAR: u128 = 1_000_000_000_000_000_000_000_000;
|
||||
|
||||
let s = s.trim();
|
||||
|
||||
// Check for explicit yoctoNEAR suffix
|
||||
if let Some(yocto_str) = s
|
||||
.strip_suffix("yoctoNEAR")
|
||||
.or_else(|| s.strip_suffix("yocto"))
|
||||
{
|
||||
return yocto_str
|
||||
.trim()
|
||||
.parse::<u128>()
|
||||
.map_err(|e| KeyError::InvalidKeyFormat {
|
||||
reason: format!("invalid yoctoNEAR amount: {}", e),
|
||||
});
|
||||
}
|
||||
|
||||
// Strip optional "NEAR" suffix
|
||||
let amount_str = s
|
||||
.strip_suffix("NEAR")
|
||||
.or_else(|| s.strip_suffix("near"))
|
||||
.unwrap_or(s)
|
||||
.trim();
|
||||
|
||||
// Parse as decimal NEAR
|
||||
if let Some((whole_str, frac_str)) = amount_str.split_once('.') {
|
||||
let whole: u128 = whole_str.parse().map_err(|e| KeyError::InvalidKeyFormat {
|
||||
reason: format!("invalid NEAR amount: {}", e),
|
||||
})?;
|
||||
|
||||
// Pad or truncate fractional part to 24 digits
|
||||
let mut frac_padded = frac_str.to_string();
|
||||
if frac_padded.len() > 24 {
|
||||
frac_padded.truncate(24);
|
||||
}
|
||||
while frac_padded.len() < 24 {
|
||||
frac_padded.push('0');
|
||||
}
|
||||
|
||||
let frac: u128 = frac_padded
|
||||
.parse()
|
||||
.map_err(|e| KeyError::InvalidKeyFormat {
|
||||
reason: format!("invalid NEAR fractional amount: {}", e),
|
||||
})?;
|
||||
|
||||
Ok(whole * ONE_NEAR + frac)
|
||||
} else {
|
||||
let whole: u128 = amount_str.parse().map_err(|e| KeyError::InvalidKeyFormat {
|
||||
reason: format!("invalid NEAR amount: {}", e),
|
||||
})?;
|
||||
Ok(whole * ONE_NEAR)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::keys::types::{
|
||||
AccessKeyPermission, KeyType, NearAccountId, NearNetwork, NearPublicKey, format_yocto,
|
||||
parse_near_amount,
|
||||
};
|
||||
|
||||
// -- NearAccountId tests --
|
||||
|
||||
#[test]
|
||||
fn test_valid_account_ids() {
|
||||
assert!(NearAccountId::new("alice.near").is_ok());
|
||||
assert!(NearAccountId::new("bob.testnet").is_ok());
|
||||
assert!(NearAccountId::new("system").is_ok());
|
||||
assert!(NearAccountId::new("ab").is_ok()); // minimum 2 chars
|
||||
assert!(NearAccountId::new("a0").is_ok());
|
||||
assert!(NearAccountId::new("alice-bob.near").is_ok());
|
||||
assert!(NearAccountId::new("alice_bob.near").is_ok());
|
||||
// 64 chars max
|
||||
let long_id = "a".repeat(64);
|
||||
assert!(NearAccountId::new(&long_id).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_invalid_account_ids() {
|
||||
// Too short
|
||||
assert!(NearAccountId::new("a").is_err());
|
||||
// Too long
|
||||
assert!(NearAccountId::new(&"a".repeat(65)).is_err());
|
||||
// Uppercase
|
||||
assert!(NearAccountId::new("Alice.near").is_err());
|
||||
// Leading separator
|
||||
assert!(NearAccountId::new(".alice").is_err());
|
||||
assert!(NearAccountId::new("-alice").is_err());
|
||||
// Trailing separator
|
||||
assert!(NearAccountId::new("alice.").is_err());
|
||||
// Invalid chars
|
||||
assert!(NearAccountId::new("alice@near").is_err());
|
||||
assert!(NearAccountId::new("alice near").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_account_id_display() {
|
||||
let id = NearAccountId::new("alice.near").unwrap();
|
||||
assert_eq!(id.to_string(), "alice.near");
|
||||
assert_eq!(id.as_str(), "alice.near");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_account_id_from_str() {
|
||||
let id: NearAccountId = "bob.testnet".parse().unwrap();
|
||||
assert_eq!(id.as_str(), "bob.testnet");
|
||||
}
|
||||
|
||||
// -- NearPublicKey tests --
|
||||
|
||||
#[test]
|
||||
fn test_public_key_roundtrip() {
|
||||
let key_str = "ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp";
|
||||
let key = NearPublicKey::from_near_format(key_str).unwrap();
|
||||
assert_eq!(key.key_type, KeyType::Ed25519);
|
||||
assert_eq!(key.to_near_format(), key_str);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_public_key_invalid_prefix() {
|
||||
assert!(NearPublicKey::from_near_format("secp256k1:abc").is_err());
|
||||
assert!(NearPublicKey::from_near_format("abc123").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_public_key_invalid_base58() {
|
||||
assert!(NearPublicKey::from_near_format("ed25519:not-valid-base58!!!").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_public_key_wrong_length() {
|
||||
// Too short (only 16 bytes encoded)
|
||||
assert!(NearPublicKey::from_near_format("ed25519:3gZNbFLLDt").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_public_key_debug_redacts() {
|
||||
let key_str = "ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp";
|
||||
let key = NearPublicKey::from_near_format(key_str).unwrap();
|
||||
let debug = format!("{:?}", key);
|
||||
// Should show first 8 chars of base58, not the whole thing
|
||||
assert!(debug.contains("..."));
|
||||
assert!(!debug.contains("6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp"));
|
||||
}
|
||||
|
||||
// -- AccessKeyPermission tests --
|
||||
|
||||
#[test]
|
||||
fn test_permission_display() {
|
||||
assert_eq!(AccessKeyPermission::FullAccess.to_string(), "FullAccess");
|
||||
|
||||
let fc = AccessKeyPermission::FunctionCall {
|
||||
allowance: None,
|
||||
receiver_id: "intents.near".to_string(),
|
||||
method_names: vec![],
|
||||
};
|
||||
assert_eq!(fc.to_string(), "FunctionCall(intents.near)");
|
||||
|
||||
let fc_methods = AccessKeyPermission::FunctionCall {
|
||||
allowance: Some(1_000_000_000_000_000_000_000_000),
|
||||
receiver_id: "contract.near".to_string(),
|
||||
method_names: vec!["deposit".to_string(), "withdraw".to_string()],
|
||||
};
|
||||
assert!(fc_methods.to_string().contains("deposit,withdraw"));
|
||||
assert!(fc_methods.to_string().contains("1 NEAR"));
|
||||
}
|
||||
|
||||
// -- NearNetwork tests --
|
||||
|
||||
#[test]
|
||||
fn test_network_rpc_urls() {
|
||||
assert_eq!(
|
||||
NearNetwork::Mainnet.rpc_url(),
|
||||
"https://rpc.mainnet.near.org"
|
||||
);
|
||||
assert_eq!(
|
||||
NearNetwork::Testnet.rpc_url(),
|
||||
"https://rpc.testnet.near.org"
|
||||
);
|
||||
let custom = NearNetwork::Custom("https://custom.rpc.dev".to_string());
|
||||
assert_eq!(custom.rpc_url(), "https://custom.rpc.dev");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_network_from_str() {
|
||||
assert_eq!(
|
||||
"mainnet".parse::<NearNetwork>().unwrap(),
|
||||
NearNetwork::Mainnet
|
||||
);
|
||||
assert_eq!(
|
||||
"testnet".parse::<NearNetwork>().unwrap(),
|
||||
NearNetwork::Testnet
|
||||
);
|
||||
assert_eq!(
|
||||
"https://custom.rpc".parse::<NearNetwork>().unwrap(),
|
||||
NearNetwork::Custom("https://custom.rpc".to_string())
|
||||
);
|
||||
assert!("garbage".parse::<NearNetwork>().is_err());
|
||||
}
|
||||
|
||||
// -- NEAR amount formatting/parsing --
|
||||
|
||||
#[test]
|
||||
fn test_format_yocto() {
|
||||
assert_eq!(format_yocto(0), "0 NEAR");
|
||||
assert_eq!(format_yocto(1_000_000_000_000_000_000_000_000), "1 NEAR");
|
||||
assert_eq!(
|
||||
format_yocto(5_500_000_000_000_000_000_000_000),
|
||||
"5.500 NEAR"
|
||||
);
|
||||
assert_eq!(format_yocto(1), "1 yoctoNEAR");
|
||||
assert_eq!(format_yocto(500_000_000_000_000_000_000_000), "0.500 NEAR");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_near_amount() {
|
||||
assert_eq!(
|
||||
parse_near_amount("1").unwrap(),
|
||||
1_000_000_000_000_000_000_000_000
|
||||
);
|
||||
assert_eq!(
|
||||
parse_near_amount("0.5").unwrap(),
|
||||
500_000_000_000_000_000_000_000
|
||||
);
|
||||
assert_eq!(
|
||||
parse_near_amount("1.5 NEAR").unwrap(),
|
||||
1_500_000_000_000_000_000_000_000
|
||||
);
|
||||
assert_eq!(parse_near_amount("100 yoctoNEAR").unwrap(), 100);
|
||||
assert_eq!(parse_near_amount("0").unwrap(), 0);
|
||||
}
|
||||
|
||||
// -- Borsh serialization tests --
|
||||
|
||||
#[test]
|
||||
fn test_account_id_borsh() {
|
||||
let id = NearAccountId::new("alice.near").unwrap();
|
||||
let bytes = borsh::to_vec(&id).unwrap();
|
||||
// Length-prefixed string: 4 bytes length + 10 bytes "alice.near"
|
||||
assert_eq!(bytes.len(), 4 + 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_public_key_borsh() {
|
||||
let key_str = "ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp";
|
||||
let key = NearPublicKey::from_near_format(key_str).unwrap();
|
||||
let bytes = borsh::to_vec(&key).unwrap();
|
||||
// 1 byte key_type + 32 bytes data
|
||||
assert_eq!(bytes.len(), 33);
|
||||
assert_eq!(bytes[0], 0); // Ed25519 = 0
|
||||
}
|
||||
}
|
||||
@@ -48,6 +48,7 @@ pub mod estimation;
|
||||
pub mod evaluation;
|
||||
pub mod extensions;
|
||||
pub mod history;
|
||||
pub mod keys;
|
||||
pub mod llm;
|
||||
pub mod safety;
|
||||
pub mod sandbox;
|
||||
|
||||
+1
-1
@@ -13,7 +13,7 @@ pub mod session;
|
||||
pub use nearai::{ModelInfo, NearAiProvider};
|
||||
pub use nearai_chat::NearAiChatProvider;
|
||||
pub use provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
|
||||
ChatMessage, CompletionRequest, CompletionResponse, LlmProvider, Role, ToolCall,
|
||||
ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
|
||||
};
|
||||
pub use reasoning::{ActionPlan, Reasoning, ReasoningContext, RespondResult, ToolSelection};
|
||||
|
||||
+19
-1
@@ -17,12 +17,14 @@ use ironclaw::{
|
||||
web::log_layer::{LogBroadcaster, WebLogLayer},
|
||||
},
|
||||
cli::{
|
||||
Cli, Command, run_mcp_command, run_memory_command, run_status_command, run_tool_command,
|
||||
Cli, Command, run_key_command, run_mcp_command, run_memory_command, run_status_command,
|
||||
run_tool_command,
|
||||
},
|
||||
config::Config,
|
||||
context::ContextManager,
|
||||
extensions::ExtensionManager,
|
||||
history::Store,
|
||||
keys::KeyManager,
|
||||
llm::{SessionConfig, create_llm_provider, create_session_manager},
|
||||
safety::SafetyLayer,
|
||||
secrets::{PostgresSecretsStore, SecretsCrypto, SecretsStore},
|
||||
@@ -52,6 +54,16 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
return run_tool_command(tool_cmd.clone()).await;
|
||||
}
|
||||
Some(Command::Key(key_cmd)) => {
|
||||
let _ = dotenvy::dotenv();
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")),
|
||||
)
|
||||
.init();
|
||||
|
||||
return run_key_command(key_cmd.clone()).await;
|
||||
}
|
||||
Some(Command::Config(config_cmd)) => {
|
||||
// Config commands don't need logging setup
|
||||
return ironclaw::cli::run_config_command(config_cmd.clone())
|
||||
@@ -318,6 +330,11 @@ async fn main() -> anyhow::Result<()> {
|
||||
None
|
||||
};
|
||||
|
||||
// Create key manager if secrets store is available.
|
||||
let key_manager: Option<Arc<KeyManager>> = secrets_store
|
||||
.as_ref()
|
||||
.map(|store| Arc::new(KeyManager::new(Arc::clone(store), "default".to_string())));
|
||||
|
||||
let mcp_session_manager = Arc::new(McpSessionManager::new());
|
||||
|
||||
// Create WASM tool runtime (sync, just builds the wasmtime engine)
|
||||
@@ -746,6 +763,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
tools,
|
||||
workspace,
|
||||
extension_manager,
|
||||
key_manager,
|
||||
};
|
||||
let agent = Agent::new(
|
||||
config.agent.clone(),
|
||||
|
||||
@@ -511,6 +511,14 @@ fn default_patterns() -> Vec<LeakPattern> {
|
||||
severity: LeakSeverity::High,
|
||||
action: LeakAction::Redact,
|
||||
},
|
||||
// NEAR ed25519 private keys (base58 encoded, ~88 chars after prefix).
|
||||
// Public keys are shorter (~44 chars), so this pattern is specific to secrets.
|
||||
LeakPattern {
|
||||
name: "near_ed25519_secret_key".to_string(),
|
||||
regex: Regex::new(r"ed25519:[1-9A-HJ-NP-Za-km-z]{80,90}").unwrap(),
|
||||
severity: LeakSeverity::Critical,
|
||||
action: LeakAction::Block,
|
||||
},
|
||||
// High entropy hex (potential secrets, warn only)
|
||||
// Uses word boundary since look-around isn't supported in the regex crate.
|
||||
// This catches standalone 64-char hex strings (like SHA256 hashes used as secrets).
|
||||
@@ -696,6 +704,39 @@ mod tests {
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_detect_near_ed25519_secret_key() {
|
||||
let detector = LeakDetector::new();
|
||||
// A realistic NEAR secret key (88 base58 chars after prefix)
|
||||
let content = "key: ed25519:3D4YudUahN1nawWogh9MFV2MXJBMHCS2RE1KU7rWAiMi3t12UiSnMYCJ7BFXbsFhKfNUWDj8CCEbifTByREAMkTi";
|
||||
|
||||
let result = detector.scan(content);
|
||||
assert!(!result.is_clean());
|
||||
assert!(result.should_block);
|
||||
assert!(
|
||||
result
|
||||
.matches
|
||||
.iter()
|
||||
.any(|m| m.pattern_name == "near_ed25519_secret_key")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_near_public_key_not_blocked() {
|
||||
let detector = LeakDetector::new();
|
||||
// Public keys are ~44 base58 chars, should NOT match the 80-90 char pattern
|
||||
let content = "pubkey: ed25519:6E8sCci9badyRkXb3JoRpBj5p8C6Tw41ELDZoiihKEtp";
|
||||
|
||||
let result = detector.scan(content);
|
||||
// Should not match near_ed25519_secret_key pattern
|
||||
assert!(
|
||||
!result
|
||||
.matches
|
||||
.iter()
|
||||
.any(|m| m.pattern_name == "near_ed25519_secret_key")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_scan_http_request_blocks_secret_in_body() {
|
||||
let detector = LeakDetector::new();
|
||||
|
||||
@@ -32,6 +32,8 @@ pub struct Capabilities {
|
||||
pub tool_invoke: Option<ToolInvokeCapability>,
|
||||
/// Check if secrets exist.
|
||||
pub secrets: Option<SecretsCapability>,
|
||||
/// Sign payloads using managed NEAR keys.
|
||||
pub signing: Option<SigningCapability>,
|
||||
}
|
||||
|
||||
impl Capabilities {
|
||||
@@ -71,6 +73,21 @@ impl Capabilities {
|
||||
});
|
||||
self
|
||||
}
|
||||
|
||||
/// Enable payload signing with the given key labels.
|
||||
pub fn with_signing(
|
||||
mut self,
|
||||
allowed_labels: Vec<String>,
|
||||
max_signs: u32,
|
||||
signer: Option<Arc<dyn PayloadSigner>>,
|
||||
) -> Self {
|
||||
self.signing = Some(SigningCapability {
|
||||
allowed_key_labels: allowed_labels,
|
||||
max_signs_per_execution: max_signs,
|
||||
signer,
|
||||
});
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
/// Workspace read capability configuration.
|
||||
@@ -301,6 +318,68 @@ impl SecretsCapability {
|
||||
}
|
||||
}
|
||||
|
||||
/// Signing capability: allows WASM tools to request payload signatures from managed keys.
|
||||
///
|
||||
/// The private keys NEVER enter WASM memory. The host performs the signing and
|
||||
/// returns only the signature bytes.
|
||||
#[derive(Clone)]
|
||||
pub struct SigningCapability {
|
||||
/// Key labels this tool is allowed to use for signing.
|
||||
pub allowed_key_labels: Vec<String>,
|
||||
/// Maximum number of sign operations per execution.
|
||||
pub max_signs_per_execution: u32,
|
||||
/// Implementation that performs the actual signing.
|
||||
/// Injected at runtime. None means signing will always return an error.
|
||||
pub signer: Option<Arc<dyn PayloadSigner>>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for SigningCapability {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("SigningCapability")
|
||||
.field("allowed_key_labels", &self.allowed_key_labels)
|
||||
.field("max_signs_per_execution", &self.max_signs_per_execution)
|
||||
.field("signer", &self.signer.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
impl SigningCapability {
|
||||
/// Check if a key label is allowed.
|
||||
pub fn is_label_allowed(&self, label: &str) -> bool {
|
||||
self.allowed_key_labels.iter().any(|l| l == label)
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a payload signing operation.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SignPayloadResult {
|
||||
/// Base64-encoded signature (set on success).
|
||||
pub signature: Option<String>,
|
||||
/// Error message (set on failure).
|
||||
pub error: Option<String>,
|
||||
/// Whether user approval is needed before signing can proceed.
|
||||
pub approval_pending: bool,
|
||||
}
|
||||
|
||||
/// Trait for performing payload signing from the host boundary.
|
||||
///
|
||||
/// This is intentionally synchronous because WASM host functions run in a
|
||||
/// blocking context. Implementations that need async should use
|
||||
/// `Handle::block_on()` internally.
|
||||
pub trait PayloadSigner: Send + Sync {
|
||||
/// Sign a payload using the specified key.
|
||||
///
|
||||
/// The payload is raw bytes (decoded from the base64 the WASM tool sent).
|
||||
/// Returns a `SignPayloadResult` which may contain a signature, an error,
|
||||
/// or an approval-pending flag.
|
||||
fn sign_payload(
|
||||
&self,
|
||||
key_label: &str,
|
||||
payload: &[u8],
|
||||
context_json: &str,
|
||||
) -> SignPayloadResult;
|
||||
}
|
||||
|
||||
/// Rate limiting configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RateLimitConfig {
|
||||
|
||||
@@ -35,7 +35,7 @@ use serde::{Deserialize, Serialize};
|
||||
use crate::secrets::{CredentialLocation, CredentialMapping};
|
||||
use crate::tools::wasm::{
|
||||
Capabilities, EndpointPattern, HttpCapability, RateLimitConfig, SecretsCapability,
|
||||
ToolInvokeCapability, WorkspaceCapability,
|
||||
SigningCapability, ToolInvokeCapability, WorkspaceCapability,
|
||||
};
|
||||
|
||||
/// Root schema for a capabilities JSON file.
|
||||
@@ -57,6 +57,10 @@ pub struct CapabilitiesFile {
|
||||
#[serde(default)]
|
||||
pub workspace: Option<WorkspaceCapabilitySchema>,
|
||||
|
||||
/// Payload signing using managed NEAR keys.
|
||||
#[serde(default)]
|
||||
pub signing: Option<SigningCapabilitySchema>,
|
||||
|
||||
/// Authentication setup instructions.
|
||||
/// Used by `ironclaw config` to guide users through auth setup.
|
||||
#[serde(default)]
|
||||
@@ -106,6 +110,14 @@ impl CapabilitiesFile {
|
||||
});
|
||||
}
|
||||
|
||||
if let Some(signing) = &self.signing {
|
||||
caps.signing = Some(SigningCapability {
|
||||
allowed_key_labels: signing.allowed_key_labels.clone(),
|
||||
max_signs_per_execution: signing.max_signs_per_execution.unwrap_or(5),
|
||||
signer: None, // Injected at runtime
|
||||
});
|
||||
}
|
||||
|
||||
caps
|
||||
}
|
||||
}
|
||||
@@ -318,6 +330,32 @@ pub struct ToolInvokeCapabilitySchema {
|
||||
pub rate_limit: Option<RateLimitSchema>,
|
||||
}
|
||||
|
||||
/// Signing capability schema.
|
||||
///
|
||||
/// Allows WASM tools to request payload signatures from managed NEAR keys.
|
||||
/// The private keys never enter WASM memory.
|
||||
///
|
||||
/// # Example
|
||||
///
|
||||
/// ```json
|
||||
/// {
|
||||
/// "signing": {
|
||||
/// "allowed_key_labels": ["intents-signer"],
|
||||
/// "max_signs_per_execution": 5
|
||||
/// }
|
||||
/// }
|
||||
/// ```
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SigningCapabilitySchema {
|
||||
/// Key labels this tool is allowed to use for signing.
|
||||
#[serde(default)]
|
||||
pub allowed_key_labels: Vec<String>,
|
||||
|
||||
/// Maximum sign operations per execution (default: 5).
|
||||
#[serde(default)]
|
||||
pub max_signs_per_execution: Option<u32>,
|
||||
}
|
||||
|
||||
/// Workspace read capability schema.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct WorkspaceCapabilitySchema {
|
||||
@@ -754,4 +792,60 @@ mod tests {
|
||||
assert!(auth.display_name.is_none());
|
||||
assert!(auth.setup_url.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_signing_capability() {
|
||||
let json = r#"{
|
||||
"signing": {
|
||||
"allowed_key_labels": ["intents-signer", "trading-key"],
|
||||
"max_signs_per_execution": 10
|
||||
}
|
||||
}"#;
|
||||
|
||||
let caps = CapabilitiesFile::from_json(json).unwrap();
|
||||
let signing = caps.signing.unwrap();
|
||||
assert_eq!(
|
||||
signing.allowed_key_labels,
|
||||
vec!["intents-signer", "trading-key"]
|
||||
);
|
||||
assert_eq!(signing.max_signs_per_execution, Some(10));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_signing_defaults() {
|
||||
let json = r#"{
|
||||
"signing": {
|
||||
"allowed_key_labels": ["default"]
|
||||
}
|
||||
}"#;
|
||||
|
||||
let caps = CapabilitiesFile::from_json(json).unwrap();
|
||||
let signing = caps.signing.as_ref().unwrap();
|
||||
assert_eq!(signing.max_signs_per_execution, None);
|
||||
|
||||
// Should default to 5 when converted
|
||||
let runtime_caps = caps.to_capabilities();
|
||||
let runtime_signing = runtime_caps.signing.unwrap();
|
||||
assert_eq!(runtime_signing.max_signs_per_execution, 5);
|
||||
assert!(runtime_signing.is_label_allowed("default"));
|
||||
assert!(!runtime_signing.is_label_allowed("other"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_signing_to_capabilities() {
|
||||
let json = r#"{
|
||||
"signing": {
|
||||
"allowed_key_labels": ["signer-1"],
|
||||
"max_signs_per_execution": 3
|
||||
}
|
||||
}"#;
|
||||
|
||||
let file = CapabilitiesFile::from_json(json).unwrap();
|
||||
let caps = file.to_capabilities();
|
||||
|
||||
let signing = caps.signing.unwrap();
|
||||
assert_eq!(signing.allowed_key_labels, vec!["signer-1"]);
|
||||
assert_eq!(signing.max_signs_per_execution, 3);
|
||||
assert!(signing.signer.is_none()); // Injected at runtime
|
||||
}
|
||||
}
|
||||
|
||||
@@ -67,6 +67,7 @@ pub struct LogEntry {
|
||||
///
|
||||
/// This is the "VMLogic" equivalent, it tracks all side effects and enforces limits.
|
||||
/// Extended in V2 to support HTTP requests, tool invocation, and secret checks.
|
||||
/// Extended in V3 to support payload signing via managed NEAR keys.
|
||||
pub struct HostState {
|
||||
/// Collected log entries.
|
||||
logs: Vec<LogEntry>,
|
||||
@@ -82,6 +83,8 @@ pub struct HostState {
|
||||
http_request_count: u32,
|
||||
/// Tool invoke count for rate limiting within this execution.
|
||||
tool_invoke_count: u32,
|
||||
/// Signing request count for rate limiting within this execution.
|
||||
sign_count: u32,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for HostState {
|
||||
@@ -93,6 +96,7 @@ impl std::fmt::Debug for HostState {
|
||||
.field("user_id", &self.user_id)
|
||||
.field("http_request_count", &self.http_request_count)
|
||||
.field("tool_invoke_count", &self.tool_invoke_count)
|
||||
.field("sign_count", &self.sign_count)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -108,6 +112,7 @@ impl HostState {
|
||||
user_id: None,
|
||||
http_request_count: 0,
|
||||
tool_invoke_count: 0,
|
||||
sign_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -121,6 +126,7 @@ impl HostState {
|
||||
user_id: Some(user_id.into()),
|
||||
http_request_count: 0,
|
||||
tool_invoke_count: 0,
|
||||
sign_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -223,6 +229,87 @@ impl HostState {
|
||||
}
|
||||
}
|
||||
|
||||
/// Sign a payload using a managed NEAR key.
|
||||
///
|
||||
/// Checks signing capability, key label allowlist, and rate limit.
|
||||
/// Delegates actual signing to the `PayloadSigner` if all checks pass.
|
||||
///
|
||||
/// Private keys NEVER enter WASM memory. Only the signature is returned.
|
||||
pub fn sign_payload(
|
||||
&mut self,
|
||||
key_label: &str,
|
||||
payload_base64: &str,
|
||||
context_json: &str,
|
||||
) -> crate::tools::wasm::capabilities::SignPayloadResult {
|
||||
use crate::tools::wasm::capabilities::SignPayloadResult;
|
||||
|
||||
let capability = match &self.capabilities.signing {
|
||||
Some(cap) => cap,
|
||||
None => {
|
||||
return SignPayloadResult {
|
||||
signature: None,
|
||||
error: Some("Signing capability not granted".to_string()),
|
||||
approval_pending: false,
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
// Check key label is allowed
|
||||
if !capability.is_label_allowed(key_label) {
|
||||
return SignPayloadResult {
|
||||
signature: None,
|
||||
error: Some(format!(
|
||||
"Key label '{}' not in allowed list for this tool",
|
||||
key_label
|
||||
)),
|
||||
approval_pending: false,
|
||||
};
|
||||
}
|
||||
|
||||
// Check rate limit
|
||||
self.sign_count += 1;
|
||||
if self.sign_count > capability.max_signs_per_execution {
|
||||
return SignPayloadResult {
|
||||
signature: None,
|
||||
error: Some(format!(
|
||||
"Sign limit exceeded ({} per execution)",
|
||||
capability.max_signs_per_execution
|
||||
)),
|
||||
approval_pending: false,
|
||||
};
|
||||
}
|
||||
|
||||
// Decode base64 payload
|
||||
let payload_bytes = match base64::Engine::decode(
|
||||
&base64::engine::general_purpose::STANDARD,
|
||||
payload_base64,
|
||||
) {
|
||||
Ok(bytes) => bytes,
|
||||
Err(e) => {
|
||||
return SignPayloadResult {
|
||||
signature: None,
|
||||
error: Some(format!("Invalid base64 payload: {}", e)),
|
||||
approval_pending: false,
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
// Delegate to signer implementation
|
||||
match &capability.signer {
|
||||
Some(signer) => signer.sign_payload(key_label, &payload_bytes, context_json),
|
||||
None => SignPayloadResult {
|
||||
signature: None,
|
||||
error: Some("No signing provider configured".to_string()),
|
||||
approval_pending: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the sign count for this execution.
|
||||
pub fn sign_count(&self) -> u32 {
|
||||
self.sign_count
|
||||
}
|
||||
|
||||
/// Get collected logs after execution.
|
||||
pub fn take_logs(&mut self) -> Vec<LogEntry> {
|
||||
std::mem::take(&mut self.logs)
|
||||
@@ -603,4 +690,107 @@ mod tests {
|
||||
let state = HostState::new_with_user(Capabilities::default(), "user123");
|
||||
assert_eq!(state.user_id(), Some("user123"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sign_payload_no_capability() {
|
||||
let mut state = HostState::minimal();
|
||||
let result = state.sign_payload("any-key", "AAAA", "{}");
|
||||
assert!(result.error.is_some());
|
||||
assert!(result.error.unwrap().contains("not granted"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sign_payload_label_not_allowed() {
|
||||
let capabilities = Capabilities {
|
||||
signing: Some(crate::tools::wasm::capabilities::SigningCapability {
|
||||
allowed_key_labels: vec!["allowed-key".to_string()],
|
||||
max_signs_per_execution: 5,
|
||||
signer: None,
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
let mut state = HostState::new(capabilities);
|
||||
let result = state.sign_payload("forbidden-key", "AAAA", "{}");
|
||||
assert!(result.error.is_some());
|
||||
assert!(result.error.unwrap().contains("not in allowed list"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sign_payload_rate_limit() {
|
||||
let capabilities = Capabilities {
|
||||
signing: Some(crate::tools::wasm::capabilities::SigningCapability {
|
||||
allowed_key_labels: vec!["key".to_string()],
|
||||
max_signs_per_execution: 2,
|
||||
signer: None,
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
let mut state = HostState::new(capabilities);
|
||||
|
||||
// First two should hit "no signer" (not rate limit)
|
||||
let r1 = state.sign_payload("key", "AAAA", "{}");
|
||||
assert!(r1.error.as_deref().unwrap().contains("No signing provider"));
|
||||
|
||||
let r2 = state.sign_payload("key", "AAAA", "{}");
|
||||
assert!(r2.error.as_deref().unwrap().contains("No signing provider"));
|
||||
|
||||
// Third should hit rate limit
|
||||
let r3 = state.sign_payload("key", "AAAA", "{}");
|
||||
assert!(r3.error.as_deref().unwrap().contains("limit exceeded"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sign_payload_invalid_base64() {
|
||||
let capabilities = Capabilities {
|
||||
signing: Some(crate::tools::wasm::capabilities::SigningCapability {
|
||||
allowed_key_labels: vec!["key".to_string()],
|
||||
max_signs_per_execution: 5,
|
||||
signer: Some(Arc::new(MockSigner)),
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
let mut state = HostState::new(capabilities);
|
||||
let result = state.sign_payload("key", "not-valid-base64!!!", "{}");
|
||||
assert!(result.error.is_some());
|
||||
assert!(result.error.unwrap().contains("Invalid base64"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sign_payload_with_mock_signer() {
|
||||
let capabilities = Capabilities {
|
||||
signing: Some(crate::tools::wasm::capabilities::SigningCapability {
|
||||
allowed_key_labels: vec!["test-key".to_string()],
|
||||
max_signs_per_execution: 5,
|
||||
signer: Some(Arc::new(MockSigner)),
|
||||
}),
|
||||
..Default::default()
|
||||
};
|
||||
let mut state = HostState::new(capabilities);
|
||||
|
||||
// Encode some payload as base64
|
||||
let payload =
|
||||
base64::Engine::encode(&base64::engine::general_purpose::STANDARD, b"sign this");
|
||||
let result = state.sign_payload("test-key", &payload, "{}");
|
||||
assert!(result.signature.is_some());
|
||||
assert!(result.error.is_none());
|
||||
assert!(!result.approval_pending);
|
||||
assert_eq!(result.signature.unwrap(), "mock-signature");
|
||||
}
|
||||
|
||||
struct MockSigner;
|
||||
|
||||
impl crate::tools::wasm::capabilities::PayloadSigner for MockSigner {
|
||||
fn sign_payload(
|
||||
&self,
|
||||
_key_label: &str,
|
||||
_payload: &[u8],
|
||||
_context_json: &str,
|
||||
) -> crate::tools::wasm::capabilities::SignPayloadResult {
|
||||
crate::tools::wasm::capabilities::SignPayloadResult {
|
||||
signature: Some("mock-signature".to_string()),
|
||||
error: None,
|
||||
approval_pending: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -98,8 +98,9 @@ pub use wrapper::WasmToolWrapper;
|
||||
|
||||
// Capabilities (V2)
|
||||
pub use capabilities::{
|
||||
Capabilities, EndpointPattern, HttpCapability, RateLimitConfig, SecretsCapability,
|
||||
ToolInvokeCapability, WorkspaceCapability, WorkspaceReader,
|
||||
Capabilities, EndpointPattern, HttpCapability, PayloadSigner, RateLimitConfig,
|
||||
SecretsCapability, SignPayloadResult, SigningCapability, ToolInvokeCapability,
|
||||
WorkspaceCapability, WorkspaceReader,
|
||||
};
|
||||
|
||||
// Security components (V2)
|
||||
@@ -120,5 +121,5 @@ pub use loader::{DiscoveredTool, LoadResults, WasmLoadError, WasmToolLoader, dis
|
||||
// Capabilities schema (for parsing *.capabilities.json files)
|
||||
pub use capabilities_schema::{
|
||||
AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema, RateLimitSchema,
|
||||
ValidationEndpointSchema,
|
||||
SigningCapabilitySchema, ValidationEndpointSchema,
|
||||
};
|
||||
|
||||
@@ -245,6 +245,26 @@ impl WasmToolWrapper {
|
||||
WasmError::ConfigError(format!("Failed to add workspace-read function: {}", e))
|
||||
})?;
|
||||
|
||||
// host.sign-payload(key-label, payload, context-json) -> sign-result
|
||||
// Returns a record { signature: option<string>, error: option<string>, approval-pending: bool }
|
||||
linker
|
||||
.root()
|
||||
.func_wrap(
|
||||
"sign-payload",
|
||||
|mut ctx: wasmtime::StoreContextMut<'_, StoreData>,
|
||||
(key_label, payload, context_json): (String, String, String)|
|
||||
-> anyhow::Result<(Option<String>, Option<String>, bool)> {
|
||||
let result =
|
||||
ctx.data_mut()
|
||||
.host_state
|
||||
.sign_payload(&key_label, &payload, &context_json);
|
||||
Ok((result.signature, result.error, result.approval_pending))
|
||||
},
|
||||
)
|
||||
.map_err(|e| {
|
||||
WasmError::ConfigError(format!("Failed to add sign-payload function: {}", e))
|
||||
})?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -98,6 +98,32 @@ interface host {
|
||||
///
|
||||
/// Returns true if the secret exists and is accessible to this tool.
|
||||
secret-exists: func(name: string) -> bool;
|
||||
|
||||
// ==================== Signing Capability ====================
|
||||
|
||||
/// Result of a payload signing request.
|
||||
record sign-result {
|
||||
/// Base64-encoded signature bytes (set on success).
|
||||
signature: option<string>,
|
||||
/// Error message (set on failure).
|
||||
error: option<string>,
|
||||
/// True if user approval is needed before signing can proceed.
|
||||
approval-pending: bool,
|
||||
}
|
||||
|
||||
/// Sign a payload using a NEAR key managed by the host (if capability granted).
|
||||
///
|
||||
/// Security:
|
||||
/// - Private keys NEVER enter WASM memory; signing happens in host code only
|
||||
/// - Only key labels declared in the tool's signing capability can be used
|
||||
/// - Rate-limited per execution
|
||||
/// - Subject to the host's transaction policy (may require user approval)
|
||||
///
|
||||
/// The payload should be base64-encoded bytes to sign.
|
||||
/// The context-json is optional metadata about what's being signed (for policy display).
|
||||
///
|
||||
/// Returns sign-result with either signature or error/approval-pending.
|
||||
sign-payload: func(key-label: string, payload: string, context-json: string) -> sign-result;
|
||||
}
|
||||
|
||||
/// Tool interface that sandboxed tools must implement.
|
||||
|
||||
Reference in New Issue
Block a user