Compare commits

..
Author SHA1 Message Date
Illia PolosukhinandClaude Opus 4.6 810ba58fd2 feat: Add Okta SSO WASM tool for profile management and app catalog
Sandboxed WASM tool that integrates with Okta's Management API and
MyAccount API. Supports user profile CRUD, listing all SSO app
chiclets, searching apps by name, retrieving SSO launch links, and
fetching org info. Uses OAuth2 with PKCE against the Org Authorization
Server, with the domain stored in workspace at okta/domain.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
2026-02-09 23:31:29 -08:00
24 changed files with 663 additions and 3281 deletions
Generated
-24
View File
@@ -2236,30 +2236,6 @@ dependencies = [
"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",
]
[[package]]
name = "is-docker"
version = "0.2.0"
-3
View File
@@ -1,6 +1,3 @@
[workspace]
members = [".", "benchmarks"]
[package]
name = "ironclaw"
version = "0.1.0"
-49
View File
@@ -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"
-237
View File
@@ -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);
}
}
-183
View File
@@ -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);
}
}
-110
View File
@@ -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,
})
}
}
}
-291
View File
@@ -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);
}
}
-227
View File
@@ -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);
}
}
-229
View File
@@ -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);
}
}
-203
View File
@@ -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"
);
}
}
-40
View File
@@ -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 },
}
-248
View File
@@ -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());
}
}
-314
View File
@@ -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(())
}
-392
View File
@@ -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"));
}
}
-460
View File
@@ -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()),
}
}
-113
View File
@@ -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")
);
}
}
-156
View File
@@ -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)
}
}
-1
View File
@@ -314,7 +314,6 @@ 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",
+1 -1
View File
@@ -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};
+21
View File
@@ -0,0 +1,21 @@
[package]
name = "okta-tool"
version = "0.1.0"
edition = "2021"
description = "Okta SSO tool for IronClaw (WASM component) — user profile, app catalog, and SSO launch links"
license = "MIT OR Apache-2.0"
publish = false
[lib]
crate-type = ["cdylib"]
[dependencies]
wit-bindgen = "=0.36"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
[profile.release]
opt-level = "s"
lto = true
strip = true
codegen-units = 1
@@ -0,0 +1,93 @@
{
"http": {
"allowlist": [
{
"host": "*.okta.com",
"path_prefix": "/api/v1/",
"methods": ["GET", "POST", "PUT"]
},
{
"host": "*.okta.com",
"path_prefix": "/idp/myaccount/",
"methods": ["GET", "PUT"]
},
{
"host": "*.okta.com",
"path_prefix": "/oauth2/v1/",
"methods": ["POST"]
},
{
"host": "*.oktapreview.com",
"path_prefix": "/api/v1/",
"methods": ["GET", "POST", "PUT"]
},
{
"host": "*.oktapreview.com",
"path_prefix": "/idp/myaccount/",
"methods": ["GET", "PUT"]
},
{
"host": "*.oktapreview.com",
"path_prefix": "/oauth2/v1/",
"methods": ["POST"]
},
{
"host": "*.okta-emea.com",
"path_prefix": "/api/v1/",
"methods": ["GET", "POST", "PUT"]
},
{
"host": "*.okta-emea.com",
"path_prefix": "/idp/myaccount/",
"methods": ["GET", "PUT"]
},
{
"host": "*.okta-emea.com",
"path_prefix": "/oauth2/v1/",
"methods": ["POST"]
}
],
"credentials": {
"okta_oauth_token": {
"secret_name": "okta_oauth_token",
"location": { "type": "bearer" },
"host_patterns": ["*.okta.com", "*.oktapreview.com", "*.okta-emea.com"]
}
},
"rate_limit": {
"requests_per_minute": 30,
"requests_per_hour": 500
},
"timeout_secs": 30
},
"workspace": {
"allowed_prefixes": ["okta/"]
},
"secrets": {
"allowed_names": ["okta_oauth_token"]
},
"auth": {
"secret_name": "okta_oauth_token",
"display_name": "Okta",
"oauth": {
"authorization_url": "https://{okta_domain}/oauth2/v1/authorize",
"token_url": "https://{okta_domain}/oauth2/v1/token",
"client_id_env": "OKTA_OAUTH_CLIENT_ID",
"client_secret_env": "OKTA_OAUTH_CLIENT_SECRET",
"scopes": [
"openid",
"profile",
"email",
"offline_access",
"okta.users.read.self",
"okta.users.manage.self",
"okta.apps.read"
],
"use_pkce": true
},
"instructions": "1. In your Okta Admin Console, go to Applications > Create App Integration\n2. Select 'OIDC - OpenID Connect', then 'Web Application'\n3. Set Sign-in redirect URI to http://localhost:9876/callback (through :9886)\n4. Under Okta API Scopes, grant: okta.users.read.self, okta.users.manage.self, okta.apps.read\n5. Copy the Client ID and Client Secret\n6. IMPORTANT: You must use the Org Authorization Server (not a custom one)\n7. Store your Okta domain in workspace at 'okta/domain' (e.g., 'mycompany.okta.com')\n8. For custom domains, add them to okta-tool.capabilities.json allowlist",
"setup_url": "https://developer.okta.com/docs/guides/implement-oauth-for-okta/main/",
"token_hint": "OAuth2 access token (JWT)",
"env_var": "OKTA_OAUTH_TOKEN"
}
}
+281
View File
@@ -0,0 +1,281 @@
use crate::near::agent::host;
use crate::types::*;
const WORKSPACE_DOMAIN_PATH: &str = "okta/domain";
/// Read the configured Okta domain from workspace, or return a helpful error.
fn get_domain() -> Result<String, String> {
host::workspace_read(WORKSPACE_DOMAIN_PATH).ok_or_else(|| {
"Okta domain not configured. Write your Okta domain to workspace path 'okta/domain' \
using the memory_write tool (e.g., memory_write with path='okta/domain' and \
content='mycompany.okta.com')."
.to_string()
})
}
/// Build the base URL for the Okta Management API.
fn management_base(domain: &str) -> String {
format!("https://{}/api/v1", domain)
}
/// Make an Okta API call.
fn okta_api_call(method: &str, url: &str, body: Option<&str>) -> Result<String, String> {
let headers = if body.is_some() {
r#"{"Content-Type": "application/json", "Accept": "application/json"}"#
} else {
r#"{"Accept": "application/json"}"#
};
let body_bytes = body.map(|b| b.as_bytes().to_vec());
host::log(
host::LogLevel::Debug,
&format!("Okta API: {} {}", method, url),
);
let response = host::http_request(method, url, headers, body_bytes.as_deref())?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
// Try to extract Okta's error summary for a better message.
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&body_text) {
if let Some(summary) = parsed["errorSummary"].as_str() {
return Err(format!("Okta API error ({}): {}", response.status, summary));
}
}
return Err(format!(
"Okta API returned status {}: {}",
response.status, body_text
));
}
String::from_utf8(response.body).map_err(|e| format!("Invalid UTF-8: {}", e))
}
// ---------------------------------------------------------------------------
// Action implementations
// ---------------------------------------------------------------------------
/// GET /api/v1/users/me
pub fn get_profile() -> Result<String, String> {
let domain = get_domain()?;
let url = format!("{}/users/me", management_base(&domain));
let response = okta_api_call("GET", &url, None)?;
let parsed: serde_json::Value =
serde_json::from_str(&response).map_err(|e| format!("Failed to parse response: {}", e))?;
let profile = parse_user_profile(&parsed)?;
serde_json::to_string(&profile).map_err(|e| e.to_string())
}
/// POST /api/v1/users/me (partial update via Management API)
pub fn update_profile(fields: &serde_json::Value) -> Result<String, String> {
let domain = get_domain()?;
let url = format!("{}/users/me", management_base(&domain));
// Wrap fields under "profile" key for Okta's expected format.
let payload = serde_json::json!({ "profile": fields });
let body = serde_json::to_string(&payload).map_err(|e| e.to_string())?;
let response = okta_api_call("POST", &url, Some(&body))?;
let parsed: serde_json::Value =
serde_json::from_str(&response).map_err(|e| format!("Failed to parse response: {}", e))?;
let profile = parse_user_profile(&parsed)?;
let result = UpdateProfileResult {
success: true,
profile,
};
serde_json::to_string(&result).map_err(|e| e.to_string())
}
/// GET /api/v1/users/me/appLinks
pub fn list_apps() -> Result<String, String> {
let domain = get_domain()?;
let url = format!("{}/users/me/appLinks", management_base(&domain));
let response = okta_api_call("GET", &url, None)?;
let parsed: serde_json::Value =
serde_json::from_str(&response).map_err(|e| format!("Failed to parse response: {}", e))?;
let apps = parse_app_links(&parsed)?;
let count = apps.len();
let result = ListAppsResult { apps, count };
serde_json::to_string(&result).map_err(|e| e.to_string())
}
/// Search apps by label (case-insensitive substring match).
pub fn search_apps(query: &str) -> Result<String, String> {
let domain = get_domain()?;
let url = format!("{}/users/me/appLinks", management_base(&domain));
let response = okta_api_call("GET", &url, None)?;
let parsed: serde_json::Value =
serde_json::from_str(&response).map_err(|e| format!("Failed to parse response: {}", e))?;
let all_apps = parse_app_links(&parsed)?;
let query_lower = query.to_lowercase();
let apps: Vec<AppLink> = all_apps
.into_iter()
.filter(|app| {
app.label.to_lowercase().contains(&query_lower)
|| app.app_name.to_lowercase().contains(&query_lower)
})
.collect();
let count = apps.len();
let result = ListAppsResult { apps, count };
serde_json::to_string(&result).map_err(|e| e.to_string())
}
/// Find an app by ID or label and return its SSO launch link.
pub fn get_app_sso_link(app: &str) -> Result<String, String> {
let domain = get_domain()?;
let url = format!("{}/users/me/appLinks", management_base(&domain));
let response = okta_api_call("GET", &url, None)?;
let parsed: serde_json::Value =
serde_json::from_str(&response).map_err(|e| format!("Failed to parse response: {}", e))?;
let all_apps = parse_app_links(&parsed)?;
let app_lower = app.to_lowercase();
// Try exact ID match first, then case-insensitive label match.
let found = all_apps
.iter()
.find(|a| a.app_instance_id == app)
.or_else(|| {
all_apps
.iter()
.find(|a| a.label.to_lowercase() == app_lower)
})
.or_else(|| {
all_apps
.iter()
.find(|a| a.label.to_lowercase().contains(&app_lower))
});
match found {
Some(app_link) => {
let result = AppSsoLinkResult {
label: app_link.label.clone(),
link_url: app_link.link_url.clone(),
app_instance_id: app_link.app_instance_id.clone(),
app_name: app_link.app_name.clone(),
};
serde_json::to_string(&result).map_err(|e| e.to_string())
}
None => {
let available: Vec<String> = all_apps.iter().map(|a| a.label.clone()).collect();
Err(format!(
"App '{}' not found. Available apps: {}",
app,
available.join(", ")
))
}
}
}
/// GET /idp/myaccount/organization
pub fn get_org_info() -> Result<String, String> {
let domain = get_domain()?;
let url = format!("https://{}/idp/myaccount/organization", domain);
// MyAccount API requires the okta-version header.
let response = okta_api_call_with_headers(
"GET",
&url,
None,
r#"{"Accept": "application/json; okta-version=1.0.0"}"#,
)?;
let parsed: serde_json::Value =
serde_json::from_str(&response).map_err(|e| format!("Failed to parse response: {}", e))?;
let result = OrgInfo {
id: parsed["id"].as_str().unwrap_or("").to_string(),
name: parsed["name"].as_str().unwrap_or("").to_string(),
subdomain: parsed["subdomain"].as_str().map(|s| s.to_string()),
website: parsed["website"].as_str().map(|s| s.to_string()),
support_phone: parsed["supportPhoneNumber"].as_str().map(|s| s.to_string()),
technical_contact: parsed["technicalContact"].as_str().map(|s| s.to_string()),
};
serde_json::to_string(&result).map_err(|e| e.to_string())
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
/// Like `okta_api_call` but with custom headers (for MyAccount API versioning).
fn okta_api_call_with_headers(
method: &str,
url: &str,
body: Option<&str>,
headers: &str,
) -> Result<String, String> {
let body_bytes = body.map(|b| b.as_bytes().to_vec());
host::log(
host::LogLevel::Debug,
&format!("Okta API: {} {}", method, url),
);
let response = host::http_request(method, url, headers, body_bytes.as_deref())?;
if response.status < 200 || response.status >= 300 {
let body_text = String::from_utf8_lossy(&response.body);
if let Ok(parsed) = serde_json::from_str::<serde_json::Value>(&body_text) {
if let Some(summary) = parsed["errorSummary"].as_str() {
return Err(format!("Okta API error ({}): {}", response.status, summary));
}
}
return Err(format!(
"Okta API returned status {}: {}",
response.status, body_text
));
}
String::from_utf8(response.body).map_err(|e| format!("Invalid UTF-8: {}", e))
}
fn parse_user_profile(v: &serde_json::Value) -> Result<UserProfile, String> {
let p = &v["profile"];
Ok(UserProfile {
id: v["id"].as_str().unwrap_or("").to_string(),
status: v["status"].as_str().unwrap_or("").to_string(),
first_name: p["firstName"].as_str().unwrap_or("").to_string(),
last_name: p["lastName"].as_str().unwrap_or("").to_string(),
email: p["email"].as_str().unwrap_or("").to_string(),
login: p["login"].as_str().unwrap_or("").to_string(),
mobile_phone: p["mobilePhone"].as_str().map(|s| s.to_string()),
display_name: p["displayName"].as_str().map(|s| s.to_string()),
nick_name: p["nickName"].as_str().map(|s| s.to_string()),
title: p["title"].as_str().map(|s| s.to_string()),
department: p["department"].as_str().map(|s| s.to_string()),
organization: p["organization"].as_str().map(|s| s.to_string()),
timezone: p["timezone"].as_str().map(|s| s.to_string()),
locale: p["locale"].as_str().map(|s| s.to_string()),
})
}
fn parse_app_links(v: &serde_json::Value) -> Result<Vec<AppLink>, String> {
let arr = v
.as_array()
.ok_or_else(|| "Expected array of app links from Okta".to_string())?;
Ok(arr
.iter()
.map(|a| AppLink {
app_instance_id: a["appInstanceId"].as_str().unwrap_or("").to_string(),
label: a["label"].as_str().unwrap_or("").to_string(),
link_url: a["linkUrl"].as_str().unwrap_or("").to_string(),
logo_url: a["logoUrl"].as_str().map(|s| s.to_string()),
app_name: a["appName"].as_str().unwrap_or("").to_string(),
hidden: a["hidden"].as_bool().unwrap_or(false),
})
.collect())
}
+148
View File
@@ -0,0 +1,148 @@
//! Okta WASM Tool for IronClaw.
//!
//! Provides user profile management, SSO app catalog browsing, and
//! launch links for all applications under Okta single sign-on.
//!
//! # Setup
//!
//! 1. Configure OAuth2 with PKCE (see capabilities.json instructions)
//! 2. Write your Okta domain to workspace: `memory_write(path="okta/domain", content="mycompany.okta.com")`
//! 3. All actions read the domain from workspace automatically
//!
//! # Capabilities Required
//!
//! - HTTP: `*.okta.com/api/v1/*`, `*.okta.com/idp/myaccount/*` (GET, POST, PUT)
//! - Secrets: `okta_oauth_token` (injected as Bearer token)
//! - Workspace: `okta/` prefix (read-only, for domain config)
//!
//! # Supported Actions
//!
//! - `get_profile`: Fetch the current user's profile
//! - `update_profile`: Update profile fields
//! - `list_apps`: List all SSO apps assigned to the user
//! - `search_apps`: Search apps by name
//! - `get_app_sso_link`: Get the SSO launch URL for a specific app
//! - `get_org_info`: Get organization details
mod api;
mod types;
use types::OktaAction;
wit_bindgen::generate!({
world: "sandboxed-tool",
path: "../../wit/tool.wit",
});
struct OktaTool;
impl exports::near::agent::tool::Guest for OktaTool {
fn execute(req: exports::near::agent::tool::Request) -> exports::near::agent::tool::Response {
match execute_inner(&req.params) {
Ok(result) => exports::near::agent::tool::Response {
output: Some(result),
error: None,
},
Err(e) => exports::near::agent::tool::Response {
output: None,
error: Some(e),
},
}
}
fn schema() -> String {
r#"{
"type": "object",
"required": ["action"],
"oneOf": [
{
"properties": {
"action": { "const": "get_profile" }
},
"required": ["action"]
},
{
"properties": {
"action": { "const": "update_profile" },
"fields": {
"type": "object",
"description": "Profile fields to update. Common: firstName, lastName, email, mobilePhone, displayName, nickName, title, department, organization"
}
},
"required": ["action", "fields"]
},
{
"properties": {
"action": { "const": "list_apps" }
},
"required": ["action"]
},
{
"properties": {
"action": { "const": "search_apps" },
"query": {
"type": "string",
"description": "Case-insensitive search query to match against app labels and names"
}
},
"required": ["action", "query"]
},
{
"properties": {
"action": { "const": "get_app_sso_link" },
"app": {
"type": "string",
"description": "App instance ID (e.g., '0oa1xxx') or app label (e.g., 'Google Workspace')"
}
},
"required": ["action", "app"]
},
{
"properties": {
"action": { "const": "get_org_info" }
},
"required": ["action"]
}
]
}"#
.to_string()
}
fn description() -> String {
"Okta SSO tool for managing your profile and accessing all applications under \
single sign-on. Supports viewing/updating your Okta profile, listing all assigned \
SSO apps, searching apps by name, and getting direct SSO launch links. Requires \
Okta domain in workspace at 'okta/domain' and an OAuth token with \
okta.users.read.self, okta.users.manage.self, and okta.apps.read scopes."
.to_string()
}
}
fn execute_inner(params: &str) -> Result<String, String> {
if !crate::near::agent::host::secret_exists("okta_oauth_token") {
return Err(
"Okta OAuth token not configured. Please add the 'okta_oauth_token' secret \
via OAuth2 flow or set the OKTA_OAUTH_TOKEN environment variable."
.to_string(),
);
}
let action: OktaAction =
serde_json::from_str(params).map_err(|e| format!("Invalid parameters: {}", e))?;
crate::near::agent::host::log(
crate::near::agent::host::LogLevel::Info,
&format!("Executing Okta action: {:?}", action),
);
match action {
OktaAction::GetProfile => api::get_profile(),
OktaAction::UpdateProfile { fields } => api::update_profile(&fields),
OktaAction::ListApps => api::list_apps(),
OktaAction::SearchApps { query } => api::search_apps(&query),
OktaAction::GetAppSsoLink { app } => api::get_app_sso_link(&app),
OktaAction::GetOrgInfo => api::get_org_info(),
}
}
export!(OktaTool);
+119
View File
@@ -0,0 +1,119 @@
use serde::{Deserialize, Serialize};
/// Input parameters for the Okta tool.
///
/// Actions map to Okta Management API (/api/v1/) and MyAccount API (/idp/myaccount/).
/// The tool reads the Okta domain from workspace at `okta/domain`.
#[derive(Debug, Deserialize)]
#[serde(tag = "action", rename_all = "snake_case")]
pub enum OktaAction {
/// Get the current user's Okta profile.
GetProfile,
/// Update fields on the current user's profile (partial update).
UpdateProfile {
/// Key-value pairs of profile fields to update.
/// Common fields: firstName, lastName, email, mobilePhone, displayName,
/// nickName, title, department, organization.
fields: serde_json::Value,
},
/// List all SSO applications assigned to the current user.
ListApps,
/// Search assigned apps by name (case-insensitive substring match).
SearchApps {
/// Search query to match against app labels.
query: String,
},
/// Get the SSO launch link for a specific app by its instance ID or label.
GetAppSsoLink {
/// App instance ID (e.g., "0oa1xxx") or app label to search for.
app: String,
},
/// Get information about the Okta organization.
GetOrgInfo,
}
// ---------------------------------------------------------------------------
// Response types
// ---------------------------------------------------------------------------
/// User profile from Okta.
#[derive(Debug, Serialize)]
pub struct UserProfile {
pub id: String,
pub status: String,
pub first_name: String,
pub last_name: String,
pub email: String,
pub login: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub mobile_phone: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub display_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub nick_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub department: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub organization: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub timezone: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub locale: Option<String>,
}
/// Result of a profile update.
#[derive(Debug, Serialize)]
pub struct UpdateProfileResult {
pub success: bool,
pub profile: UserProfile,
}
/// An SSO app link (chiclet) assigned to the user.
#[derive(Debug, Serialize)]
pub struct AppLink {
pub app_instance_id: String,
pub label: String,
pub link_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub logo_url: Option<String>,
pub app_name: String,
pub hidden: bool,
}
/// Result of listing or searching apps.
#[derive(Debug, Serialize)]
pub struct ListAppsResult {
pub apps: Vec<AppLink>,
pub count: usize,
}
/// SSO launch link for a specific app.
#[derive(Debug, Serialize)]
pub struct AppSsoLinkResult {
pub label: String,
pub link_url: String,
pub app_instance_id: String,
pub app_name: String,
}
/// Okta organization info.
#[derive(Debug, Serialize)]
pub struct OrgInfo {
pub id: String,
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub subdomain: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub website: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub support_phone: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub technical_contact: Option<String>,
}