mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-27 08:00:17 +00:00
Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 github-actions[bot]GitHubgithub-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> |
542268fde5 | ||
|
|
3b6105d5ea | ||
|
Pierre LE GUENandGitHub
|
df8616b604 | ||
|
|
e8dcb52fda | ||
|
github-actions[bot]GitHubgithub-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>Illia Polosukhin |
1f18422b88 | ||
|
|
448383cfb0 | ||
|
|
7df356c109 |
+8
-10
@@ -6,21 +6,19 @@ DATABASE_POOL_SIZE=10
|
||||
# LLM_BACKEND=nearai # default
|
||||
# Possible values: nearai, ollama, openai_compatible, openai, anthropic, tinfoil
|
||||
|
||||
# === NEAR AI Chat (Responses API, session token auth) ===
|
||||
# Default mode. Uses browser OAuth (GitHub/Google) on first run.
|
||||
# Session token stored in ~/.ironclaw/session.json automatically.
|
||||
# For hosting providers: set NEARAI_SESSION_TOKEN env var directly.
|
||||
# === NEAR AI (Chat Completions API) ===
|
||||
# Two auth modes:
|
||||
# 1. Session token (default): Uses browser OAuth (GitHub/Google) on first run.
|
||||
# Session token stored in ~/.ironclaw/session.json automatically.
|
||||
# Base URL defaults to https://private.near.ai
|
||||
# 2. API key: Set NEARAI_API_KEY to use API key auth from cloud.near.ai.
|
||||
# Base URL defaults to https://cloud-api.near.ai
|
||||
NEARAI_MODEL=zai-org/GLM-5-FP8
|
||||
NEARAI_BASE_URL=https://private.near.ai
|
||||
NEARAI_AUTH_URL=https://private.near.ai
|
||||
# NEARAI_SESSION_TOKEN=sess_... # hosting providers: set this
|
||||
# NEARAI_SESSION_PATH=~/.ironclaw/session.json # optional, default shown
|
||||
|
||||
# === NEAR AI Cloud (Chat Completions API, API key auth) ===
|
||||
# Auto-selected when NEARAI_API_KEY is set. Get a key from cloud.near.ai.
|
||||
# NEARAI_API_KEY=...
|
||||
# NEARAI_BASE_URL=https://cloud-api.near.ai # default for cloud mode
|
||||
# NEARAI_API_MODE=chat_completions # auto-detected from API key
|
||||
# NEARAI_API_KEY=... # API key from cloud.near.ai
|
||||
|
||||
# Local LLM Providers (Ollama, LM Studio, vLLM, LiteLLM)
|
||||
|
||||
|
||||
@@ -7,6 +7,40 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
## [0.9.0](https://github.com/nearai/ironclaw/compare/v0.8.0...v0.9.0) - 2026-02-21
|
||||
|
||||
### Added
|
||||
|
||||
- add TEE attestation shield to web gateway UI ([#275](https://github.com/nearai/ironclaw/pull/275))
|
||||
- configurable tool iterations, auto-approve, and policy fix ([#251](https://github.com/nearai/ironclaw/pull/251))
|
||||
|
||||
### Fixed
|
||||
|
||||
- add X-Accel-Buffering header to SSE endpoints ([#277](https://github.com/nearai/ironclaw/pull/277))
|
||||
|
||||
## [0.8.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.7.0...ironclaw-v0.8.0) - 2026-02-20
|
||||
|
||||
### Added
|
||||
|
||||
- extension registry with metadata catalog and onboarding integration ([#238](https://github.com/nearai/ironclaw/pull/238))
|
||||
- *(models)* add GPT-5.3 Codex, full GPT-5.x family, Claude 4.x series, o4-mini ([#197](https://github.com/nearai/ironclaw/pull/197))
|
||||
- wire memory hygiene into the heartbeat loop ([#195](https://github.com/nearai/ironclaw/pull/195))
|
||||
|
||||
### Fixed
|
||||
|
||||
- persist WASM channel workspace writes across callbacks ([#264](https://github.com/nearai/ironclaw/pull/264))
|
||||
- consolidate per-module ENV_MUTEX into crate-wide test lock ([#246](https://github.com/nearai/ironclaw/pull/246))
|
||||
- remove auto-proceed fake user message injection from agent loop ([#255](https://github.com/nearai/ironclaw/pull/255))
|
||||
- onboarding errors reset flow and remote server auth (#185, #186) ([#248](https://github.com/nearai/ironclaw/pull/248))
|
||||
- parallelize tool call execution via JoinSet ([#219](https://github.com/nearai/ironclaw/pull/219)) ([#252](https://github.com/nearai/ironclaw/pull/252))
|
||||
- prevent pipe deadlock in shell command execution ([#140](https://github.com/nearai/ironclaw/pull/140))
|
||||
- persist turns after approval and add agent-level tests ([#250](https://github.com/nearai/ironclaw/pull/250))
|
||||
|
||||
### Other
|
||||
|
||||
- add automated PR labeling system ([#253](https://github.com/nearai/ironclaw/pull/253))
|
||||
- update CLAUDE.md for recently merged features ([#183](https://github.com/nearai/ironclaw/pull/183))
|
||||
|
||||
## [0.7.0](https://github.com/nearai/ironclaw/compare/ironclaw-v0.6.0...ironclaw-v0.7.0) - 2026-02-19
|
||||
|
||||
### Added
|
||||
|
||||
@@ -120,8 +120,7 @@ src/
|
||||
├── llm/ # LLM integration (multi-provider)
|
||||
│ ├── mod.rs # Provider factory, LlmBackend enum
|
||||
│ ├── provider.rs # LlmProvider trait, message types
|
||||
│ ├── nearai.rs # NEAR AI Responses API provider
|
||||
│ ├── nearai_chat.rs # NEAR AI Chat Completions fallback
|
||||
│ ├── nearai_chat.rs # NEAR AI Chat Completions provider (session token + API key auth)
|
||||
│ ├── reasoning.rs # Planning, tool selection, evaluation
|
||||
│ ├── session.rs # Session token management with auto-renewal
|
||||
│ ├── circuit_breaker.rs # Circuit breaker for provider failures
|
||||
@@ -339,13 +338,12 @@ LIBSQL_PATH=~/.ironclaw/ironclaw.db # libSQL local path (default)
|
||||
# LIBSQL_AUTH_TOKEN=xxx # Required with LIBSQL_URL
|
||||
|
||||
# NEAR AI (when LLM_BACKEND=nearai, the default)
|
||||
# Two modes: "NEAR AI Chat" (session token) or "NEAR AI Cloud" (API key)
|
||||
# NEAR AI Chat (Responses API, default):
|
||||
NEARAI_SESSION_TOKEN=sess_... # session token for chat-api
|
||||
# Two auth modes: session token (default) or API key
|
||||
# Session token auth (default): uses browser OAuth on first run
|
||||
NEARAI_SESSION_TOKEN=sess_... # hosting providers: set this
|
||||
NEARAI_BASE_URL=https://private.near.ai
|
||||
# NEAR AI Cloud (Chat Completions API, auto-selected when API key is set):
|
||||
# API key auth: set NEARAI_API_KEY, base URL defaults to cloud-api.near.ai
|
||||
# NEARAI_API_KEY=... # API key from cloud.near.ai
|
||||
# NEARAI_BASE_URL=https://cloud-api.near.ai
|
||||
NEARAI_MODEL=claude-3-5-sonnet-20241022
|
||||
|
||||
# Agent settings
|
||||
@@ -408,11 +406,9 @@ TINFOIL_MODEL=kimi-k2-5 # Default model
|
||||
|
||||
IronClaw supports multiple LLM backends via the `LLM_BACKEND` env var: `nearai` (default), `openai`, `anthropic`, `ollama`, `openai_compatible`, and `tinfoil`.
|
||||
|
||||
**NEAR AI Chat** -- Uses the NEAR AI Responses API (`https://private.near.ai/v1/responses`). Authenticates with session tokens (`sess_xxx`) obtained via browser OAuth (GitHub/Google). Supports response chaining (delta-only follow-up messages) for efficient multi-turn conversations. This is the default mode when no `NEARAI_API_KEY` is set. Set `NEARAI_SESSION_TOKEN` env var for hosting providers that inject tokens via environment. Configure with `NEARAI_BASE_URL` (default: `https://private.near.ai`).
|
||||
**NEAR AI** -- Uses the Chat Completions API with dual auth support. Session token auth (default): authenticates with session tokens (`sess_xxx`) obtained via browser OAuth (GitHub/Google), base URL defaults to `https://private.near.ai`. API key auth: set `NEARAI_API_KEY` (from `cloud.near.ai`), base URL defaults to `https://cloud-api.near.ai`. Both modes use the same Chat Completions endpoint. Tool messages are flattened to plain text for compatibility. Set `NEARAI_SESSION_TOKEN` env var for hosting providers that inject tokens via environment.
|
||||
|
||||
**NEAR AI Cloud** -- Uses the OpenAI-compatible Chat Completions API (`https://cloud-api.near.ai/v1/chat/completions`). Authenticates with API keys from `cloud.near.ai`. Auto-selected when `NEARAI_API_KEY` is set (or explicitly via `NEARAI_API_MODE=chat_completions`). Tool messages are flattened to plain text for compatibility. Configure with `NEARAI_API_KEY` and `NEARAI_BASE_URL` (default: `https://cloud-api.near.ai`).
|
||||
|
||||
**Tinfoil** -- Private inference via `https://inference.tinfoil.sh/v1`. Runs models inside hardware-attested TEEs so neither Tinfoil nor the cloud provider can see prompts or responses. Uses the OpenAI-compatible Chat Completions API only (not the Responses API, so tool calls are adapted to chat format). Configure with `TINFOIL_API_KEY` and `TINFOIL_MODEL` (default: `kimi-k2-5`).
|
||||
**Tinfoil** -- Private inference via `https://inference.tinfoil.sh/v1`. Runs models inside hardware-attested TEEs so neither Tinfoil nor the cloud provider can see prompts or responses. Uses the OpenAI-compatible Chat Completions API. Configure with `TINFOIL_API_KEY` and `TINFOIL_MODEL` (default: `kimi-k2-5`).
|
||||
|
||||
## Database
|
||||
|
||||
|
||||
Generated
+1
-25
@@ -2490,7 +2490,7 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "ironclaw"
|
||||
version = "0.7.0"
|
||||
version = "0.9.0"
|
||||
dependencies = [
|
||||
"aes-gcm",
|
||||
"aho-corasick",
|
||||
@@ -2559,30 +2559,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"
|
||||
|
||||
+2
-2
@@ -1,5 +1,5 @@
|
||||
[workspace]
|
||||
members = [".", "benchmarks"]
|
||||
members = ["."]
|
||||
exclude = [
|
||||
"channels-src/discord",
|
||||
"channels-src/telegram",
|
||||
@@ -19,7 +19,7 @@ exclude = [
|
||||
|
||||
[package]
|
||||
name = "ironclaw"
|
||||
version = "0.7.0"
|
||||
version = "0.9.0"
|
||||
edition = "2024"
|
||||
rust-version = "1.92"
|
||||
description = "Secure personal AI assistant that protects your data and expands its capabilities on the fly"
|
||||
|
||||
@@ -1,50 +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"
|
||||
publish = false
|
||||
|
||||
[[bin]]
|
||||
name = "ironclaw-bench"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
ironclaw = { path = ".." }
|
||||
|
||||
# Async runtime
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
tokio-stream = { version = "0.1", features = ["sync"] }
|
||||
futures = "0.3"
|
||||
|
||||
# Serialization
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
toml = "0.8"
|
||||
|
||||
# CLI
|
||||
clap = { version = "4", features = ["derive"] }
|
||||
|
||||
# Core types
|
||||
uuid = { version = "1", features = ["v4", "serde"] }
|
||||
chrono = { version = "0.4", features = ["serde"] }
|
||||
rust_decimal = { version = "1", features = ["serde", "serde-with-str"] }
|
||||
|
||||
# Error handling
|
||||
thiserror = "2"
|
||||
anyhow = "1"
|
||||
|
||||
# Async traits
|
||||
async-trait = "0.1"
|
||||
|
||||
# Logging
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
|
||||
# Scoring
|
||||
regex = "1"
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = "3"
|
||||
@@ -1,15 +0,0 @@
|
||||
{
|
||||
"run_id": "8a48de1f-09ce-4c35-9ad1-dad98fb83a1a",
|
||||
"suite_id": "spot",
|
||||
"config_label": "default",
|
||||
"model": "openai/gpt-5.2",
|
||||
"commit_hash": "2c43b83",
|
||||
"pass_rate": 1.0,
|
||||
"avg_score": 1.0,
|
||||
"total_tasks": 21,
|
||||
"completed_tasks": 21,
|
||||
"total_cost_usd": 0.307053,
|
||||
"total_wall_time_ms": 111009,
|
||||
"started_at": "2026-02-17T22:02:08.206112Z",
|
||||
"finished_at": "2026-02-17T22:03:59.270325Z"
|
||||
}
|
||||
@@ -1,21 +0,0 @@
|
||||
{"task_id":"smoke-greeting","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":2037,"llm_calls":1,"input_tokens":1359,"output_tokens":65,"estimated_cost_usd":0.005052,"tool_calls":[],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"I’m NEAR AI Agent—an autonomous assistant that helps you plan, write, debug, and execute tasks using available tools (like reading/writing files, running shell commands, and making HTTP requests). Tell me what you’re working on and what outcome you want, and I’ll drive it to completion.","started_at":"2026-02-17T22:02:08.206165Z","finished_at":"2026-02-17T22:02:10.243552Z","config_label":"default","error":null}
|
||||
{"task_id":"smoke-math","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":2038,"llm_calls":1,"input_tokens":1366,"output_tokens":5,"estimated_cost_usd":0.004173,"tool_calls":[],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"1081","started_at":"2026-02-17T22:02:10.243933Z","finished_at":"2026-02-17T22:02:12.282133Z","config_label":"default","error":null}
|
||||
{"task_id":"tool-echo","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":2646,"llm_calls":2,"input_tokens":2786,"output_tokens":24,"estimated_cost_usd":0.008718,"tool_calls":[{"name":"echo","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"Spot check passed","started_at":"2026-02-17T22:02:12.282294Z","finished_at":"2026-02-17T22:02:14.928540Z","config_label":"default","error":null}
|
||||
{"task_id":"tool-time","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":3460,"llm_calls":2,"input_tokens":2839,"output_tokens":78,"estimated_cost_usd":0.009687,"tool_calls":[{"name":"time","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"Current date/time (UTC): **2026-02-17 22:02:16** \nISO 8601: **2026-02-17T22:02:16.045933+00:00**","started_at":"2026-02-17T22:02:14.928810Z","finished_at":"2026-02-17T22:02:18.389188Z","config_label":"default","error":null}
|
||||
{"task_id":"tool-json-query","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":2751,"llm_calls":2,"input_tokens":2839,"output_tokens":43,"estimated_cost_usd":0.009162,"tool_calls":[{"name":"json","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"Bob","started_at":"2026-02-17T22:02:18.389812Z","finished_at":"2026-02-17T22:02:21.141146Z","config_label":"default","error":null}
|
||||
{"task_id":"tool-shell-echo","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":3460,"llm_calls":2,"input_tokens":2811,"output_tokens":25,"estimated_cost_usd":0.008808,"tool_calls":[{"name":"shell","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"benchmark test","started_at":"2026-02-17T22:02:21.141384Z","finished_at":"2026-02-17T22:02:24.601812Z","config_label":"default","error":null}
|
||||
{"task_id":"tool-list-dir","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":19146,"llm_calls":2,"input_tokens":4370,"output_tokens":1337,"estimated_cost_usd":0.033165,"tool_calls":[{"name":"list_dir","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"```\nclaude-501/\nclaude-505/\nclaude/\ncodex/\ncom.apple.launchd.HqYyNM7hoK/\ncom.apple.launchd.Z0bTtlsPEN/\ncom.apple.launchd.oY8y5gkzN4/\ncom.apple.launchd.sdo3WGu0S7/\ndata-gym-cache/\ndocker-desktop-privileged2203342862/\ngrammers-test/\nnear-agent-target/\nnear-market-bot-push/\nnear-market-bot/\nopenclaw/\npowerlog/\ntest_grpo/\ntest_grpo2/\ntest_grpo_entry/\ntest_grpo_mix/\ntest_grpo_singleton/\ntmpybtjiuc0/\n.d8bd37babaf2f1f7-00000000.node (5.9MB)\n.d8bd37fbfbb0f7ff-00000000.node (5.9MB)\n.d8bd3d9bdf6dd3f7-00000000.node (5.9MB)\n.d8bd3fbfbf21ddf7-00000000.node (5.9MB)\n.d8bd7dfebb38d5ff-00000000.node (5.9MB)\n.d8bdf5deba3bf1f7-00000000.node (5.9MB)\n.d8bdf7cfff76f3f7-00000000.node (5.9MB)\n.d8bdfd8ffe26d5ff-00000000.node (5.9MB)\n.d8bdfdeb9ba8dbff-00000000.node (5.9MB)\n.d8bdfffe9ab2ddff-00000000.node (5.9MB)\n.s.PGSQL.5432 (0B)\n.s.PGSQL.5432.lock (56B)\n__KMP_REGISTERED_LIB_75079 (1.0KB)\nagent_loop_new.rs (28.5KB)\nagent_mod.rs (1.6KB)\nauth_trace.md (10.8KB)\nbench-daily.md (72B)\nbench-log.md (109B)\nbench-meeting.md (160B)\nbench-monday.md (44B)\nbench-prefs.md (61B)\nbench-project.md (213B)\nbench-reminder.md (68B)\nbench-todo.md (153B)\nbench-tuesday.md (40B)\ncac-deck.html (151.5KB)\ncircuit_breaker.rs (22.1KB)\ncli_config.rs (9.1KB)\ncli_service.rs (1.1KB)\ncommands.rs (17.5KB)\nconfig.rs (58.4KB)\nconflicts_summary.md (21.4KB)\ncost_guard.rs (11.3KB)\ndebug_forc2.py (2.8KB)\ndebug_forc3.py (3.0KB)\ndebug_forc4.py (3.0KB)\ndebug_forc5.py (4.0KB)\ndebug_forc6.py (3.9KB)\ndebug_forc7.py (2.7KB)\ndebug_forc8.py (2.9KB)\ndebug_forc_prove.py (3.0KB)\ndispatcher.rs (26.2KB)\ndoctor.rs (8.4KB)\nhygiene.rs (7.3KB)\nironclaw_blog_test.png (21.4KB)\nironclaw_browser_test_viewport.png (156.1KB)\nironclaw_linkedin_debug.png (6.5KB)\nironclaw_spot_test.txt (19B)\nkeys_chain_signatures.rs (6.1KB)\nkeys_error.rs (1.6KB)\nkeys_intents.rs (5.8KB)\nkeys_mod.rs (31.6KB)\nkeys_policy.rs (31.4KB)\nkeys_rpc.rs (8.5KB)\nkeys_signer.rs (8.0KB)\nkeys_spending.rs (5.9KB)\nkeys_transaction.rs (13.7KB)\nkeys_types.rs (17.4KB)\nleak_detection_research_summary.md (13.0KB)\nleak_detector.rs (25.3KB)\nlib_new.rs (5.0KB)\nllm_mod.rs (10.0KB)\nmain.rs (56.8KB)\nmain_bootstrap.rs (11.0KB)\nmain_rs.txt (34.1KB)\nnear_resp.json (39B)\nobs_log.rs (5.9KB)\nobs_mod.rs (2.7KB)\nobs_multi.rs (4.0KB)\nobs_noop.rs (1.5KB)\nobs_traits.rs (4.0KB)\nour_bootstrap.rs (12.6KB)\nour_config.rs (45.2KB)\npost_final_review.py (6.4KB)\npost_review.py (4.6KB)\npost_review2.py (4.8KB)\npost_review3.py (2.2KB)\npr-35.diff (137.4KB)\npr115.diff (1.6MB)\npr126.diff (336.9KB)\npr14_full_diff.txt (193.4KB)\npr39_files.txt (243B)\npr51_files.txt (548B)\npr93_diff.txt (160.4KB)\nprs.json (4.6KB)\nresponse_cache.rs (15.1KB)\nrun_hard_proof_search.sh (545B)\nrun_hard_proof_search2.sh (515B)\nrun_proof_search_v3.sh (523B)\nrun_proof_search_v4.sh (524B)\nservice.rs (11.9KB)\nsettings.rs (34.7KB)\ntest_line.py (789B)\ntest_search.txt (1.8KB)\nthread_ops.rs (38.5KB)\ntool_wit.txt (6.0KB)\ntunnel_cf.rs (3.9KB)\ntunnel_custom.rs (6.2KB)\ntunnel_mod.rs (10.4KB)\ntunnel_ngrok.rs (4.2KB)\ntunnel_none.rs (1.3KB)\ntunnel_ts.rs (3.5KB)\nwasm_capabilities.rs (15.2KB)\nwasm_capabilities_schema.rs (26.2KB)\nwasm_host.rs (25.3KB)\nwasm_wrapper.rs (14.9KB)\n```","started_at":"2026-02-17T22:02:24.602248Z","finished_at":"2026-02-17T22:02:43.748463Z","config_label":"default","error":null}
|
||||
{"task_id":"chain-write-read","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":5712,"llm_calls":3,"input_tokens":4489,"output_tokens":73,"estimated_cost_usd":0.014562,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"It says: `ironclaw spot check`.","started_at":"2026-02-17T22:02:43.749933Z","finished_at":"2026-02-17T22:02:49.462007Z","config_label":"default","error":null}
|
||||
{"task_id":"chain-shell-json","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":4582,"llm_calls":3,"input_tokens":4433,"output_tokens":94,"estimated_cost_usd":0.014709,"tool_calls":[{"name":"shell","duration_ms":0,"success":true},{"name":"json","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"[Called tool `json` with arguments: {\"data\":{\"code\":200,\"status\":\"ok\"},\"operation\":\"query\",\"path\":\"status\"}]","started_at":"2026-02-17T22:02:49.463794Z","finished_at":"2026-02-17T22:02:54.046092Z","config_label":"default","error":null}
|
||||
{"task_id":"chain-time-echo","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":4583,"llm_calls":3,"input_tokens":4404,"output_tokens":76,"estimated_cost_usd":0.014352,"tool_calls":[{"name":"time","duration_ms":0,"success":true},{"name":"echo","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"2026-02-17T22:02:55.786456+00:00","started_at":"2026-02-17T22:02:54.046306Z","finished_at":"2026-02-17T22:02:58.629389Z","config_label":"default","error":null}
|
||||
{"task_id":"robust-no-tool","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":1020,"llm_calls":1,"input_tokens":1366,"output_tokens":5,"estimated_cost_usd":0.004173,"tool_calls":[],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"Paris.","started_at":"2026-02-17T22:02:58.630502Z","finished_at":"2026-02-17T22:02:59.650655Z","config_label":"default","error":null}
|
||||
{"task_id":"robust-correct-tool","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":3980,"llm_calls":2,"input_tokens":2825,"output_tokens":50,"estimated_cost_usd":0.009225,"tool_calls":[{"name":"time","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"It’s **22:03:01 UTC** right now (**2026-02-17**).","started_at":"2026-02-17T22:02:59.650901Z","finished_at":"2026-02-17T22:03:03.631122Z","config_label":"default","error":null}
|
||||
{"task_id":"robust-json-validate","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":2954,"llm_calls":2,"input_tokens":2822,"output_tokens":43,"estimated_cost_usd":0.009111,"tool_calls":[{"name":"json","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"The JSON is valid (`valid: true`).","started_at":"2026-02-17T22:03:03.631527Z","finished_at":"2026-02-17T22:03:06.585971Z","config_label":"default","error":null}
|
||||
{"task_id":"memory-save-daily","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":4890,"llm_calls":3,"input_tokens":4547,"output_tokens":114,"estimated_cost_usd":0.015351,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"```md\n# Daily Tasks\n\n1. Review PR #42\n2. Update API docs\n3. Deploy to staging\n```","started_at":"2026-02-17T22:03:06.586556Z","finished_at":"2026-02-17T22:03:11.476648Z","config_label":"default","error":null}
|
||||
{"task_id":"memory-save-reminder","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":5607,"llm_calls":3,"input_tokens":4531,"output_tokens":95,"estimated_cost_usd":0.015018,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"The appointment is on **March 5th at 2pm**, with **Dr. Smith**.","started_at":"2026-02-17T22:03:11.476873Z","finished_at":"2026-02-17T22:03:17.084343Z","config_label":"default","error":null}
|
||||
{"task_id":"memory-save-meeting","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":6113,"llm_calls":3,"input_tokens":4750,"output_tokens":122,"estimated_cost_usd":0.01608,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"Frontend owner: **Bob** \nLaunch date: **April 15th**","started_at":"2026-02-17T22:03:17.085432Z","finished_at":"2026-02-17T22:03:23.198905Z","config_label":"default","error":null}
|
||||
{"task_id":"memory-append-log","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":9880,"llm_calls":6,"input_tokens":9941,"output_tokens":208,"estimated_cost_usd":0.032943,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true},{"name":"shell","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"Contents of `/tmp/bench-log.md`:\n1. Morning standup: discussed blockers on auth service \n2. Afternoon: resolved API timeout issue in payment gateway \n\nSummary: You discussed blockers affecting the auth service during the morning standup, and later in the afternoon you resolved an API timeout issue in the payment gateway.","started_at":"2026-02-17T22:03:23.199155Z","finished_at":"2026-02-17T22:03:33.079396Z","config_label":"default","error":null}
|
||||
{"task_id":"memory-recall-detail","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":4989,"llm_calls":3,"input_tokens":4721,"output_tokens":121,"estimated_cost_usd":0.015978,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"The project uses **PostgreSQL** as the database, and the PM is **James**.","started_at":"2026-02-17T22:03:33.080611Z","finished_at":"2026-02-17T22:03:38.069896Z","config_label":"default","error":null}
|
||||
{"task_id":"memory-todo-priority","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":5916,"llm_calls":3,"input_tokens":4676,"output_tokens":110,"estimated_cost_usd":0.015678,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"High priority tasks:\n- Fix login bug\n- Security audit","started_at":"2026-02-17T22:03:38.070162Z","finished_at":"2026-02-17T22:03:43.986831Z","config_label":"default","error":null}
|
||||
{"task_id":"memory-multifile","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":6081,"llm_calls":4,"input_tokens":6442,"output_tokens":196,"estimated_cost_usd":0.022266,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"On Tuesday you have **Sprint retro at 11am**, then **team lunch at noon**.","started_at":"2026-02-17T22:03:43.988013Z","finished_at":"2026-02-17T22:03:50.069452Z","config_label":"default","error":null}
|
||||
{"task_id":"memory-update-context","suite_id":"spot","score":{"value":1.0,"label":"pass","details":null},"trace":{"wall_time_ms":9164,"llm_calls":5,"input_tokens":8354,"output_tokens":252,"estimated_cost_usd":0.028842,"tool_calls":[{"name":"write_file","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true},{"name":"apply_patch","duration_ms":0,"success":true},{"name":"read_file","duration_ms":0,"success":true}],"turns":1,"hit_iteration_limit":false,"hit_timeout":false},"response":"Final read confirms the file now says `timezone: EST`.","started_at":"2026-02-17T22:03:50.070277Z","finished_at":"2026-02-17T22:03:59.235196Z","config_label":"default","error":null}
|
||||
@@ -1,21 +0,0 @@
|
||||
{"id": "smoke-greeting", "prompt": "Hello! Introduce yourself briefly.", "tags": ["smoke"], "assertions": {"response_matches": "(?i)(hello|hi|hey|assistant|agent|help)", "no_error": true, "max_tool_calls": 0}}
|
||||
{"id": "smoke-math", "prompt": "What is 47 * 23? Reply with just the number.", "tags": ["smoke"], "assertions": {"response_contains": ["1081"], "no_error": true, "max_tool_calls": 0}}
|
||||
{"id": "tool-echo", "prompt": "Use the echo tool to repeat the message: 'Spot check passed'", "tags": ["tool"], "assertions": {"tools_used": ["echo"], "response_contains": ["Spot check passed"], "no_error": true}}
|
||||
{"id": "tool-time", "prompt": "What is the current date and time? Use the time tool.", "tags": ["tool"], "assertions": {"tools_used": ["time"], "response_matches": "20\\d{2}", "no_error": true}}
|
||||
{"id": "tool-json-query", "prompt": "Given this JSON: {\"users\": [{\"name\": \"Alice\"}, {\"name\": \"Bob\"}]}, use the json tool to extract the second user's name.", "tags": ["tool"], "assertions": {"tools_used": ["json"], "response_contains": ["Bob"], "no_error": true}}
|
||||
{"id": "tool-shell-echo", "prompt": "Use the shell tool to run: echo 'benchmark test'", "tags": ["tool"], "assertions": {"tools_used": ["shell"], "response_contains": ["benchmark test"], "no_error": true}}
|
||||
{"id": "tool-list-dir", "prompt": "Use the list_dir tool to list the contents of the /tmp directory.", "tags": ["tool"], "assertions": {"tools_used": ["list_dir"], "no_error": true}}
|
||||
{"id": "chain-write-read", "prompt": "Write the text 'ironclaw spot check' to /tmp/ironclaw_spot_test.txt using the write_file tool, then read it back using the read_file tool and tell me what it says.", "tags": ["chain"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["ironclaw spot check"], "no_error": true}}
|
||||
{"id": "chain-shell-json", "prompt": "Run a shell command to output the JSON string '{\"status\": \"ok\", \"code\": 200}', then use the json tool to extract the status field.", "tags": ["chain"], "assertions": {"response_contains": ["ok"], "min_tool_calls": 1, "no_error": true}}
|
||||
{"id": "chain-time-echo", "prompt": "First get the current time using the time tool, then use the echo tool to repeat it back.", "tags": ["chain"], "assertions": {"tools_used": ["time", "echo"], "no_error": true}}
|
||||
{"id": "robust-no-tool", "prompt": "What is the capital of France? Answer directly without using any tools.", "tags": ["robust"], "assertions": {"response_contains": ["Paris"], "max_tool_calls": 0, "no_error": true}}
|
||||
{"id": "robust-correct-tool", "prompt": "What time is it right now?", "tags": ["robust"], "assertions": {"tools_used": ["time"], "tools_not_used": ["shell", "echo"], "no_error": true}}
|
||||
{"id": "robust-json-validate", "prompt": "Use the json tool to validate whether this is valid JSON: {\"key\": \"value\", \"num\": 42}", "tags": ["robust"], "assertions": {"tools_used": ["json"], "tools_not_used": ["shell"], "no_error": true}}
|
||||
{"id": "memory-save-daily", "prompt": "Save these daily tasks to /tmp/bench-daily.md:\n1. Review PR #42\n2. Update API docs\n3. Deploy to staging\nThen read the file back and confirm what was saved.", "tags": ["memory"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["PR #42", "docs", "staging"], "no_error": true}}
|
||||
{"id": "memory-save-reminder", "prompt": "Write a reminder to /tmp/bench-reminder.md: Dentist appointment on March 5th at 2pm with Dr. Smith. Then read the file back and tell me when the appointment is and with whom.", "tags": ["memory"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["March 5", "Smith"], "response_matches": "2(:00)?\\s*[Pp][Mm]", "no_error": true}}
|
||||
{"id": "memory-save-meeting", "prompt": "Save these meeting notes to /tmp/bench-meeting.md:\nMeeting: Project Phoenix sync\nAttendees: Alice, Bob, Carol\nDecisions:\n- Launch date: April 15th\n- Budget: $50k approved\n- Bob owns frontend, Carol owns backend\nThen read the file back and tell me who owns the frontend and what the launch date is.", "tags": ["memory"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["Bob", "frontend", "April 15"], "no_error": true}}
|
||||
{"id": "memory-append-log", "prompt": "Write 'Morning standup: discussed blockers on auth service' to /tmp/bench-log.md. Then append a new line 'Afternoon: resolved API timeout issue in payment gateway' to the same file. Finally read the full file and summarize what happened.", "tags": ["memory"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["auth", "timeout"], "min_tool_calls": 3, "no_error": true}}
|
||||
{"id": "memory-recall-detail", "prompt": "Save the following project context to /tmp/bench-project.md:\nProject Ironclad uses Rust for the backend, React for the frontend, and PostgreSQL for the database. The API is deployed on AWS ECS. The lead developer is Sarah and the PM is James. The sprint ends on March 20th.\nThen read it back and answer: What database does the project use, and who is the PM?", "tags": ["memory"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["PostgreSQL", "James"], "no_error": true}}
|
||||
{"id": "memory-todo-priority", "prompt": "Write the following to /tmp/bench-todo.md:\n- [ ] Fix login bug (priority: HIGH)\n- [ ] Write unit tests (priority: medium)\n- [ ] Update README (priority: low)\n- [ ] Security audit (priority: HIGH)\nThen read it back and tell me which tasks are high priority.", "tags": ["memory"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["login bug", "security audit"], "no_error": true}}
|
||||
{"id": "memory-multifile", "prompt": "Save 'Team standup at 9am, then client demo at 2pm' to /tmp/bench-monday.md and 'Sprint retro at 11am, team lunch at noon' to /tmp/bench-tuesday.md. Then read both files and tell me what's happening on Tuesday.", "tags": ["memory"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["retro", "lunch"], "min_tool_calls": 3, "no_error": true}}
|
||||
{"id": "memory-update-context", "prompt": "Write 'User preference: dark mode, timezone: PST, language: English' to /tmp/bench-prefs.md. Then read it back, and rewrite the file changing the timezone to EST. Finally read it one more time and confirm the timezone is now EST.", "tags": ["memory"], "assertions": {"tools_used": ["write_file", "read_file"], "response_contains": ["EST"], "min_tool_calls": 4, "no_error": true}}
|
||||
@@ -1,8 +0,0 @@
|
||||
task_timeout = "120s"
|
||||
parallelism = 1
|
||||
|
||||
[[matrix]]
|
||||
label = "default"
|
||||
|
||||
[suite_config]
|
||||
dataset_path = "benchmarks/data/spot.jsonl"
|
||||
@@ -1,243 +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
|
||||
tracing::warn!(
|
||||
task_id = %task.id,
|
||||
"LLM-as-judge scoring not implemented, returning placeholder 0.5"
|
||||
);
|
||||
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![],
|
||||
error: None,
|
||||
};
|
||||
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![],
|
||||
error: None,
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
}
|
||||
@@ -1,183 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::scoring;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskResource, TaskSubmission};
|
||||
|
||||
/// GAIA dataset entry (Hugging Face JSONL format).
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GaiaEntry {
|
||||
task_id: String,
|
||||
#[serde(alias = "Question")]
|
||||
question: String,
|
||||
#[serde(alias = "Final answer", alias = "final_answer")]
|
||||
final_answer: String,
|
||||
#[serde(alias = "Level", default)]
|
||||
level: Option<u32>,
|
||||
#[serde(alias = "file_name", default)]
|
||||
file_name: Option<String>,
|
||||
}
|
||||
|
||||
/// 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![],
|
||||
error: None,
|
||||
};
|
||||
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![],
|
||||
error: None,
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
}
|
||||
@@ -1,124 +0,0 @@
|
||||
pub mod custom;
|
||||
pub mod gaia;
|
||||
pub mod spot;
|
||||
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)"),
|
||||
("spot", "Spot checks (end-to-end user workflows)"),
|
||||
("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,
|
||||
)))
|
||||
}
|
||||
"spot" => {
|
||||
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 'spot' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
Ok(Box::new(spot::SpotSuite::new(dataset_path)))
|
||||
}
|
||||
"tau_bench" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'tau_bench' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
let domain = suite_map
|
||||
.get("domain")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("retail")
|
||||
.to_string();
|
||||
Ok(Box::new(tau_bench::TauBenchSuite::new(
|
||||
dataset_path,
|
||||
domain,
|
||||
)))
|
||||
}
|
||||
"swe_bench" => {
|
||||
let dataset_path = suite_map
|
||||
.get("dataset_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string())
|
||||
.ok_or_else(|| {
|
||||
BenchError::Config(
|
||||
"suite_config.dataset_path is required for 'swe_bench' suite".to_string(),
|
||||
)
|
||||
})?;
|
||||
let workspace_dir = suite_map
|
||||
.get("workspace_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("/tmp/swe-bench")
|
||||
.to_string();
|
||||
let use_docker = suite_map
|
||||
.get("use_docker")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(false);
|
||||
Ok(Box::new(swe_bench::SweBenchSuite::new(
|
||||
dataset_path,
|
||||
workspace_dir,
|
||||
use_docker,
|
||||
)))
|
||||
}
|
||||
_ => {
|
||||
let available = KNOWN_SUITES
|
||||
.iter()
|
||||
.map(|(id, _)| *id)
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
Err(BenchError::SuiteNotFound {
|
||||
name: name.to_string(),
|
||||
available,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,504 +0,0 @@
|
||||
use std::collections::HashSet;
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskSubmission};
|
||||
|
||||
/// Multi-criterion assertions for a spot check scenario.
|
||||
///
|
||||
/// Each field generates one or more individual checks. The final score is
|
||||
/// `passed_checks / total_checks`, giving a value between 0.0 and 1.0.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SpotAssertions {
|
||||
/// All must appear in the response (case-insensitive).
|
||||
#[serde(default)]
|
||||
pub response_contains: Vec<String>,
|
||||
|
||||
/// None may appear in the response (case-insensitive).
|
||||
#[serde(default)]
|
||||
pub response_not_contains: Vec<String>,
|
||||
|
||||
/// Each tool name must appear in the tool_calls list (checked by name,
|
||||
/// not by count; duplicates in tool_calls are collapsed).
|
||||
#[serde(default)]
|
||||
pub tools_used: Vec<String>,
|
||||
|
||||
/// None of these tool names may appear in the tool_calls list.
|
||||
#[serde(default)]
|
||||
pub tools_not_used: Vec<String>,
|
||||
|
||||
/// Regex pattern the response must match.
|
||||
#[serde(default)]
|
||||
pub response_matches: Option<String>,
|
||||
|
||||
/// Hard fail if the task produced an error.
|
||||
#[serde(default)]
|
||||
pub no_error: bool,
|
||||
|
||||
/// Minimum number of tool calls expected (counts duplicates).
|
||||
#[serde(default)]
|
||||
pub min_tool_calls: Option<usize>,
|
||||
|
||||
/// Maximum number of tool calls allowed (counts duplicates).
|
||||
#[serde(default)]
|
||||
pub max_tool_calls: Option<usize>,
|
||||
}
|
||||
|
||||
impl SpotAssertions {
|
||||
/// Evaluate all assertions against a submission, returning (score, failure_details).
|
||||
pub fn evaluate(&self, submission: &TaskSubmission) -> (f64, Vec<String>) {
|
||||
let mut passed: usize = 0;
|
||||
let mut total: usize = 0;
|
||||
let mut failures: Vec<String> = Vec::new();
|
||||
|
||||
// Hard fail: error check
|
||||
if self.no_error {
|
||||
total += 1;
|
||||
if let Some(ref err) = submission.error {
|
||||
failures.push(format!("no_error: task errored with: {err}"));
|
||||
// Hard fail: return 0.0 immediately
|
||||
return (0.0, failures);
|
||||
}
|
||||
passed += 1;
|
||||
}
|
||||
|
||||
let response_lower = submission.response.to_lowercase();
|
||||
|
||||
// response_contains: all must appear
|
||||
for needle in &self.response_contains {
|
||||
total += 1;
|
||||
if response_lower.contains(&needle.to_lowercase()) {
|
||||
passed += 1;
|
||||
} else {
|
||||
failures.push(format!("response_contains: missing \"{needle}\""));
|
||||
}
|
||||
}
|
||||
|
||||
// response_not_contains: none may appear
|
||||
for needle in &self.response_not_contains {
|
||||
total += 1;
|
||||
if response_lower.contains(&needle.to_lowercase()) {
|
||||
failures.push(format!("response_not_contains: found \"{needle}\""));
|
||||
} else {
|
||||
passed += 1;
|
||||
}
|
||||
}
|
||||
|
||||
let tool_set: HashSet<&str> = submission.tool_calls.iter().map(|s| s.as_str()).collect();
|
||||
|
||||
// tools_used: each must appear
|
||||
for tool in &self.tools_used {
|
||||
total += 1;
|
||||
if tool_set.contains(tool.as_str()) {
|
||||
passed += 1;
|
||||
} else {
|
||||
failures.push(format!("tools_used: \"{tool}\" not called"));
|
||||
}
|
||||
}
|
||||
|
||||
// tools_not_used: none may appear
|
||||
for tool in &self.tools_not_used {
|
||||
total += 1;
|
||||
if tool_set.contains(tool.as_str()) {
|
||||
failures.push(format!("tools_not_used: \"{tool}\" was called"));
|
||||
} else {
|
||||
passed += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// response_matches: regex pattern
|
||||
if let Some(ref pattern) = self.response_matches {
|
||||
total += 1;
|
||||
match Regex::new(pattern) {
|
||||
Ok(re) => {
|
||||
if re.is_match(&submission.response) {
|
||||
passed += 1;
|
||||
} else {
|
||||
failures.push(format!("response_matches: /{pattern}/ did not match"));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
failures.push(format!("response_matches: bad regex: {e}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let call_count = submission.tool_calls.len();
|
||||
|
||||
// min_tool_calls
|
||||
if let Some(min) = self.min_tool_calls {
|
||||
total += 1;
|
||||
if call_count >= min {
|
||||
passed += 1;
|
||||
} else {
|
||||
failures.push(format!(
|
||||
"min_tool_calls: expected >= {min}, got {call_count}"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// max_tool_calls
|
||||
if let Some(max) = self.max_tool_calls {
|
||||
total += 1;
|
||||
if call_count <= max {
|
||||
passed += 1;
|
||||
} else {
|
||||
failures.push(format!(
|
||||
"max_tool_calls: expected <= {max}, got {call_count}"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if total == 0 {
|
||||
return (1.0, failures);
|
||||
}
|
||||
|
||||
let score = passed as f64 / total as f64;
|
||||
(score, failures)
|
||||
}
|
||||
}
|
||||
|
||||
/// JSONL entry for a spot check scenario.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct SpotEntry {
|
||||
id: String,
|
||||
prompt: String,
|
||||
#[serde(default)]
|
||||
context: Option<String>,
|
||||
#[serde(default)]
|
||||
tags: Vec<String>,
|
||||
#[serde(default)]
|
||||
assertions: SpotAssertions,
|
||||
}
|
||||
|
||||
/// Spot benchmark suite: end-to-end checks for real user workflows.
|
||||
///
|
||||
/// Tests conversation, individual tool use, multi-tool chaining, and robustness.
|
||||
/// Each task declares multi-criterion assertions scored as passed/total.
|
||||
pub struct SpotSuite {
|
||||
dataset_path: PathBuf,
|
||||
}
|
||||
|
||||
impl SpotSuite {
|
||||
pub fn new(dataset_path: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
dataset_path: dataset_path.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BenchSuite for SpotSuite {
|
||||
fn name(&self) -> &str {
|
||||
"Spot Checks"
|
||||
}
|
||||
|
||||
fn id(&self) -> &str {
|
||||
"spot"
|
||||
}
|
||||
|
||||
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: SpotEntry = serde_json::from_str(trimmed)
|
||||
.map_err(|e| BenchError::Config(format!("spot line {}: {}", line_num + 1, e)))?;
|
||||
|
||||
let metadata = serde_json::json!({
|
||||
"assertions": serde_json::to_value(&entry.assertions)
|
||||
.map_err(|e| BenchError::Config(format!("spot {}: {}", entry.id, e)))?,
|
||||
});
|
||||
|
||||
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 assertions: SpotAssertions = task
|
||||
.metadata
|
||||
.get("assertions")
|
||||
.ok_or_else(|| BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: "missing assertions in metadata".to_string(),
|
||||
})
|
||||
.and_then(|v| {
|
||||
serde_json::from_value(v.clone()).map_err(|e| BenchError::Scoring {
|
||||
task_id: task.id.clone(),
|
||||
reason: format!("bad assertions: {e}"),
|
||||
})
|
||||
})?;
|
||||
|
||||
let (score, failures) = assertions.evaluate(submission);
|
||||
|
||||
if score >= 1.0 {
|
||||
Ok(BenchScore::pass())
|
||||
} else if score <= 0.0 {
|
||||
Ok(BenchScore::fail(failures.join("; ")))
|
||||
} else {
|
||||
Ok(BenchScore::partial(score, failures.join("; ")))
|
||||
}
|
||||
}
|
||||
|
||||
fn additional_tools(&self) -> Vec<Arc<dyn ironclaw::tools::Tool>> {
|
||||
vec![
|
||||
Arc::new(ironclaw::tools::builtin::ShellTool::new()),
|
||||
Arc::new(ironclaw::tools::builtin::ReadFileTool::new()),
|
||||
Arc::new(ironclaw::tools::builtin::WriteFileTool::new()),
|
||||
Arc::new(ironclaw::tools::builtin::ListDirTool::new()),
|
||||
Arc::new(ironclaw::tools::builtin::ApplyPatchTool::new()),
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
fn make_submission(
|
||||
response: &str,
|
||||
tool_calls: Vec<&str>,
|
||||
error: Option<&str>,
|
||||
) -> TaskSubmission {
|
||||
TaskSubmission {
|
||||
response: response.to_string(),
|
||||
conversation: vec![],
|
||||
tool_calls: tool_calls.into_iter().map(|s| s.to_string()).collect(),
|
||||
error: error.map(|s| s.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_all_pass() {
|
||||
let assertions = SpotAssertions {
|
||||
response_contains: vec!["hello".to_string()],
|
||||
tools_used: vec!["echo".to_string()],
|
||||
no_error: true,
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("Hello, world!", vec!["echo"], None);
|
||||
let (score, failures) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 1.0);
|
||||
assert!(failures.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_hard_fail_on_error() {
|
||||
let assertions = SpotAssertions {
|
||||
no_error: true,
|
||||
response_contains: vec!["hello".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("Hello!", vec![], Some("timeout after 60s"));
|
||||
let (score, failures) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 0.0);
|
||||
assert!(failures[0].contains("no_error"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_partial_score() {
|
||||
let assertions = SpotAssertions {
|
||||
response_contains: vec!["alpha".to_string(), "beta".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("alpha is here but not the other", vec![], None);
|
||||
let (score, failures) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 0.5);
|
||||
assert_eq!(failures.len(), 1);
|
||||
assert!(failures[0].contains("beta"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_response_not_contains() {
|
||||
let assertions = SpotAssertions {
|
||||
response_not_contains: vec!["error".to_string(), "fail".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("This is an error message", vec![], None);
|
||||
let (score, failures) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 0.5);
|
||||
assert_eq!(failures.len(), 1);
|
||||
assert!(failures[0].contains("error"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tools_used_and_not_used() {
|
||||
let assertions = SpotAssertions {
|
||||
tools_used: vec!["time".to_string()],
|
||||
tools_not_used: vec!["shell".to_string(), "echo".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("The time is now", vec!["time"], None);
|
||||
let (score, failures) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 1.0);
|
||||
assert!(failures.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tools_not_used_fails() {
|
||||
let assertions = SpotAssertions {
|
||||
tools_not_used: vec!["shell".to_string()],
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("result", vec!["shell", "time"], None);
|
||||
let (score, _) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_response_matches_regex() {
|
||||
let assertions = SpotAssertions {
|
||||
response_matches: Some(r"\d{4}".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("The year is 2026", vec![], None);
|
||||
let (score, failures) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 1.0);
|
||||
assert!(failures.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_response_matches_regex_fail() {
|
||||
let assertions = SpotAssertions {
|
||||
response_matches: Some(r"^\d+$".to_string()),
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("not a number", vec![], None);
|
||||
let (score, _) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_min_max_tool_calls() {
|
||||
let assertions = SpotAssertions {
|
||||
min_tool_calls: Some(2),
|
||||
max_tool_calls: Some(4),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Within range
|
||||
let sub = make_submission("ok", vec!["a", "b", "c"], None);
|
||||
let (score, _) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 1.0);
|
||||
|
||||
// Too few
|
||||
let sub = make_submission("ok", vec!["a"], None);
|
||||
let (score, failures) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 0.5);
|
||||
assert!(failures[0].contains("min_tool_calls"));
|
||||
|
||||
// Too many
|
||||
let sub = make_submission("ok", vec!["a", "b", "c", "d", "e"], None);
|
||||
let (score, failures) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 0.5);
|
||||
assert!(failures[0].contains("max_tool_calls"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_max_zero_tool_calls() {
|
||||
let assertions = SpotAssertions {
|
||||
max_tool_calls: Some(0),
|
||||
..Default::default()
|
||||
};
|
||||
let sub = make_submission("just talking", vec![], None);
|
||||
let (score, _) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 1.0);
|
||||
|
||||
let sub = make_submission("oops", vec!["echo"], None);
|
||||
let (score, _) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_assertions() {
|
||||
let assertions = SpotAssertions::default();
|
||||
let sub = make_submission("anything", vec!["whatever"], None);
|
||||
let (score, _) = assertions.evaluate(&sub);
|
||||
assert_eq!(score, 1.0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_spot_load_tasks() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("spot.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "s1", "prompt": "Hello", "tags": ["smoke"], "assertions": {{"response_contains": ["hello"], "no_error": true}}}}"#
|
||||
)
|
||||
.unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "s2", "prompt": "Echo test", "assertions": {{"tools_used": ["echo"]}}}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = SpotSuite::new(&path);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
assert_eq!(tasks.len(), 2);
|
||||
assert_eq!(tasks[0].id, "s1");
|
||||
assert_eq!(tasks[1].id, "s2");
|
||||
assert!(tasks[0].tags.contains(&"smoke".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_spot_scoring() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("spot.jsonl");
|
||||
let mut file = std::fs::File::create(&path).unwrap();
|
||||
writeln!(
|
||||
file,
|
||||
r#"{{"id": "s1", "prompt": "Hello", "assertions": {{"response_contains": ["hello", "world"], "no_error": true}}}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = SpotSuite::new(&path);
|
||||
let tasks = suite.load_tasks().await.unwrap();
|
||||
|
||||
// Full pass
|
||||
let sub = make_submission("Hello World!", vec![], None);
|
||||
let score = suite.score(&tasks[0], &sub).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
assert_eq!(score.label, "pass");
|
||||
|
||||
// Partial
|
||||
let sub = make_submission("Hello there", vec![], None);
|
||||
let score = suite.score(&tasks[0], &sub).await.unwrap();
|
||||
assert!(score.value > 0.0 && score.value < 1.0);
|
||||
assert_eq!(score.label, "partial");
|
||||
|
||||
// Error hard fail
|
||||
let sub = make_submission("Hello World!", vec![], Some("boom"));
|
||||
let score = suite.score(&tasks[0], &sub).await.unwrap();
|
||||
assert_eq!(score.value, 0.0);
|
||||
assert_eq!(score.label, "fail");
|
||||
}
|
||||
}
|
||||
@@ -1,416 +0,0 @@
|
||||
use std::io::BufRead;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use regex::Regex;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::error::BenchError;
|
||||
use crate::suite::{BenchScore, BenchSuite, BenchTask, TaskSubmission};
|
||||
|
||||
/// Validate that a string is safe for use as a filesystem path component.
|
||||
/// Allows alphanumerics, hyphens, underscores, dots, and forward slashes (for nested paths).
|
||||
/// Rejects absolute paths, `..` traversal, and shell metacharacters.
|
||||
fn is_safe_path_component(s: &str) -> bool {
|
||||
!s.is_empty()
|
||||
&& !s.starts_with('/')
|
||||
&& !s.contains("..")
|
||||
&& s.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.' | '/'))
|
||||
}
|
||||
|
||||
/// Validate that a repo string matches the expected `owner/repo` GitHub format.
|
||||
fn is_valid_github_repo(repo: &str) -> bool {
|
||||
// Match "owner/repo" where both parts are alphanumeric with hyphens/underscores/dots
|
||||
static REPO_PATTERN: std::sync::LazyLock<Regex> =
|
||||
std::sync::LazyLock::new(|| Regex::new(r"^[a-zA-Z0-9._-]+/[a-zA-Z0-9._-]+$").unwrap());
|
||||
REPO_PATTERN.is_match(repo)
|
||||
}
|
||||
|
||||
/// Validate that a string looks like a git ref (hex SHA or valid ref name).
|
||||
fn is_valid_git_ref(s: &str) -> bool {
|
||||
!s.is_empty()
|
||||
&& s.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.' | '/'))
|
||||
&& !s.contains("..")
|
||||
}
|
||||
|
||||
/// 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))
|
||||
})?;
|
||||
|
||||
if !is_safe_path_component(&entry.instance_id) {
|
||||
return Err(BenchError::Config(format!(
|
||||
"swe_bench line {}: unsafe instance_id \"{}\"",
|
||||
line_num + 1,
|
||||
entry.instance_id,
|
||||
)));
|
||||
}
|
||||
if !is_valid_github_repo(&entry.repo) {
|
||||
return Err(BenchError::Config(format!(
|
||||
"swe_bench line {}: invalid repo format \"{}\"",
|
||||
line_num + 1,
|
||||
entry.repo,
|
||||
)));
|
||||
}
|
||||
if !is_valid_git_ref(&entry.base_commit) {
|
||||
return Err(BenchError::Config(format!(
|
||||
"swe_bench line {}: invalid base_commit \"{}\"",
|
||||
line_num + 1,
|
||||
entry.base_commit,
|
||||
)));
|
||||
}
|
||||
|
||||
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.
|
||||
tracing::warn!(
|
||||
task_id = %task.id,
|
||||
"SWE-bench test execution not implemented, returning placeholder 0.25"
|
||||
);
|
||||
Ok(BenchScore::partial(
|
||||
0.25,
|
||||
"test execution not yet implemented; partial credit for response",
|
||||
))
|
||||
} else {
|
||||
tracing::warn!(
|
||||
task_id = %task.id,
|
||||
"no test_patch available, returning placeholder 0.25"
|
||||
);
|
||||
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![],
|
||||
error: None,
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_safe_path_component() {
|
||||
assert!(is_safe_path_component("django__django-12345"));
|
||||
assert!(is_safe_path_component("org/repo"));
|
||||
assert!(is_safe_path_component("abc123"));
|
||||
assert!(!is_safe_path_component(""));
|
||||
assert!(!is_safe_path_component("../../etc/passwd"));
|
||||
assert!(!is_safe_path_component("/etc/passwd"));
|
||||
assert!(!is_safe_path_component("foo;rm -rf /"));
|
||||
assert!(!is_safe_path_component("foo bar"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_valid_github_repo() {
|
||||
assert!(is_valid_github_repo("django/django"));
|
||||
assert!(is_valid_github_repo("org/repo-name"));
|
||||
assert!(is_valid_github_repo("Org.Name/Repo_v2"));
|
||||
assert!(!is_valid_github_repo(""));
|
||||
assert!(!is_valid_github_repo("no-slash"));
|
||||
assert!(!is_valid_github_repo("too/many/slashes"));
|
||||
assert!(!is_valid_github_repo("spa ce/repo"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_valid_git_ref() {
|
||||
assert!(is_valid_git_ref("abc123"));
|
||||
assert!(is_valid_git_ref("deadbeef0123456789abcdef0123456789abcdef"));
|
||||
assert!(is_valid_git_ref("v1.2.3"));
|
||||
assert!(is_valid_git_ref("main"));
|
||||
assert!(!is_valid_git_ref(""));
|
||||
assert!(!is_valid_git_ref("bad..ref"));
|
||||
assert!(!is_valid_git_ref("has space"));
|
||||
assert!(!is_valid_git_ref("semi;colon"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_swe_bench_rejects_path_traversal() {
|
||||
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": "../../etc/passwd", "repo": "org/repo", "base_commit": "abc", "problem_statement": "evil"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = SweBenchSuite::new(&path, "/tmp/swe-test", false);
|
||||
let err = suite.load_tasks().await.unwrap_err();
|
||||
assert!(err.to_string().contains("unsafe instance_id"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_swe_bench_rejects_bad_repo() {
|
||||
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": "task1", "repo": "not-a-repo-format", "base_commit": "abc", "problem_statement": "bad"}}"#
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let suite = SweBenchSuite::new(&path, "/tmp/swe-test", false);
|
||||
let err = suite.load_tasks().await.unwrap_err();
|
||||
assert!(err.to_string().contains("invalid repo format"));
|
||||
}
|
||||
}
|
||||
@@ -1,233 +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);
|
||||
}
|
||||
}
|
||||
|
||||
// Multi-turn simulation requires an LLM to play the customer role.
|
||||
// Until that's implemented, every scenario is single-turn only.
|
||||
// TODO: Use LLM to simulate customer based on user_persona.
|
||||
tracing::warn!(
|
||||
task_id = %task.id,
|
||||
"multi-turn simulation not implemented, ending after first turn"
|
||||
);
|
||||
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()],
|
||||
error: None,
|
||||
};
|
||||
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()],
|
||||
error: None,
|
||||
};
|
||||
let score = suite.score(&tasks[0], &submission).await.unwrap();
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
}
|
||||
@@ -1,259 +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;
|
||||
|
||||
/// Truncate a string to at most `max_bytes` without splitting a UTF-8 character.
|
||||
fn truncate_str(s: &str, max_bytes: usize) -> &str {
|
||||
if s.len() <= max_bytes {
|
||||
return s;
|
||||
}
|
||||
let mut end = max_bytes;
|
||||
while !s.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
&s[..end]
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
}
|
||||
|
||||
#[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} -> {}",
|
||||
truncate_str(preview, 100)
|
||||
));
|
||||
}
|
||||
StatusUpdate::StreamChunk(_) => {}
|
||||
StatusUpdate::Status(ref msg) => {
|
||||
cap.status_log.push(format!("status: {msg}"));
|
||||
}
|
||||
StatusUpdate::JobStarted {
|
||||
ref job_id,
|
||||
ref title,
|
||||
..
|
||||
} => {
|
||||
cap.status_log
|
||||
.push(format!("job_started: {job_id} ({title})"));
|
||||
}
|
||||
StatusUpdate::AuthRequired {
|
||||
ref extension_name, ..
|
||||
} => {
|
||||
cap.status_log
|
||||
.push(format!("auth_required: {extension_name} (auto-skipped)"));
|
||||
}
|
||||
StatusUpdate::AuthCompleted {
|
||||
ref extension_name,
|
||||
success,
|
||||
..
|
||||
} => {
|
||||
cap.status_log.push(format!(
|
||||
"auth_completed: {extension_name} success={success}"
|
||||
));
|
||||
}
|
||||
}
|
||||
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: {}",
|
||||
truncate_str(&response.content, 100)
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> Result<(), ChannelError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn shutdown(&self) -> Result<(), ChannelError> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_captures_responses() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
let capture = channel.capture();
|
||||
|
||||
let msg = IncomingMessage::new("bench", "user", "hello");
|
||||
let response = OutgoingResponse::text("world");
|
||||
channel.respond(&msg, response).await.unwrap();
|
||||
|
||||
let cap = capture.lock().await;
|
||||
assert_eq!(cap.responses.len(), 1);
|
||||
assert_eq!(cap.responses[0], "world");
|
||||
assert_eq!(cap.conversation.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_auto_approves() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
// start() to consume the receiver
|
||||
let _stream = channel.start().await.unwrap();
|
||||
|
||||
let status = StatusUpdate::ApprovalNeeded {
|
||||
request_id: "req-1".to_string(),
|
||||
tool_name: "shell".to_string(),
|
||||
description: "run ls".to_string(),
|
||||
parameters: serde_json::json!({}),
|
||||
};
|
||||
channel
|
||||
.send_status(status, &serde_json::Value::Null)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// The approval message was sent through msg_tx,
|
||||
// which means the stream would receive it.
|
||||
// We can't easily read from the stream in this test without
|
||||
// consuming it, but we can verify the status log.
|
||||
let capture_arc = channel.capture();
|
||||
let cap = capture_arc.lock().await;
|
||||
assert!(cap.status_log.iter().any(|s| s.contains("auto_approved")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bench_channel_captures_tool_events() {
|
||||
let (channel, _tx) = BenchChannel::new();
|
||||
|
||||
let status = StatusUpdate::ToolCompleted {
|
||||
name: "echo".to_string(),
|
||||
success: true,
|
||||
};
|
||||
channel
|
||||
.send_status(status, &serde_json::Value::Null)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let capture_arc = channel.capture();
|
||||
let cap = capture_arc.lock().await;
|
||||
assert_eq!(cap.tool_calls.len(), 1);
|
||||
assert_eq!(cap.tool_calls[0].name, "echo");
|
||||
assert!(cap.tool_calls[0].success);
|
||||
}
|
||||
}
|
||||
@@ -1,205 +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,
|
||||
|
||||
/// 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>,
|
||||
}
|
||||
|
||||
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)?;
|
||||
if config.matrix.is_empty() {
|
||||
return Err(BenchError::Config(
|
||||
"config must have at least one [[matrix]] entry".to_string(),
|
||||
));
|
||||
}
|
||||
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(),
|
||||
parallelism: default_parallelism(),
|
||||
matrix: vec![MatrixEntry { label, model }],
|
||||
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_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.parallelism, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_rejects_empty_matrix() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("empty.toml");
|
||||
std::fs::write(
|
||||
&path,
|
||||
r#"
|
||||
results_dir = "./results"
|
||||
task_timeout = "60s"
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let err = BenchConfig::from_file(&path).unwrap_err();
|
||||
assert!(
|
||||
err.to_string().contains("at least one [[matrix]]"),
|
||||
"got: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_from_toml() {
|
||||
let toml_str = r#"
|
||||
results_dir = "./my-results"
|
||||
task_timeout = "60s"
|
||||
parallelism = 2
|
||||
|
||||
[[matrix]]
|
||||
label = "fast"
|
||||
model = "gpt-4o-mini"
|
||||
|
||||
[[matrix]]
|
||||
label = "full"
|
||||
model = "claude-3-5-sonnet"
|
||||
|
||||
[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.parallelism, 2);
|
||||
assert_eq!(config.matrix.len(), 2);
|
||||
assert_eq!(
|
||||
config.suite_config_str("dataset_path").unwrap(),
|
||||
"./data/test.jsonl"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,31 +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("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),
|
||||
}
|
||||
@@ -1,251 +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 rust_decimal::prelude::ToPrimitive;
|
||||
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)]
|
||||
#[allow(dead_code)]
|
||||
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;
|
||||
total.to_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,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
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,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_records_metrics() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
|
||||
assert_eq!(instrumented.call_count(), 1);
|
||||
assert_eq!(instrumented.total_input_tokens(), 100);
|
||||
assert_eq!(instrumented.total_output_tokens(), 50);
|
||||
|
||||
let records = instrumented.take_records().await;
|
||||
assert_eq!(records.len(), 1);
|
||||
assert_eq!(records[0].input_tokens, 100);
|
||||
assert!(!records[0].had_tool_calls);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_cost_calculation() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
|
||||
// 100 * 0.000003 + 50 * 0.000015 = 0.0003 + 0.00075 = 0.00105
|
||||
let cost = instrumented.estimated_cost();
|
||||
assert!((cost - 0.00105).abs() < 0.0001);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_instrumented_reset() {
|
||||
let inner = Arc::new(FakeLlm);
|
||||
let instrumented = InstrumentedLlm::new(inner);
|
||||
|
||||
let request = CompletionRequest::new(vec![ChatMessage::user("hello")]);
|
||||
let _ = instrumented.complete(request).await.unwrap();
|
||||
assert_eq!(instrumented.call_count(), 1);
|
||||
|
||||
instrumented.reset().await;
|
||||
assert_eq!(instrumented.call_count(), 0);
|
||||
assert_eq!(instrumented.total_input_tokens(), 0);
|
||||
|
||||
let records = instrumented.take_records().await;
|
||||
assert!(records.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -1,313 +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, spot, 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().await.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(),
|
||||
})
|
||||
.await;
|
||||
session.ensure_authenticated().await?;
|
||||
|
||||
let llm = ironclaw::llm::create_llm_provider(&ironclaw_config.llm, session)?;
|
||||
let safety = Arc::new(ironclaw::safety::SafetyLayer::new(&ironclaw_config.safety));
|
||||
|
||||
let runner = runner::BenchRunner::new(bench_suite, bench_config.clone(), llm, safety);
|
||||
|
||||
// Run for each matrix entry
|
||||
for matrix_entry in &bench_config.matrix {
|
||||
let run_id = runner
|
||||
.run(
|
||||
matrix_entry,
|
||||
sample,
|
||||
task_ids.as_deref(),
|
||||
tags.as_deref(),
|
||||
resume,
|
||||
)
|
||||
.await?;
|
||||
println!("Run complete: {}", run_id);
|
||||
}
|
||||
}
|
||||
Commands::Results {
|
||||
run_id,
|
||||
format,
|
||||
results_dir,
|
||||
} => {
|
||||
let base = results_dir.unwrap_or_else(|| PathBuf::from("./bench-results"));
|
||||
let uuid = if run_id == "latest" {
|
||||
results::find_latest_run(&base)?
|
||||
.ok_or_else(|| anyhow::anyhow!("No runs found in {}", base.display()))?
|
||||
} else {
|
||||
Uuid::parse_str(&run_id)?
|
||||
};
|
||||
|
||||
let json_path = results::run_json_path(&base, uuid);
|
||||
let jsonl_path = results::tasks_jsonl_path(&base, uuid);
|
||||
|
||||
let run = results::read_run_result(&json_path)?;
|
||||
let tasks = results::read_task_results(&jsonl_path)?;
|
||||
|
||||
match format {
|
||||
ResultsFormat::Table => {
|
||||
results::print_results_table(&tasks, &run);
|
||||
}
|
||||
ResultsFormat::Json => {
|
||||
let output = serde_json::json!({
|
||||
"run": run,
|
||||
"tasks": tasks,
|
||||
});
|
||||
println!("{}", serde_json::to_string_pretty(&output)?);
|
||||
}
|
||||
ResultsFormat::Csv => {
|
||||
println!("task_id,score,label,tokens,cost,turns,time_s");
|
||||
for task in &tasks {
|
||||
println!(
|
||||
"{},{:.3},{},{},{:.4},{},{:.1}",
|
||||
task.task_id,
|
||||
task.score.value,
|
||||
task.score.label,
|
||||
task.trace.input_tokens + task.trace.output_tokens,
|
||||
task.trace.estimated_cost_usd,
|
||||
task.trace.turns,
|
||||
task.trace.wall_time_ms as f64 / 1000.0,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Commands::Compare {
|
||||
baseline,
|
||||
comparison,
|
||||
results_dir,
|
||||
} => {
|
||||
let base = results_dir.unwrap_or_else(|| PathBuf::from("./bench-results"));
|
||||
|
||||
let baseline_run = results::read_run_result(&results::run_json_path(&base, baseline))?;
|
||||
let comparison_run =
|
||||
results::read_run_result(&results::run_json_path(&base, comparison))?;
|
||||
|
||||
println!("\nComparison: {} vs {}\n", baseline, comparison);
|
||||
println!(
|
||||
"{:<20} {:>12} {:>12} {:>10}",
|
||||
"Metric", "Baseline", "Comparison", "Delta"
|
||||
);
|
||||
println!("{}", "-".repeat(58));
|
||||
|
||||
let pass_delta = comparison_run.pass_rate - baseline_run.pass_rate;
|
||||
println!(
|
||||
"{:<20} {:>11.1}% {:>11.1}% {:>+9.1}%",
|
||||
"Pass rate",
|
||||
baseline_run.pass_rate * 100.0,
|
||||
comparison_run.pass_rate * 100.0,
|
||||
pass_delta * 100.0,
|
||||
);
|
||||
|
||||
let score_delta = comparison_run.avg_score - baseline_run.avg_score;
|
||||
println!(
|
||||
"{:<20} {:>12.3} {:>12.3} {:>+10.3}",
|
||||
"Avg score", baseline_run.avg_score, comparison_run.avg_score, score_delta,
|
||||
);
|
||||
|
||||
let cost_delta = comparison_run.total_cost_usd - baseline_run.total_cost_usd;
|
||||
println!(
|
||||
"{:<20} {:>11.4}$ {:>11.4}$ {:>+9.4}$",
|
||||
"Total cost",
|
||||
baseline_run.total_cost_usd,
|
||||
comparison_run.total_cost_usd,
|
||||
cost_delta,
|
||||
);
|
||||
|
||||
let time_b = baseline_run.total_wall_time_ms as f64 / 1000.0;
|
||||
let time_c = comparison_run.total_wall_time_ms as f64 / 1000.0;
|
||||
println!(
|
||||
"{:<20} {:>11.1}s {:>11.1}s {:>+9.1}s",
|
||||
"Total time",
|
||||
time_b,
|
||||
time_c,
|
||||
time_c - time_b,
|
||||
);
|
||||
|
||||
println!(
|
||||
"{:<20} {:>12} {:>12}",
|
||||
"Model", baseline_run.model, comparison_run.model,
|
||||
);
|
||||
println!();
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -1,473 +0,0 @@
|
||||
use std::collections::HashSet;
|
||||
use std::io::{BufRead, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
#[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,
|
||||
/// Short git commit hash at the time of the run.
|
||||
#[serde(default)]
|
||||
pub commit_hash: 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.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn from_tasks(
|
||||
run_id: Uuid,
|
||||
suite_id: &str,
|
||||
config_label: &str,
|
||||
model: &str,
|
||||
commit_hash: &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(),
|
||||
commit_hash: commit_hash.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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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(())
|
||||
}
|
||||
|
||||
/// Overwrite the JSONL file with the given results (used after scoring).
|
||||
pub fn write_task_results(path: &Path, results: &[TaskResult]) -> Result<(), BenchError> {
|
||||
let mut file = std::fs::File::create(path)?;
|
||||
for result in results {
|
||||
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).
|
||||
///
|
||||
/// Only includes tasks that have been scored (label != "pending"). Tasks that
|
||||
/// were written but not scored (e.g., from an interrupted run) will be re-executed.
|
||||
pub fn completed_task_ids(path: &Path) -> Result<HashSet<String>, BenchError> {
|
||||
let results = read_task_results(path)?;
|
||||
Ok(results
|
||||
.into_iter()
|
||||
.filter(|r| r.score.label != "pending")
|
||||
.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 the modification time of its `run.json`.
|
||||
///
|
||||
/// Falls back to `tasks.jsonl` mtime, then directory mtime. This avoids the
|
||||
/// issue where modifying files inside a directory doesn't update the directory's
|
||||
/// mtime on many filesystems.
|
||||
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 dir_path = e.path();
|
||||
// Prefer run.json mtime, fall back to tasks.jsonl, then directory
|
||||
let modified = std::fs::metadata(dir_path.join("run.json"))
|
||||
.and_then(|m| m.modified())
|
||||
.or_else(|_| {
|
||||
std::fs::metadata(dir_path.join("tasks.jsonl")).and_then(|m| m.modified())
|
||||
})
|
||||
.or_else(|_| e.metadata().and_then(|m| m.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!();
|
||||
let commit_suffix = if run.commit_hash.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(" | Commit: {}", run.commit_hash)
|
||||
};
|
||||
println!(
|
||||
"Run: {} | Suite: {} | Model: {}{}",
|
||||
run.run_id, run.suite_id, run.model, commit_suffix
|
||||
);
|
||||
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 {
|
||||
let truncated: String = task.task_id.chars().take(25).collect();
|
||||
format!("{truncated}...")
|
||||
} 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",
|
||||
"abc1234",
|
||||
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"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_task_results_overwrites() {
|
||||
let dir = tempfile::tempdir().expect("temp dir");
|
||||
let path = dir.path().join("tasks.jsonl");
|
||||
|
||||
// Write initial "pending" result via append
|
||||
let pending = TaskResult {
|
||||
task_id: "t1".to_string(),
|
||||
suite_id: "spot".to_string(),
|
||||
score: BenchScore {
|
||||
value: 0.0,
|
||||
label: "pending".to_string(),
|
||||
details: None,
|
||||
},
|
||||
trace: Trace {
|
||||
wall_time_ms: 100,
|
||||
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: "42".to_string(),
|
||||
started_at: Utc::now(),
|
||||
finished_at: Utc::now(),
|
||||
config_label: "default".to_string(),
|
||||
error: None,
|
||||
};
|
||||
append_task_result(&path, &pending).expect("append");
|
||||
|
||||
// Verify pending score
|
||||
let before = read_task_results(&path).expect("read");
|
||||
assert_eq!(before.len(), 1);
|
||||
assert_eq!(before[0].score.label, "pending");
|
||||
|
||||
// Overwrite with scored result
|
||||
let mut scored = pending;
|
||||
scored.score = BenchScore::pass();
|
||||
write_task_results(&path, &[scored]).expect("write");
|
||||
|
||||
// Verify scored result replaced pending
|
||||
let after = read_task_results(&path).expect("read");
|
||||
assert_eq!(after.len(), 1);
|
||||
assert_eq!(after[0].score.label, "pass");
|
||||
assert_eq!(after[0].score.value, 1.0);
|
||||
}
|
||||
}
|
||||
@@ -1,550 +0,0 @@
|
||||
use std::collections::{HashMap, 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, write_task_results,
|
||||
};
|
||||
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: Arc<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: Arc::from(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 all tasks once (used for both execution and scoring)
|
||||
let all_tasks = self.suite.load_tasks().await?;
|
||||
let task_index: HashMap<String, BenchTask> = all_tasks
|
||||
.iter()
|
||||
.map(|t| (t.id.clone(), t.clone()))
|
||||
.collect();
|
||||
|
||||
// Filter tasks for execution
|
||||
let mut tasks = all_tasks;
|
||||
|
||||
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();
|
||||
let model_label = matrix.model.as_deref().unwrap_or(self.llm.model_name());
|
||||
let commit_hash = git_short_hash();
|
||||
tracing::info!(
|
||||
"[{} @ {}] Running {} tasks for suite '{}' (run: {})",
|
||||
model_label,
|
||||
commit_hash,
|
||||
tasks.len(),
|
||||
self.suite.id(),
|
||||
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
|
||||
);
|
||||
if let Err(e) = self.suite.setup_task(task).await {
|
||||
tracing::warn!("setup_task failed for {}: {}", task.id, e);
|
||||
let result = make_error_result(
|
||||
task,
|
||||
self.suite.id(),
|
||||
&matrix.label,
|
||||
Utc::now(),
|
||||
&format!("setup_task failed: {e}"),
|
||||
);
|
||||
append_task_result(&jsonl_path, &result)?;
|
||||
all_results.lock().await.push(result);
|
||||
continue;
|
||||
}
|
||||
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;
|
||||
if let Err(e) = self.suite.teardown_task(task).await {
|
||||
tracing::warn!("teardown_task failed for {}: {}", task.id, e);
|
||||
}
|
||||
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 shared_tools: Arc<[Arc<dyn ironclaw::tools::Tool>]> =
|
||||
Arc::from(self.suite.additional_tools());
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for (i, task) in tasks.into_iter().enumerate() {
|
||||
let sem = Arc::clone(&semaphore);
|
||||
let suite = Arc::clone(&self.suite);
|
||||
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 completed_count = completed.len();
|
||||
let total = total_tasks;
|
||||
let additional_tools = Arc::clone(&shared_tools);
|
||||
|
||||
handles.push(tokio::spawn(async move {
|
||||
let _permit = match sem.acquire().await {
|
||||
Ok(p) => p,
|
||||
Err(_) => {
|
||||
tracing::error!("Semaphore closed for task {}", task.id);
|
||||
return;
|
||||
}
|
||||
};
|
||||
tracing::info!(
|
||||
"[{}/{}] Running task: {}",
|
||||
i + 1 + completed_count,
|
||||
total,
|
||||
task.id
|
||||
);
|
||||
if let Err(e) = suite.setup_task(&task).await {
|
||||
tracing::warn!("setup_task failed for {}: {}", task.id, e);
|
||||
let result = make_error_result(
|
||||
&task,
|
||||
suite.id(),
|
||||
&config_label,
|
||||
Utc::now(),
|
||||
&format!("setup_task failed: {e}"),
|
||||
);
|
||||
results_ref.lock().await.push(result);
|
||||
return;
|
||||
}
|
||||
let suite_id = suite.id().to_string();
|
||||
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) = suite.teardown_task(&task).await {
|
||||
tracing::warn!("teardown_task failed 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);
|
||||
}
|
||||
}
|
||||
|
||||
// Write all results to JSONL after parallel execution completes.
|
||||
// This avoids the race condition of concurrent file appends.
|
||||
let results = all_results.lock().await;
|
||||
for result in results.iter() {
|
||||
append_task_result(&jsonl_path, result)?;
|
||||
}
|
||||
}
|
||||
|
||||
// Score all results using the cached task index
|
||||
let results = all_results.lock().await;
|
||||
let mut scored: Vec<TaskResult> = Vec::with_capacity(results.len());
|
||||
for result in results.iter() {
|
||||
if let Some(task) = task_index.get(&result.task_id) {
|
||||
let submission = TaskSubmission {
|
||||
response: result.response.clone(),
|
||||
conversation: vec![],
|
||||
tool_calls: result
|
||||
.trace
|
||||
.tool_calls
|
||||
.iter()
|
||||
.map(|tc| tc.name.clone())
|
||||
.collect(),
|
||||
error: result.error.clone(),
|
||||
};
|
||||
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);
|
||||
|
||||
// Rewrite JSONL with scored results so `results` command shows final scores
|
||||
write_task_results(&jsonl_path, &all_for_aggregate)?;
|
||||
|
||||
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,
|
||||
&commit_hash,
|
||||
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",
|
||||
model_name,
|
||||
commit_hash,
|
||||
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.
|
||||
///
|
||||
/// # Current limitations
|
||||
///
|
||||
/// - **Single-turn only**: After the first assistant response, `/quit` is sent.
|
||||
/// Multi-turn suites (e.g., Tau-bench's `next_user_message()`) are not yet wired.
|
||||
/// - **Resources not injected**: `BenchTask.resources` (e.g., GAIA file attachments)
|
||||
/// are not included in the prompt or made available via the workspace.
|
||||
/// - **Conversation not captured**: `TaskSubmission.conversation` is always empty,
|
||||
/// which prevents multi-turn scoring hooks from working.
|
||||
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,
|
||||
allow_local_tools: true,
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
};
|
||||
|
||||
let cost_guard = Arc::new(ironclaw::agent::cost_guard::CostGuard::new(
|
||||
ironclaw::agent::cost_guard::CostGuardConfig::default(),
|
||||
));
|
||||
|
||||
let deps = AgentDeps {
|
||||
store: None,
|
||||
llm: instrumented.clone() as Arc<dyn LlmProvider>,
|
||||
cheap_llm: None,
|
||||
safety,
|
||||
tools,
|
||||
workspace: None,
|
||||
extension_manager: None,
|
||||
skill_registry: None,
|
||||
skills_config: ironclaw::config::SkillsConfig::default(),
|
||||
hooks: Arc::new(ironclaw::hooks::HookRegistry::new()),
|
||||
cost_guard,
|
||||
};
|
||||
|
||||
let mut channels = ChannelManager::new();
|
||||
channels.add(Box::new(bench_channel));
|
||||
|
||||
let agent = Agent::new(agent_config, deps, channels, None, None, 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()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the short git commit hash of HEAD, or "unknown" if not in a repo.
|
||||
fn git_short_hash() -> String {
|
||||
std::process::Command::new("git")
|
||||
.args(["rev-parse", "--short", "HEAD"])
|
||||
.output()
|
||||
.ok()
|
||||
.and_then(|o| {
|
||||
if o.status.success() {
|
||||
Some(String::from_utf8_lossy(&o.stdout).trim().to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.unwrap_or_else(|| "unknown".to_string())
|
||||
}
|
||||
@@ -1,113 +0,0 @@
|
||||
use regex::Regex;
|
||||
|
||||
use crate::suite::BenchScore;
|
||||
|
||||
/// Normalize an answer string for comparison: lowercase, trim whitespace,
|
||||
/// strip trailing punctuation, collapse internal whitespace.
|
||||
pub fn normalize_answer(s: &str) -> String {
|
||||
let trimmed = s.trim().to_lowercase();
|
||||
let collapsed: String = trimmed.split_whitespace().collect::<Vec<_>>().join(" ");
|
||||
collapsed.trim_end_matches(['.', ',', ';', '!']).to_string()
|
||||
}
|
||||
|
||||
/// Exact match after normalization.
|
||||
pub fn exact_match(expected: &str, actual: &str) -> BenchScore {
|
||||
let norm_expected = normalize_answer(expected);
|
||||
let norm_actual = normalize_answer(actual);
|
||||
if norm_expected == norm_actual {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!(
|
||||
"expected \"{norm_expected}\", got \"{norm_actual}\""
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the actual answer contains the expected substring (normalized).
|
||||
pub fn contains_match(expected_substring: &str, actual: &str) -> BenchScore {
|
||||
let norm_expected = normalize_answer(expected_substring);
|
||||
let norm_actual = normalize_answer(actual);
|
||||
if norm_actual.contains(&norm_expected) {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!("response does not contain \"{norm_expected}\""))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if the actual answer matches a regex pattern.
|
||||
pub fn regex_match(pattern: &str, actual: &str) -> BenchScore {
|
||||
match Regex::new(pattern) {
|
||||
Ok(re) => {
|
||||
if re.is_match(actual) {
|
||||
BenchScore::pass()
|
||||
} else {
|
||||
BenchScore::fail(format!("response does not match pattern /{pattern}/"))
|
||||
}
|
||||
}
|
||||
Err(e) => BenchScore::fail(format!("invalid regex pattern: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_normalize_answer() {
|
||||
assert_eq!(normalize_answer(" Hello World. "), "hello world");
|
||||
assert_eq!(normalize_answer("Yes!"), "yes");
|
||||
assert_eq!(normalize_answer("42"), "42");
|
||||
assert_eq!(normalize_answer(" "), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exact_match_pass() {
|
||||
let score = exact_match("Hello World", " hello world. ");
|
||||
assert_eq!(score.value, 1.0);
|
||||
assert_eq!(score.label, "pass");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_exact_match_fail() {
|
||||
let score = exact_match("hello", "world");
|
||||
assert_eq!(score.value, 0.0);
|
||||
assert_eq!(score.label, "fail");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_contains_match_pass() {
|
||||
let score = contains_match("world", "Hello World!");
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_contains_match_fail() {
|
||||
let score = contains_match("xyz", "Hello World!");
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_pass() {
|
||||
let score = regex_match(r"\d{4}", "The year is 2024.");
|
||||
assert_eq!(score.value, 1.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_fail() {
|
||||
let score = regex_match(r"\d{4}", "No numbers here.");
|
||||
assert_eq!(score.value, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_regex_match_invalid_pattern() {
|
||||
let score = regex_match(r"[invalid", "anything");
|
||||
assert_eq!(score.value, 0.0);
|
||||
assert!(
|
||||
score
|
||||
.details
|
||||
.as_deref()
|
||||
.unwrap_or("")
|
||||
.contains("invalid regex")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,154 +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)]
|
||||
#[allow(dead_code)]
|
||||
pub struct TaskSubmission {
|
||||
pub response: String,
|
||||
pub conversation: Vec<ConversationTurn>,
|
||||
pub tool_calls: Vec<String>,
|
||||
pub error: Option<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, PartialEq, 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]
|
||||
#[allow(dead_code)]
|
||||
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![]
|
||||
}
|
||||
|
||||
/// 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)
|
||||
}
|
||||
}
|
||||
@@ -358,10 +358,6 @@ impl Agent {
|
||||
}
|
||||
});
|
||||
|
||||
tracing::info!(
|
||||
"Heartbeat enabled with {}s interval",
|
||||
hb_config.interval_secs
|
||||
);
|
||||
let hygiene = self
|
||||
.hygiene_config
|
||||
.as_ref()
|
||||
@@ -373,6 +369,7 @@ impl Agent {
|
||||
hygiene,
|
||||
workspace.clone(),
|
||||
self.cheap_llm().clone(),
|
||||
self.safety().clone(),
|
||||
Some(notify_tx),
|
||||
))
|
||||
} else {
|
||||
|
||||
+10
-7
@@ -13,7 +13,7 @@ use crate::agent::submission::SubmissionResult;
|
||||
use crate::agent::{Agent, MessageIntent};
|
||||
use crate::channels::{IncomingMessage, StatusUpdate};
|
||||
use crate::error::Error;
|
||||
use crate::llm::ChatMessage;
|
||||
use crate::llm::{ChatMessage, Reasoning};
|
||||
|
||||
impl Agent {
|
||||
/// Handle job-related intents without turn tracking.
|
||||
@@ -235,6 +235,7 @@ impl Agent {
|
||||
crate::workspace::hygiene::HygieneConfig::default(),
|
||||
workspace.clone(),
|
||||
self.llm().clone(),
|
||||
self.safety().clone(),
|
||||
);
|
||||
|
||||
match runner.check_heartbeat().await {
|
||||
@@ -295,10 +296,11 @@ impl Agent {
|
||||
.with_max_tokens(512)
|
||||
.with_temperature(0.3);
|
||||
|
||||
match self.llm().complete(request).await {
|
||||
Ok(response) => Ok(SubmissionResult::response(format!(
|
||||
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
|
||||
match reasoning.complete(request).await {
|
||||
Ok((text, _usage)) => Ok(SubmissionResult::response(format!(
|
||||
"Thread Summary:\n\n{}",
|
||||
response.content.trim()
|
||||
text.trim()
|
||||
))),
|
||||
Err(e) => Ok(SubmissionResult::error(format!("Summarize failed: {}", e))),
|
||||
}
|
||||
@@ -342,10 +344,11 @@ impl Agent {
|
||||
.with_max_tokens(512)
|
||||
.with_temperature(0.5);
|
||||
|
||||
match self.llm().complete(request).await {
|
||||
Ok(response) => Ok(SubmissionResult::response(format!(
|
||||
let reasoning = Reasoning::new(self.llm().clone(), self.safety().clone());
|
||||
match reasoning.complete(request).await {
|
||||
Ok((text, _usage)) => Ok(SubmissionResult::response(format!(
|
||||
"Suggested Next Steps:\n\n{}",
|
||||
response.content.trim()
|
||||
text.trim()
|
||||
))),
|
||||
Err(e) => Ok(SubmissionResult::error(format!("Suggest failed: {}", e))),
|
||||
}
|
||||
|
||||
@@ -12,7 +12,8 @@ use chrono::Utc;
|
||||
use crate::agent::context_monitor::{CompactionStrategy, ContextBreakdown};
|
||||
use crate::agent::session::Thread;
|
||||
use crate::error::Error;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider};
|
||||
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
/// Result of a compaction operation.
|
||||
@@ -33,12 +34,13 @@ pub struct CompactionResult {
|
||||
/// Compacts conversation context to stay within limits.
|
||||
pub struct ContextCompactor {
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
}
|
||||
|
||||
impl ContextCompactor {
|
||||
/// Create a new context compactor.
|
||||
pub fn new(llm: Arc<dyn LlmProvider>) -> Self {
|
||||
Self { llm }
|
||||
pub fn new(llm: Arc<dyn LlmProvider>, safety: Arc<SafetyLayer>) -> Self {
|
||||
Self { llm, safety }
|
||||
}
|
||||
|
||||
/// Compact a thread's context using the given strategy.
|
||||
@@ -231,8 +233,9 @@ Be brief but capture all important details. Use bullet points."#,
|
||||
.with_max_tokens(1024)
|
||||
.with_temperature(0.3);
|
||||
|
||||
let response = self.llm.complete(request).await?;
|
||||
Ok(response.content)
|
||||
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
|
||||
let (text, _) = reasoning.complete(request).await?;
|
||||
Ok(text)
|
||||
}
|
||||
|
||||
/// Write a summary to the workspace daily log.
|
||||
|
||||
+39
-33
@@ -33,16 +33,12 @@ impl Agent {
|
||||
/// Returns `AgenticLoopResult::Response` on completion, or
|
||||
/// `AgenticLoopResult::NeedApproval` if a tool requires user approval.
|
||||
///
|
||||
/// When `resume_after_tool` is true the loop already knows a tool was
|
||||
/// executed earlier in this turn (e.g. an approved tool), so it won't
|
||||
/// force the LLM to use tools if it responds with text.
|
||||
pub(super) async fn run_agentic_loop(
|
||||
&self,
|
||||
message: &IncomingMessage,
|
||||
session: Arc<Mutex<Session>>,
|
||||
thread_id: Uuid,
|
||||
initial_messages: Vec<ChatMessage>,
|
||||
resume_after_tool: bool,
|
||||
) -> Result<AgenticLoopResult, Error> {
|
||||
// Load workspace system prompt (identity files: AGENTS.md, SOUL.md, etc.)
|
||||
let system_prompt = if let Some(ws) = self.workspace() {
|
||||
@@ -112,16 +108,21 @@ impl Agent {
|
||||
// Create a JobContext for tool execution (chat doesn't have a real job)
|
||||
let job_ctx = JobContext::with_user(&message.user_id, "chat", "Interactive chat session");
|
||||
|
||||
const MAX_TOOL_ITERATIONS: usize = 10;
|
||||
let max_tool_iterations = self.config.max_tool_iterations;
|
||||
// Force a text-only response on the last iteration to guarantee termination
|
||||
// instead of hard-erroring. The penultimate iteration also gets a nudge
|
||||
// message so the LLM knows it should wrap up.
|
||||
let force_text_at = max_tool_iterations;
|
||||
let nudge_at = max_tool_iterations.saturating_sub(1);
|
||||
let mut iteration = 0;
|
||||
let mut tools_executed = resume_after_tool;
|
||||
|
||||
loop {
|
||||
iteration += 1;
|
||||
if iteration > MAX_TOOL_ITERATIONS {
|
||||
// Hard ceiling one past the forced-text iteration (should never be reached
|
||||
// since force_text_at guarantees a text response, but kept as a safety net).
|
||||
if iteration > max_tool_iterations + 1 {
|
||||
return Err(crate::error::LlmError::InvalidResponse {
|
||||
provider: "agent".to_string(),
|
||||
reason: format!("Exceeded maximum tool iterations ({})", MAX_TOOL_ITERATIONS),
|
||||
reason: format!("Exceeded maximum tool iterations ({max_tool_iterations})"),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
@@ -149,6 +150,19 @@ impl Agent {
|
||||
.into());
|
||||
}
|
||||
|
||||
// Inject a nudge message when approaching the iteration limit so the
|
||||
// LLM is aware it should produce a final answer on the next turn.
|
||||
if iteration == nudge_at {
|
||||
context_messages.push(ChatMessage::system(
|
||||
"You are approaching the tool call limit. \
|
||||
Provide your best final answer on the next response \
|
||||
using the information you have gathered so far. \
|
||||
Do not call any more tools.",
|
||||
));
|
||||
}
|
||||
|
||||
let force_text = iteration >= force_text_at;
|
||||
|
||||
// Refresh tool definitions each iteration so newly built tools become visible
|
||||
let tool_defs = self.tools().tool_definitions().await;
|
||||
|
||||
@@ -168,8 +182,9 @@ impl Agent {
|
||||
tool_defs
|
||||
};
|
||||
|
||||
// Call LLM with current context
|
||||
let context = ReasoningContext::new()
|
||||
// Call LLM with current context; force_text drops tools to guarantee a
|
||||
// text response on the final iteration.
|
||||
let mut context = ReasoningContext::new()
|
||||
.with_messages(context_messages.clone())
|
||||
.with_tools(tool_defs)
|
||||
.with_metadata({
|
||||
@@ -177,6 +192,14 @@ impl Agent {
|
||||
m.insert("thread_id".to_string(), thread_id.to_string());
|
||||
m
|
||||
});
|
||||
context.force_text = force_text;
|
||||
|
||||
if force_text {
|
||||
tracing::info!(
|
||||
iteration,
|
||||
"Forcing text-only response (iteration limit reached)"
|
||||
);
|
||||
}
|
||||
|
||||
let output = reasoning.respond_with_tools(&context).await?;
|
||||
|
||||
@@ -199,30 +222,12 @@ impl Agent {
|
||||
|
||||
match output.result {
|
||||
RespondResult::Text(text) => {
|
||||
// If no tools have been executed yet, prompt the LLM to use tools
|
||||
// This handles the case where the model explains what it will do
|
||||
// instead of actually calling tools
|
||||
if !tools_executed && iteration < 3 {
|
||||
tracing::debug!(
|
||||
"No tools executed yet (iteration {}), prompting for tool use",
|
||||
iteration
|
||||
);
|
||||
context_messages.push(ChatMessage::assistant(&text));
|
||||
context_messages.push(ChatMessage::user(
|
||||
"Please proceed and use the available tools to complete this task.",
|
||||
));
|
||||
continue;
|
||||
}
|
||||
|
||||
// Tools have been executed or we've tried multiple times, return response
|
||||
return Ok(AgenticLoopResult::Response(text));
|
||||
}
|
||||
RespondResult::ToolCalls {
|
||||
tool_calls,
|
||||
content,
|
||||
} => {
|
||||
tools_executed = true;
|
||||
|
||||
// Add the assistant message with tool_calls to context.
|
||||
// OpenAI protocol requires this before tool-result messages.
|
||||
context_messages.push(ChatMessage::assistant_with_tool_calls(
|
||||
@@ -279,8 +284,9 @@ impl Agent {
|
||||
for (idx, original_tc) in tool_calls.iter().enumerate() {
|
||||
let mut tc = original_tc.clone();
|
||||
|
||||
// Check if tool requires approval
|
||||
if let Some(tool) = self.tools().get(&tc.name).await
|
||||
// Check if tool requires approval (skipped when auto_approve_tools is set)
|
||||
if !self.config.auto_approve_tools
|
||||
&& let Some(tool) = self.tools().get(&tc.name).await
|
||||
&& tool.requires_approval()
|
||||
{
|
||||
let mut is_auto_approved = {
|
||||
@@ -823,7 +829,6 @@ mod tests {
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -837,7 +842,6 @@ mod tests {
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -874,6 +878,8 @@ mod tests {
|
||||
allow_local_tools: false,
|
||||
max_cost_per_day_cents: None,
|
||||
max_actions_per_hour: None,
|
||||
max_tool_iterations: 50,
|
||||
auto_approve_tools: false,
|
||||
},
|
||||
deps,
|
||||
ChannelManager::new(),
|
||||
|
||||
+11
-13
@@ -29,7 +29,8 @@ use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::channels::OutgoingResponse;
|
||||
use crate::llm::{ChatMessage, CompletionRequest, FinishReason, LlmProvider};
|
||||
use crate::llm::{ChatMessage, CompletionRequest, LlmProvider, Reasoning};
|
||||
use crate::safety::SafetyLayer;
|
||||
use crate::workspace::Workspace;
|
||||
use crate::workspace::hygiene::HygieneConfig;
|
||||
|
||||
@@ -100,6 +101,7 @@ pub struct HeartbeatRunner {
|
||||
hygiene_config: HygieneConfig,
|
||||
workspace: Arc<Workspace>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||
consecutive_failures: u32,
|
||||
}
|
||||
@@ -111,12 +113,14 @@ impl HeartbeatRunner {
|
||||
hygiene_config: HygieneConfig,
|
||||
workspace: Arc<Workspace>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
hygiene_config,
|
||||
workspace,
|
||||
llm,
|
||||
safety,
|
||||
response_tx: None,
|
||||
consecutive_failures: 0,
|
||||
}
|
||||
@@ -258,25 +262,18 @@ impl HeartbeatRunner {
|
||||
.with_max_tokens(max_tokens)
|
||||
.with_temperature(0.3);
|
||||
|
||||
let response = match self.llm.complete(request).await {
|
||||
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
|
||||
let (content, _usage) = match reasoning.complete(request).await {
|
||||
Ok(r) => r,
|
||||
Err(e) => return HeartbeatResult::Failed(format!("LLM call failed: {}", e)),
|
||||
};
|
||||
|
||||
let content = response.content.trim();
|
||||
let content = content.trim();
|
||||
|
||||
// Guard against empty content. Reasoning models (e.g. GLM-4.7) may
|
||||
// burn all output tokens on chain-of-thought and return content: null.
|
||||
if content.is_empty() {
|
||||
return if response.finish_reason == FinishReason::Length {
|
||||
HeartbeatResult::Failed(
|
||||
"LLM response was truncated (finish_reason=length) with no content. \
|
||||
The model may have exhausted its token budget on reasoning."
|
||||
.to_string(),
|
||||
)
|
||||
} else {
|
||||
HeartbeatResult::Failed("LLM returned empty content.".to_string())
|
||||
};
|
||||
return HeartbeatResult::Failed("LLM returned empty content.".to_string());
|
||||
}
|
||||
|
||||
// Check if nothing needs attention
|
||||
@@ -355,9 +352,10 @@ pub fn spawn_heartbeat(
|
||||
hygiene_config: HygieneConfig,
|
||||
workspace: Arc<Workspace>,
|
||||
llm: Arc<dyn LlmProvider>,
|
||||
safety: Arc<SafetyLayer>,
|
||||
response_tx: Option<mpsc::Sender<OutgoingResponse>>,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm);
|
||||
let mut runner = HeartbeatRunner::new(config, hygiene_config, workspace, llm, safety);
|
||||
if let Some(tx) = response_tx {
|
||||
runner = runner.with_response_channel(tx);
|
||||
}
|
||||
|
||||
@@ -185,10 +185,6 @@ pub struct Thread {
|
||||
/// Pending auth token request (thread is in auth mode).
|
||||
#[serde(default)]
|
||||
pub pending_auth: Option<PendingAuth>,
|
||||
/// Last NEAR AI response ID for response chaining. Persisted to DB
|
||||
/// metadata so we can resume chaining across restarts.
|
||||
#[serde(default)]
|
||||
pub last_response_id: Option<String>,
|
||||
}
|
||||
|
||||
impl Thread {
|
||||
@@ -205,7 +201,6 @@ impl Thread {
|
||||
metadata: serde_json::Value::Null,
|
||||
pending_approval: None,
|
||||
pending_auth: None,
|
||||
last_response_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -222,7 +217,6 @@ impl Thread {
|
||||
metadata: serde_json::Value::Null,
|
||||
pending_approval: None,
|
||||
pending_auth: None,
|
||||
last_response_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -863,7 +857,6 @@ mod tests {
|
||||
|
||||
thread.start_turn("hello");
|
||||
thread.complete_turn("world");
|
||||
thread.last_response_id = Some("resp_abc123".to_string());
|
||||
|
||||
let json = serde_json::to_string(&thread).unwrap();
|
||||
let restored: Thread = serde_json::from_str(&json).unwrap();
|
||||
@@ -873,7 +866,6 @@ mod tests {
|
||||
assert_eq!(restored.turns.len(), 1);
|
||||
assert_eq!(restored.turns[0].user_input, "hello");
|
||||
assert_eq!(restored.turns[0].response, Some("world".to_string()));
|
||||
assert_eq!(restored.last_response_id, Some("resp_abc123".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
+41
-92
@@ -87,20 +87,6 @@ impl Agent {
|
||||
thread.restore_from_messages(chat_messages);
|
||||
}
|
||||
|
||||
// Restore response chain from conversation metadata
|
||||
if let Some(store) = self.store()
|
||||
&& let Ok(Some(metadata)) = store.get_conversation_metadata(thread_uuid).await
|
||||
&& let Some(rid) = metadata
|
||||
.get("last_response_id")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
{
|
||||
thread.last_response_id = Some(rid.clone());
|
||||
self.llm()
|
||||
.seed_response_chain(&thread_uuid.to_string(), rid);
|
||||
tracing::debug!("Restored response chain for thread {}", thread_uuid);
|
||||
}
|
||||
|
||||
// Insert into session and register with session manager
|
||||
{
|
||||
let mut sess = session.lock().await;
|
||||
@@ -228,7 +214,7 @@ impl Agent {
|
||||
)
|
||||
.await;
|
||||
|
||||
let compactor = ContextCompactor::new(self.llm().clone());
|
||||
let compactor = ContextCompactor::new(self.llm().clone(), self.safety().clone());
|
||||
if let Err(e) = compactor
|
||||
.compact(thread, strategy, self.workspace().map(|w| w.as_ref()))
|
||||
.await
|
||||
@@ -278,7 +264,7 @@ impl Agent {
|
||||
|
||||
// Run the agentic tool execution loop
|
||||
let result = self
|
||||
.run_agentic_loop(message, session.clone(), thread_id, turn_messages, false)
|
||||
.run_agentic_loop(message, session.clone(), thread_id, turn_messages)
|
||||
.await;
|
||||
|
||||
// Re-acquire lock and check if interrupted
|
||||
@@ -325,7 +311,6 @@ impl Agent {
|
||||
};
|
||||
|
||||
thread.complete_turn(&response);
|
||||
self.persist_response_chain(thread);
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
@@ -335,8 +320,10 @@ impl Agent {
|
||||
)
|
||||
.await;
|
||||
|
||||
// Fire-and-forget: persist turn to DB
|
||||
self.persist_turn(thread_id, &message.user_id, content, Some(&response));
|
||||
// Persist turn to DB before returning so the write
|
||||
// completes even if the process shuts down right after.
|
||||
self.persist_turn(thread_id, &message.user_id, content, Some(&response))
|
||||
.await;
|
||||
|
||||
Ok(SubmissionResult::response(response))
|
||||
}
|
||||
@@ -366,15 +353,16 @@ impl Agent {
|
||||
thread.fail_turn(e.to_string());
|
||||
|
||||
// Persist the user message even on failure
|
||||
self.persist_turn(thread_id, &message.user_id, content, None);
|
||||
self.persist_turn(thread_id, &message.user_id, content, None)
|
||||
.await;
|
||||
|
||||
Ok(SubmissionResult::error(e.to_string()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fire-and-forget: persist a turn (user message + optional assistant response) to the DB.
|
||||
pub(super) fn persist_turn(
|
||||
/// Persist a turn (user message + optional assistant response) to the DB.
|
||||
pub(super) async fn persist_turn(
|
||||
&self,
|
||||
thread_id: Uuid,
|
||||
user_id: &str,
|
||||
@@ -386,70 +374,29 @@ impl Agent {
|
||||
None => return,
|
||||
};
|
||||
|
||||
let user_id = user_id.to_string();
|
||||
let user_input = user_input.to_string();
|
||||
let response = response.map(String::from);
|
||||
if let Err(e) = store
|
||||
.ensure_conversation(thread_id, "gateway", user_id, None)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e);
|
||||
return;
|
||||
}
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = store
|
||||
.ensure_conversation(thread_id, "gateway", &user_id, None)
|
||||
if let Err(e) = store
|
||||
.add_conversation_message(thread_id, "user", user_input)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to persist user message: {}", e);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(resp) = response
|
||||
&& let Err(e) = store
|
||||
.add_conversation_message(thread_id, "assistant", resp)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to ensure conversation {}: {}", thread_id, e);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Err(e) = store
|
||||
.add_conversation_message(thread_id, "user", &user_input)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to persist user message: {}", e);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(ref resp) = response
|
||||
&& let Err(e) = store
|
||||
.add_conversation_message(thread_id, "assistant", resp)
|
||||
.await
|
||||
{
|
||||
tracing::warn!("Failed to persist assistant message: {}", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Sync the provider's response chain ID to the thread and DB metadata.
|
||||
///
|
||||
/// Call after a successful agentic loop to persist the latest
|
||||
/// `previous_response_id` so chaining survives restarts.
|
||||
pub(super) fn persist_response_chain(&self, thread: &mut crate::agent::session::Thread) {
|
||||
let tid = thread.id.to_string();
|
||||
let response_id = match self.llm().get_response_chain_id(&tid) {
|
||||
Some(rid) => rid,
|
||||
None => return,
|
||||
};
|
||||
|
||||
// Update in-memory thread
|
||||
thread.last_response_id = Some(response_id.clone());
|
||||
|
||||
// Fire-and-forget DB write
|
||||
let store = match self.store() {
|
||||
Some(s) => Arc::clone(s),
|
||||
None => return,
|
||||
};
|
||||
let thread_id = thread.id;
|
||||
tokio::spawn(async move {
|
||||
let val = serde_json::json!(response_id);
|
||||
if let Err(e) = store
|
||||
.update_conversation_metadata_field(thread_id, "last_response_id", &val)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"Failed to persist response chain for thread {}: {}",
|
||||
thread_id,
|
||||
e
|
||||
);
|
||||
}
|
||||
});
|
||||
{
|
||||
tracing::warn!("Failed to persist assistant message: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn process_undo(
|
||||
@@ -562,7 +509,7 @@ impl Agent {
|
||||
crate::agent::context_monitor::CompactionStrategy::Summarize { keep_recent: 5 },
|
||||
);
|
||||
|
||||
let compactor = ContextCompactor::new(self.llm().clone());
|
||||
let compactor = ContextCompactor::new(self.llm().clone(), self.safety().clone());
|
||||
match compactor
|
||||
.compact(thread, strategy, self.workspace().map(|w| w.as_ref()))
|
||||
.await
|
||||
@@ -1057,7 +1004,7 @@ impl Agent {
|
||||
|
||||
// Continue the agentic loop (a tool was already executed this turn)
|
||||
let result = self
|
||||
.run_agentic_loop(message, session.clone(), thread_id, context_messages, true)
|
||||
.run_agentic_loop(message, session.clone(), thread_id, context_messages)
|
||||
.await;
|
||||
|
||||
// Handle the result
|
||||
@@ -1072,9 +1019,9 @@ impl Agent {
|
||||
let user_input = thread.last_turn().map(|t| t.user_input.clone());
|
||||
thread.complete_turn(&response);
|
||||
if let Some(input) = user_input {
|
||||
self.persist_turn(thread_id, &message.user_id, &input, Some(&response));
|
||||
self.persist_turn(thread_id, &message.user_id, &input, Some(&response))
|
||||
.await;
|
||||
}
|
||||
self.persist_response_chain(thread);
|
||||
let _ = self
|
||||
.channels
|
||||
.send_status(
|
||||
@@ -1112,7 +1059,8 @@ impl Agent {
|
||||
let user_input = thread.last_turn().map(|t| t.user_input.clone());
|
||||
thread.fail_turn(e.to_string());
|
||||
if let Some(input) = user_input {
|
||||
self.persist_turn(thread_id, &message.user_id, &input, None);
|
||||
self.persist_turn(thread_id, &message.user_id, &input, None)
|
||||
.await;
|
||||
}
|
||||
Ok(SubmissionResult::error(e.to_string()))
|
||||
}
|
||||
@@ -1131,7 +1079,8 @@ impl Agent {
|
||||
thread.clear_pending_approval();
|
||||
thread.complete_turn(&rejection);
|
||||
if let Some(input) = user_input {
|
||||
self.persist_turn(thread_id, &message.user_id, &input, Some(&rejection));
|
||||
self.persist_turn(thread_id, &message.user_id, &input, Some(&rejection))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1171,9 +1120,9 @@ impl Agent {
|
||||
thread.enter_auth_mode(ext_name.clone());
|
||||
thread.complete_turn(&instructions);
|
||||
if let Some(input) = user_input {
|
||||
self.persist_turn(thread_id, &message.user_id, &input, Some(&instructions));
|
||||
self.persist_turn(thread_id, &message.user_id, &input, Some(&instructions))
|
||||
.await;
|
||||
}
|
||||
self.persist_response_chain(thread);
|
||||
}
|
||||
}
|
||||
let _ = self
|
||||
|
||||
+5
-9
@@ -396,7 +396,6 @@ impl AppBuilder {
|
||||
|
||||
let tools = Arc::new(ToolRegistry::new());
|
||||
tools.register_builtin_tools();
|
||||
tracing::info!("Registered {} built-in tools", tools.count());
|
||||
|
||||
// Create embeddings provider if configured
|
||||
let embeddings: Option<Arc<dyn EmbeddingProvider>> = if self.config.embeddings.enabled {
|
||||
@@ -681,12 +680,12 @@ impl AppBuilder {
|
||||
None
|
||||
};
|
||||
|
||||
// Register dev tools if local tools are enabled
|
||||
if self.config.agent.allow_local_tools {
|
||||
// register_builder_tool() already calls register_dev_tools() internally,
|
||||
// so only register them here when the builder didn't already do it.
|
||||
let builder_registered_dev_tools = self.config.builder.enabled
|
||||
&& (self.config.agent.allow_local_tools || !self.config.sandbox.enabled);
|
||||
if self.config.agent.allow_local_tools && !builder_registered_dev_tools {
|
||||
tools.register_dev_tools();
|
||||
tracing::info!(
|
||||
"Local tools enabled (allow_local_tools=true), dev tools registered directly"
|
||||
);
|
||||
}
|
||||
|
||||
Ok((mcp_session_manager, wasm_tool_runtime, extension_manager))
|
||||
@@ -709,9 +708,6 @@ impl AppBuilder {
|
||||
// Seed workspace and backfill embeddings
|
||||
if let Some(ref ws) = workspace {
|
||||
match ws.seed_if_empty().await {
|
||||
Ok(count) if count > 0 => {
|
||||
tracing::info!("Workspace seeded with {} core files", count);
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to seed workspace: {}", e);
|
||||
|
||||
@@ -300,6 +300,51 @@ impl ChannelHostState {
|
||||
}
|
||||
}
|
||||
|
||||
/// In-memory workspace store for WASM channels.
|
||||
///
|
||||
/// Persists workspace writes across callback invocations within a single
|
||||
/// channel lifetime. This allows WASM channels to maintain state (e.g.,
|
||||
/// Telegram polling offsets) between poll ticks without requiring a
|
||||
/// full database-backed workspace.
|
||||
///
|
||||
/// Uses `std::sync::RwLock` (not tokio) because WASM execution runs
|
||||
/// inside `spawn_blocking`.
|
||||
pub struct ChannelWorkspaceStore {
|
||||
data: std::sync::RwLock<std::collections::HashMap<String, String>>,
|
||||
}
|
||||
|
||||
impl ChannelWorkspaceStore {
|
||||
/// Create a new empty workspace store.
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
data: std::sync::RwLock::new(std::collections::HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Commit pending writes from a callback execution into the store.
|
||||
pub fn commit_writes(&self, writes: &[PendingWorkspaceWrite]) {
|
||||
if writes.is_empty() {
|
||||
return;
|
||||
}
|
||||
if let Ok(mut data) = self.data.write() {
|
||||
for write in writes {
|
||||
tracing::debug!(
|
||||
path = %write.path,
|
||||
content_len = write.content.len(),
|
||||
"Committing workspace write to channel store"
|
||||
);
|
||||
data.insert(write.path.clone(), write.content.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::tools::wasm::WorkspaceReader for ChannelWorkspaceStore {
|
||||
fn read(&self, path: &str) -> Option<String> {
|
||||
self.data.read().ok()?.get(path).cloned()
|
||||
}
|
||||
}
|
||||
|
||||
/// Rate limiter for channel message emission.
|
||||
///
|
||||
/// Tracks emission rates across multiple executions.
|
||||
@@ -497,4 +542,56 @@ mod tests {
|
||||
|
||||
assert_eq!(state.channel_name(), "telegram");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_channel_workspace_store_commit_and_read() {
|
||||
use crate::channels::wasm::host::{ChannelWorkspaceStore, PendingWorkspaceWrite};
|
||||
use crate::tools::wasm::WorkspaceReader;
|
||||
|
||||
let store = ChannelWorkspaceStore::new();
|
||||
|
||||
// Initially empty
|
||||
assert!(store.read("channels/telegram/offset").is_none());
|
||||
|
||||
// Commit some writes
|
||||
let writes = vec![
|
||||
PendingWorkspaceWrite {
|
||||
path: "channels/telegram/offset".to_string(),
|
||||
content: "103".to_string(),
|
||||
},
|
||||
PendingWorkspaceWrite {
|
||||
path: "channels/telegram/state.json".to_string(),
|
||||
content: r#"{"ok":true}"#.to_string(),
|
||||
},
|
||||
];
|
||||
store.commit_writes(&writes);
|
||||
|
||||
// Should be readable
|
||||
assert_eq!(
|
||||
store.read("channels/telegram/offset"),
|
||||
Some("103".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
store.read("channels/telegram/state.json"),
|
||||
Some(r#"{"ok":true}"#.to_string())
|
||||
);
|
||||
|
||||
// Overwrite a value
|
||||
let writes2 = vec![PendingWorkspaceWrite {
|
||||
path: "channels/telegram/offset".to_string(),
|
||||
content: "200".to_string(),
|
||||
}];
|
||||
store.commit_writes(&writes2);
|
||||
assert_eq!(
|
||||
store.read("channels/telegram/offset"),
|
||||
Some("200".to_string())
|
||||
);
|
||||
|
||||
// Empty writes are a no-op
|
||||
store.commit_writes(&[]);
|
||||
assert_eq!(
|
||||
store.read("channels/telegram/offset"),
|
||||
Some("200".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -42,7 +42,9 @@ use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
|
||||
|
||||
use crate::channels::wasm::capabilities::ChannelCapabilities;
|
||||
use crate::channels::wasm::error::WasmChannelError;
|
||||
use crate::channels::wasm::host::{ChannelEmitRateLimiter, ChannelHostState, EmittedMessage};
|
||||
use crate::channels::wasm::host::{
|
||||
ChannelEmitRateLimiter, ChannelHostState, ChannelWorkspaceStore, EmittedMessage,
|
||||
};
|
||||
use crate::channels::wasm::router::RegisteredEndpoint;
|
||||
use crate::channels::wasm::runtime::{PreparedChannelModule, WasmChannelRuntime};
|
||||
use crate::channels::wasm::schema::ChannelConfig;
|
||||
@@ -547,6 +549,10 @@ pub struct WasmChannel {
|
||||
|
||||
/// Pairing store for DM pairing (guest access control).
|
||||
pairing_store: Arc<PairingStore>,
|
||||
|
||||
/// In-memory workspace store persisting writes across callback invocations.
|
||||
/// Ensures WASM channels can maintain state (e.g., polling offsets) between ticks.
|
||||
workspace_store: Arc<ChannelWorkspaceStore>,
|
||||
}
|
||||
|
||||
impl WasmChannel {
|
||||
@@ -577,6 +583,7 @@ impl WasmChannel {
|
||||
credentials: Arc::new(RwLock::new(HashMap::new())),
|
||||
typing_task: RwLock::new(None),
|
||||
pairing_store,
|
||||
workspace_store: Arc::new(ChannelWorkspaceStore::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -634,6 +641,26 @@ impl WasmChannel {
|
||||
self.endpoints.read().await.clone()
|
||||
}
|
||||
|
||||
/// Inject the workspace store as the reader into a capabilities clone.
|
||||
///
|
||||
/// Ensures `workspace_read` capability is present with the store as its reader,
|
||||
/// so WASM callbacks can read previously written workspace state.
|
||||
fn inject_workspace_reader(
|
||||
capabilities: &ChannelCapabilities,
|
||||
store: &Arc<ChannelWorkspaceStore>,
|
||||
) -> ChannelCapabilities {
|
||||
let mut caps = capabilities.clone();
|
||||
let ws_cap = caps
|
||||
.tool_capabilities
|
||||
.workspace_read
|
||||
.get_or_insert_with(|| crate::tools::wasm::WorkspaceCapability {
|
||||
allowed_prefixes: Vec::new(),
|
||||
reader: None,
|
||||
});
|
||||
ws_cap.reader = Some(Arc::clone(store) as Arc<dyn crate::tools::wasm::WorkspaceReader>);
|
||||
caps
|
||||
}
|
||||
|
||||
/// Add channel host functions to the linker using generated bindings.
|
||||
///
|
||||
/// Uses the wasmtime::component::bindgen! generated `add_to_linker` function
|
||||
@@ -765,12 +792,13 @@ impl WasmChannel {
|
||||
|
||||
let runtime = Arc::clone(&self.runtime);
|
||||
let prepared = Arc::clone(&self.prepared);
|
||||
let capabilities = self.capabilities.clone();
|
||||
let capabilities = Self::inject_workspace_reader(&self.capabilities, &self.workspace_store);
|
||||
let config_json = self.config_json.read().await.clone();
|
||||
let timeout = self.runtime.config().callback_timeout;
|
||||
let channel_name = self.name.clone();
|
||||
let credentials = self.get_credentials().await;
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
let workspace_store = self.workspace_store.clone();
|
||||
|
||||
// Execute in blocking task with timeout
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
@@ -801,8 +829,13 @@ impl WasmChannel {
|
||||
}
|
||||
};
|
||||
|
||||
let host_state =
|
||||
let mut host_state =
|
||||
Self::extract_host_state(&mut store, &prepared.name, &capabilities);
|
||||
|
||||
// Commit pending workspace writes to the persistent store
|
||||
let pending_writes = host_state.take_pending_writes();
|
||||
workspace_store.commit_writes(&pending_writes);
|
||||
|
||||
Ok((config, host_state))
|
||||
})
|
||||
.await
|
||||
@@ -897,10 +930,11 @@ impl WasmChannel {
|
||||
|
||||
let runtime = Arc::clone(&self.runtime);
|
||||
let prepared = Arc::clone(&self.prepared);
|
||||
let capabilities = self.capabilities.clone();
|
||||
let capabilities = Self::inject_workspace_reader(&self.capabilities, &self.workspace_store);
|
||||
let timeout = self.runtime.config().callback_timeout;
|
||||
let credentials = self.get_credentials().await;
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
let workspace_store = self.workspace_store.clone();
|
||||
|
||||
// Prepare request data
|
||||
let method = method.to_string();
|
||||
@@ -940,8 +974,13 @@ impl WasmChannel {
|
||||
.map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel))?;
|
||||
|
||||
let response = convert_http_response(wit_response);
|
||||
let host_state =
|
||||
let mut host_state =
|
||||
Self::extract_host_state(&mut store, &prepared.name, &capabilities);
|
||||
|
||||
// Commit pending workspace writes to the persistent store
|
||||
let pending_writes = host_state.take_pending_writes();
|
||||
workspace_store.commit_writes(&pending_writes);
|
||||
|
||||
Ok((response, host_state))
|
||||
})
|
||||
.await
|
||||
@@ -989,11 +1028,12 @@ impl WasmChannel {
|
||||
|
||||
let runtime = Arc::clone(&self.runtime);
|
||||
let prepared = Arc::clone(&self.prepared);
|
||||
let capabilities = self.capabilities.clone();
|
||||
let capabilities = Self::inject_workspace_reader(&self.capabilities, &self.workspace_store);
|
||||
let timeout = self.runtime.config().callback_timeout;
|
||||
let channel_name = self.name.clone();
|
||||
let credentials = self.get_credentials().await;
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
let workspace_store = self.workspace_store.clone();
|
||||
|
||||
// Execute in blocking task with timeout
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
@@ -1013,8 +1053,13 @@ impl WasmChannel {
|
||||
.call_on_poll(&mut store)
|
||||
.map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel))?;
|
||||
|
||||
let host_state =
|
||||
let mut host_state =
|
||||
Self::extract_host_state(&mut store, &prepared.name, &capabilities);
|
||||
|
||||
// Commit pending workspace writes to the persistent store
|
||||
let pending_writes = host_state.take_pending_writes();
|
||||
workspace_store.commit_writes(&pending_writes);
|
||||
|
||||
Ok(((), host_state))
|
||||
})
|
||||
.await
|
||||
@@ -1501,6 +1546,7 @@ impl WasmChannel {
|
||||
let credentials = self.credentials.clone();
|
||||
let pairing_store = self.pairing_store.clone();
|
||||
let callback_timeout = self.runtime.config().callback_timeout;
|
||||
let workspace_store = self.workspace_store.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
let mut interval_timer = tokio::time::interval(interval);
|
||||
@@ -1523,6 +1569,7 @@ impl WasmChannel {
|
||||
&credentials,
|
||||
pairing_store.clone(),
|
||||
callback_timeout,
|
||||
&workspace_store,
|
||||
).await;
|
||||
|
||||
match result {
|
||||
@@ -1565,7 +1612,10 @@ impl WasmChannel {
|
||||
|
||||
/// Execute a single poll callback with a fresh WASM instance.
|
||||
///
|
||||
/// Returns any emitted messages from the callback.
|
||||
/// Returns any emitted messages from the callback. Pending workspace writes
|
||||
/// are committed to the shared `ChannelWorkspaceStore` so state persists
|
||||
/// across poll ticks (e.g., Telegram polling offset).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn execute_poll(
|
||||
channel_name: &str,
|
||||
runtime: &Arc<WasmChannelRuntime>,
|
||||
@@ -1574,6 +1624,7 @@ impl WasmChannel {
|
||||
credentials: &RwLock<HashMap<String, String>>,
|
||||
pairing_store: Arc<PairingStore>,
|
||||
timeout: Duration,
|
||||
workspace_store: &Arc<ChannelWorkspaceStore>,
|
||||
) -> Result<Vec<EmittedMessage>, WasmChannelError> {
|
||||
// Skip if no WASM bytes (testing mode)
|
||||
if prepared.component_bytes.is_empty() {
|
||||
@@ -1586,9 +1637,10 @@ impl WasmChannel {
|
||||
|
||||
let runtime = Arc::clone(runtime);
|
||||
let prepared = Arc::clone(prepared);
|
||||
let capabilities = capabilities.clone();
|
||||
let capabilities = Self::inject_workspace_reader(capabilities, workspace_store);
|
||||
let credentials_snapshot = credentials.read().await.clone();
|
||||
let channel_name_owned = channel_name.to_string();
|
||||
let workspace_store = Arc::clone(workspace_store);
|
||||
|
||||
// Execute in blocking task with timeout
|
||||
let result = tokio::time::timeout(timeout, async move {
|
||||
@@ -1608,8 +1660,13 @@ impl WasmChannel {
|
||||
.call_on_poll(&mut store)
|
||||
.map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel))?;
|
||||
|
||||
let host_state =
|
||||
let mut host_state =
|
||||
Self::extract_host_state(&mut store, &prepared.name, &capabilities);
|
||||
|
||||
// Commit pending workspace writes to the persistent store
|
||||
let pending_writes = host_state.take_pending_writes();
|
||||
workspace_store.commit_writes(&pending_writes);
|
||||
|
||||
Ok(host_state)
|
||||
})
|
||||
.await
|
||||
@@ -2230,6 +2287,8 @@ mod tests {
|
||||
let credentials = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
|
||||
let timeout = std::time::Duration::from_secs(5);
|
||||
|
||||
let workspace_store = Arc::new(crate::channels::wasm::host::ChannelWorkspaceStore::new());
|
||||
|
||||
let result = WasmChannel::execute_poll(
|
||||
"poll-test",
|
||||
&runtime,
|
||||
@@ -2238,6 +2297,7 @@ mod tests {
|
||||
&credentials,
|
||||
Arc::new(PairingStore::new()),
|
||||
timeout,
|
||||
&workspace_store,
|
||||
)
|
||||
.await;
|
||||
|
||||
|
||||
@@ -22,7 +22,9 @@ use std::sync::{Arc, Mutex};
|
||||
use serde::Serialize;
|
||||
use tokio::sync::broadcast;
|
||||
use tracing::field::{Field, Visit};
|
||||
use tracing_subscriber::Layer;
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
use tracing_subscriber::util::SubscriberInitExt;
|
||||
use tracing_subscriber::{EnvFilter, Layer, reload};
|
||||
|
||||
use crate::safety::LeakDetector;
|
||||
|
||||
@@ -102,6 +104,115 @@ impl Default for LogBroadcaster {
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle for changing the tracing `EnvFilter` at runtime.
|
||||
///
|
||||
/// Wraps a `reload::Handle` so the gateway can switch between log levels
|
||||
/// (e.g. `ironclaw=debug`) without restarting the process.
|
||||
pub struct LogLevelHandle {
|
||||
handle: reload::Handle<EnvFilter, tracing_subscriber::Registry>,
|
||||
current_level: Mutex<String>,
|
||||
base_filter: String,
|
||||
}
|
||||
|
||||
impl LogLevelHandle {
|
||||
pub fn new(
|
||||
handle: reload::Handle<EnvFilter, tracing_subscriber::Registry>,
|
||||
initial_level: String,
|
||||
base_filter: String,
|
||||
) -> Self {
|
||||
Self {
|
||||
handle,
|
||||
current_level: Mutex::new(initial_level),
|
||||
base_filter,
|
||||
}
|
||||
}
|
||||
|
||||
/// Change the `ironclaw=<level>` directive at runtime.
|
||||
///
|
||||
/// `level` must be one of: trace, debug, info, warn, error.
|
||||
pub fn set_level(&self, level: &str) -> Result<(), String> {
|
||||
const VALID: &[&str] = &["trace", "debug", "info", "warn", "error"];
|
||||
let level = level.to_lowercase();
|
||||
if !VALID.contains(&level.as_str()) {
|
||||
return Err(format!(
|
||||
"invalid level '{}', must be one of: {}",
|
||||
level,
|
||||
VALID.join(", ")
|
||||
));
|
||||
}
|
||||
|
||||
let filter_str = if self.base_filter.is_empty() {
|
||||
format!("ironclaw={}", level)
|
||||
} else {
|
||||
format!("ironclaw={},{}", level, self.base_filter)
|
||||
};
|
||||
|
||||
let new_filter = EnvFilter::new(&filter_str);
|
||||
self.handle
|
||||
.reload(new_filter)
|
||||
.map_err(|e| format!("failed to reload filter: {}", e))?;
|
||||
|
||||
if let Ok(mut current) = self.current_level.lock() {
|
||||
*current = level;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the current ironclaw log level (e.g. "info", "debug").
|
||||
pub fn current_level(&self) -> String {
|
||||
self.current_level
|
||||
.lock()
|
||||
.map(|l| l.clone())
|
||||
.unwrap_or_else(|_| "info".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Initialise the tracing subscriber with a reloadable `EnvFilter`.
|
||||
///
|
||||
/// Returns the `LogLevelHandle` so callers can swap the filter at runtime.
|
||||
/// The fmt layer and `WebLogLayer` are attached alongside the reloadable filter.
|
||||
pub fn init_tracing(log_broadcaster: Arc<LogBroadcaster>) -> Arc<LogLevelHandle> {
|
||||
let raw_filter =
|
||||
std::env::var("RUST_LOG").unwrap_or_else(|_| "ironclaw=info,tower_http=warn".to_string());
|
||||
|
||||
// Split into the ironclaw directive and "everything else" (base_filter).
|
||||
let mut ironclaw_level = String::from("info");
|
||||
let mut base_parts: Vec<&str> = Vec::new();
|
||||
|
||||
for part in raw_filter.split(',') {
|
||||
let trimmed = part.trim();
|
||||
if trimmed.starts_with("ironclaw=") {
|
||||
if let Some(lvl) = trimmed.strip_prefix("ironclaw=") {
|
||||
ironclaw_level = lvl.to_string();
|
||||
}
|
||||
} else if !trimmed.is_empty() {
|
||||
base_parts.push(trimmed);
|
||||
}
|
||||
}
|
||||
let base_filter = base_parts.join(",");
|
||||
|
||||
let env_filter = EnvFilter::new(&raw_filter);
|
||||
let (reload_layer, reload_handle) = reload::Layer::new(env_filter);
|
||||
|
||||
let handle = Arc::new(LogLevelHandle::new(
|
||||
reload_handle,
|
||||
ironclaw_level,
|
||||
base_filter,
|
||||
));
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(reload_layer)
|
||||
.with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.with_target(false)
|
||||
.with_writer(crate::tracing_fmt::TruncatingStderr::default()),
|
||||
)
|
||||
.with(WebLogLayer::new(log_broadcaster))
|
||||
.init();
|
||||
|
||||
handle
|
||||
}
|
||||
|
||||
/// Visitor that extracts the `message` field and all extra key-value
|
||||
/// fields from a tracing event.
|
||||
///
|
||||
|
||||
@@ -41,7 +41,7 @@ use crate::skills::registry::SkillRegistry;
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::workspace::Workspace;
|
||||
|
||||
use self::log_layer::LogBroadcaster;
|
||||
use self::log_layer::{LogBroadcaster, LogLevelHandle};
|
||||
|
||||
use self::server::GatewayState;
|
||||
use self::sse::SseManager;
|
||||
@@ -76,6 +76,7 @@ impl GatewayChannel {
|
||||
workspace: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
@@ -105,6 +106,7 @@ impl GatewayChannel {
|
||||
workspace: self.state.workspace.clone(),
|
||||
session_manager: self.state.session_manager.clone(),
|
||||
log_broadcaster: self.state.log_broadcaster.clone(),
|
||||
log_level_handle: self.state.log_level_handle.clone(),
|
||||
extension_manager: self.state.extension_manager.clone(),
|
||||
tool_registry: self.state.tool_registry.clone(),
|
||||
store: self.state.store.clone(),
|
||||
@@ -140,6 +142,12 @@ impl GatewayChannel {
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the log level handle for runtime log level control.
|
||||
pub fn with_log_level_handle(mut self, h: Arc<LogLevelHandle>) -> Self {
|
||||
self.rebuild_state(|s| s.log_level_handle = Some(h));
|
||||
self
|
||||
}
|
||||
|
||||
/// Inject the extension manager for the extensions API.
|
||||
pub fn with_extension_manager(mut self, em: Arc<ExtensionManager>) -> Self {
|
||||
self.rebuild_state(|s| s.extension_manager = Some(em));
|
||||
|
||||
+54
-11
@@ -122,6 +122,8 @@ pub struct GatewayState {
|
||||
pub session_manager: Option<Arc<SessionManager>>,
|
||||
/// Log broadcaster for the logs SSE endpoint.
|
||||
pub log_broadcaster: Option<Arc<LogBroadcaster>>,
|
||||
/// Handle for changing the tracing log level at runtime.
|
||||
pub log_level_handle: Option<Arc<crate::channels::web::log_layer::LogLevelHandle>>,
|
||||
/// Extension manager for extension management API.
|
||||
pub extension_manager: Option<Arc<ExtensionManager>>,
|
||||
/// Tool registry for listing registered tools.
|
||||
@@ -204,6 +206,11 @@ pub async fn start_server(
|
||||
.route("/api/jobs/{id}/files/read", get(job_files_read_handler))
|
||||
// Logs
|
||||
.route("/api/logs/events", get(logs_events_handler))
|
||||
.route("/api/logs/level", get(logs_level_get_handler))
|
||||
.route(
|
||||
"/api/logs/level",
|
||||
axum::routing::put(logs_level_set_handler),
|
||||
)
|
||||
// Extensions
|
||||
.route("/api/extensions", get(extensions_list_handler))
|
||||
.route("/api/extensions/tools", get(extensions_tools_handler))
|
||||
@@ -557,9 +564,13 @@ pub async fn clear_auth_mode(state: &GatewayState) {
|
||||
async fn chat_events_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||
state.sse.subscribe().ok_or((
|
||||
let sse = state.sse.subscribe().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Too many connections".to_string(),
|
||||
))?;
|
||||
Ok((
|
||||
[("X-Accel-Buffering", "no"), ("Cache-Control", "no-cache")],
|
||||
sse,
|
||||
))
|
||||
}
|
||||
|
||||
@@ -1585,10 +1596,7 @@ async fn job_files_read_handler(
|
||||
|
||||
async fn logs_events_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Result<
|
||||
Sse<impl futures::Stream<Item = Result<Event, Infallible>> + Send + 'static>,
|
||||
(StatusCode, String),
|
||||
> {
|
||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||
let broadcaster = state.log_broadcaster.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Log broadcaster not available".to_string(),
|
||||
@@ -1601,25 +1609,60 @@ async fn logs_events_handler(
|
||||
|
||||
let history_stream = futures::stream::iter(history).map(|entry| {
|
||||
let data = serde_json::to_string(&entry).unwrap_or_default();
|
||||
Ok(Event::default().event("log").data(data))
|
||||
Ok::<_, Infallible>(Event::default().event("log").data(data))
|
||||
});
|
||||
|
||||
let live_stream = tokio_stream::wrappers::BroadcastStream::new(rx)
|
||||
.filter_map(|result| result.ok())
|
||||
.map(|entry| {
|
||||
let data = serde_json::to_string(&entry).unwrap_or_default();
|
||||
Ok(Event::default().event("log").data(data))
|
||||
Ok::<_, Infallible>(Event::default().event("log").data(data))
|
||||
});
|
||||
|
||||
let stream = history_stream.chain(live_stream);
|
||||
|
||||
Ok(Sse::new(stream).keep_alive(
|
||||
KeepAlive::new()
|
||||
.interval(std::time::Duration::from_secs(30))
|
||||
.text(""),
|
||||
Ok((
|
||||
[("X-Accel-Buffering", "no"), ("Cache-Control", "no-cache")],
|
||||
Sse::new(stream).keep_alive(
|
||||
KeepAlive::new()
|
||||
.interval(std::time::Duration::from_secs(30))
|
||||
.text(""),
|
||||
),
|
||||
))
|
||||
}
|
||||
|
||||
async fn logs_level_get_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let handle = state.log_level_handle.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Log level control not available".to_string(),
|
||||
))?;
|
||||
Ok(Json(serde_json::json!({ "level": handle.current_level() })))
|
||||
}
|
||||
|
||||
async fn logs_level_set_handler(
|
||||
State(state): State<Arc<GatewayState>>,
|
||||
Json(body): Json<serde_json::Value>,
|
||||
) -> Result<Json<serde_json::Value>, (StatusCode, String)> {
|
||||
let handle = state.log_level_handle.as_ref().ok_or((
|
||||
StatusCode::SERVICE_UNAVAILABLE,
|
||||
"Log level control not available".to_string(),
|
||||
))?;
|
||||
|
||||
let level = body
|
||||
.get("level")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or((StatusCode::BAD_REQUEST, "missing 'level' field".to_string()))?;
|
||||
|
||||
handle
|
||||
.set_level(level)
|
||||
.map_err(|e| (StatusCode::BAD_REQUEST, e))?;
|
||||
|
||||
tracing::info!("Log level changed to '{}'", handle.current_level());
|
||||
Ok(Json(serde_json::json!({ "level": handle.current_level() })))
|
||||
}
|
||||
|
||||
// --- Extension handlers ---
|
||||
|
||||
async fn extensions_list_handler(
|
||||
|
||||
@@ -29,16 +29,25 @@ function authenticate() {
|
||||
sessionStorage.setItem('ironclaw_token', token);
|
||||
document.getElementById('auth-screen').style.display = 'none';
|
||||
document.getElementById('app').style.display = 'flex';
|
||||
// Strip token from URL so it's not visible in the address bar
|
||||
// Strip token and log_level from URL so they're not visible in the address bar
|
||||
const cleaned = new URL(window.location);
|
||||
const urlLogLevel = cleaned.searchParams.get('log_level');
|
||||
cleaned.searchParams.delete('token');
|
||||
cleaned.searchParams.delete('log_level');
|
||||
window.history.replaceState({}, '', cleaned.pathname + cleaned.search);
|
||||
connectSSE();
|
||||
connectLogSSE();
|
||||
startGatewayStatusPolling();
|
||||
checkTeeStatus();
|
||||
loadThreads();
|
||||
loadMemoryTree();
|
||||
loadJobs();
|
||||
// Apply URL log_level param if present, otherwise just sync the dropdown
|
||||
if (urlLogLevel) {
|
||||
setServerLogLevel(urlLogLevel);
|
||||
} else {
|
||||
loadServerLogLevel();
|
||||
}
|
||||
})
|
||||
.catch(() => {
|
||||
sessionStorage.removeItem('ironclaw_token');
|
||||
@@ -1167,6 +1176,30 @@ function applyLogFilters() {
|
||||
}
|
||||
}
|
||||
|
||||
// --- Server-side log level control ---
|
||||
|
||||
function setServerLogLevel(level) {
|
||||
apiFetch('/api/logs/level', {
|
||||
method: 'PUT',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ level: level }),
|
||||
})
|
||||
.then(r => r.json())
|
||||
.then(data => {
|
||||
document.getElementById('logs-server-level').value = data.level;
|
||||
})
|
||||
.catch(err => console.error('Failed to set server log level:', err));
|
||||
}
|
||||
|
||||
function loadServerLogLevel() {
|
||||
apiFetch('/api/logs/level')
|
||||
.then(r => r.json())
|
||||
.then(data => {
|
||||
document.getElementById('logs-server-level').value = data.level;
|
||||
})
|
||||
.catch(() => {}); // ignore if not available
|
||||
}
|
||||
|
||||
// --- Extensions ---
|
||||
|
||||
function loadExtensions() {
|
||||
@@ -2066,6 +2099,100 @@ document.getElementById('gateway-status-trigger').addEventListener('mouseleave',
|
||||
document.getElementById('gateway-popover').classList.remove('visible');
|
||||
});
|
||||
|
||||
// --- TEE attestation ---
|
||||
|
||||
let teeInfo = null;
|
||||
let teeReportCache = null;
|
||||
let teeReportLoading = false;
|
||||
|
||||
function teeApiBase() {
|
||||
var parts = window.location.hostname.split('.');
|
||||
if (parts.length < 2) return null;
|
||||
var domain = parts.slice(1).join('.');
|
||||
return window.location.protocol + '//api.' + domain;
|
||||
}
|
||||
|
||||
function teeInstanceName() {
|
||||
return window.location.hostname.split('.')[0];
|
||||
}
|
||||
|
||||
function checkTeeStatus() {
|
||||
var base = teeApiBase();
|
||||
if (!base) return;
|
||||
var name = teeInstanceName();
|
||||
fetch(base + '/instances/' + encodeURIComponent(name) + '/attestation').then(function(res) {
|
||||
if (!res.ok) throw new Error(res.status);
|
||||
return res.json();
|
||||
}).then(function(data) {
|
||||
teeInfo = data;
|
||||
document.getElementById('tee-shield').style.display = 'flex';
|
||||
}).catch(function() {});
|
||||
}
|
||||
|
||||
function fetchTeeReport() {
|
||||
if (teeReportCache) {
|
||||
renderTeePopover(teeReportCache);
|
||||
return;
|
||||
}
|
||||
if (teeReportLoading) return;
|
||||
teeReportLoading = true;
|
||||
var base = teeApiBase();
|
||||
if (!base) return;
|
||||
var popover = document.getElementById('tee-popover');
|
||||
popover.innerHTML = '<div class="tee-popover-loading">Loading attestation report...</div>';
|
||||
fetch(base + '/attestation/report').then(function(res) {
|
||||
if (!res.ok) throw new Error(res.status);
|
||||
return res.json();
|
||||
}).then(function(data) {
|
||||
teeReportCache = data;
|
||||
renderTeePopover(data);
|
||||
}).catch(function() {
|
||||
popover.innerHTML = '<div class="tee-popover-loading">Could not load attestation report</div>';
|
||||
}).finally(function() {
|
||||
teeReportLoading = false;
|
||||
});
|
||||
}
|
||||
|
||||
function renderTeePopover(report) {
|
||||
var popover = document.getElementById('tee-popover');
|
||||
var digest = (teeInfo && teeInfo.image_digest) || 'N/A';
|
||||
var fingerprint = report.tls_certificate_fingerprint || 'N/A';
|
||||
var reportData = report.report_data || '';
|
||||
var vmConfig = report.vm_config || 'N/A';
|
||||
var truncated = reportData.length > 32 ? reportData.slice(0, 32) + '...' : reportData;
|
||||
popover.innerHTML = '<div class="tee-popover-title">'
|
||||
+ '<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M12 22s8-4 8-10V5l-8-3-8 3v7c0 6 8 10 8 10z"/></svg>'
|
||||
+ 'TEE Attestation</div>'
|
||||
+ '<div class="tee-field"><div class="tee-field-label">Image Digest</div>'
|
||||
+ '<div class="tee-field-value">' + escapeHtml(digest) + '</div></div>'
|
||||
+ '<div class="tee-field"><div class="tee-field-label">TLS Certificate Fingerprint</div>'
|
||||
+ '<div class="tee-field-value">' + escapeHtml(fingerprint) + '</div></div>'
|
||||
+ '<div class="tee-field"><div class="tee-field-label">Report Data</div>'
|
||||
+ '<div class="tee-field-value">' + escapeHtml(truncated) + '</div></div>'
|
||||
+ '<div class="tee-field"><div class="tee-field-label">VM Config</div>'
|
||||
+ '<div class="tee-field-value">' + escapeHtml(vmConfig) + '</div></div>'
|
||||
+ '<div class="tee-popover-actions">'
|
||||
+ '<button class="tee-btn-copy" onclick="copyTeeReport()">Copy Full Report</button></div>';
|
||||
}
|
||||
|
||||
function copyTeeReport() {
|
||||
if (!teeReportCache) return;
|
||||
var combined = Object.assign({}, teeReportCache, teeInfo || {});
|
||||
navigator.clipboard.writeText(JSON.stringify(combined, null, 2)).then(function() {
|
||||
showToast('Attestation report copied', 'success');
|
||||
}).catch(function() {
|
||||
showToast('Failed to copy report', 'error');
|
||||
});
|
||||
}
|
||||
|
||||
document.getElementById('tee-shield').addEventListener('mouseenter', function() {
|
||||
fetchTeeReport();
|
||||
document.getElementById('tee-popover').classList.add('visible');
|
||||
});
|
||||
document.getElementById('tee-shield').addEventListener('mouseleave', function() {
|
||||
document.getElementById('tee-popover').classList.remove('visible');
|
||||
});
|
||||
|
||||
// --- Extension install ---
|
||||
|
||||
function installExtension() {
|
||||
|
||||
@@ -40,6 +40,13 @@
|
||||
<button data-tab="routines">Routines</button>
|
||||
<button data-tab="extensions">Extensions</button>
|
||||
<div class="spacer"></div>
|
||||
<div class="tee-shield" id="tee-shield" style="display:none" title="Running in a Trusted Execution Environment">
|
||||
<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
|
||||
<path d="M12 22s8-4 8-10V5l-8-3-8 3v7c0 6 8 10 8 10z"/>
|
||||
</svg>
|
||||
<span id="tee-shield-label">TEE Verified</span>
|
||||
<div class="tee-popover" id="tee-popover"></div>
|
||||
</div>
|
||||
<div class="status" id="gateway-status-trigger">
|
||||
<div class="dot" id="sse-dot"></div>
|
||||
<span id="sse-status">Connected</span>
|
||||
@@ -127,6 +134,12 @@
|
||||
<div class="tab-panel" id="tab-logs">
|
||||
<div class="logs-container">
|
||||
<div class="logs-toolbar">
|
||||
<select id="logs-server-level" onchange="setServerLogLevel(this.value)" title="Server-side log level (changes what the server emits)">
|
||||
<option value="error">Server: ERROR</option>
|
||||
<option value="warn">Server: WARN</option>
|
||||
<option value="info" selected>Server: INFO</option>
|
||||
<option value="debug">Server: DEBUG</option>
|
||||
</select>
|
||||
<select id="logs-level-filter">
|
||||
<option value="all">All Levels</option>
|
||||
<option value="ERROR">Error</option>
|
||||
|
||||
@@ -189,6 +189,126 @@ body {
|
||||
background: var(--danger);
|
||||
}
|
||||
|
||||
/* TEE Shield */
|
||||
.tee-shield {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 6px;
|
||||
font-size: 12px;
|
||||
color: var(--success);
|
||||
padding: 4px 10px;
|
||||
border-radius: 12px;
|
||||
background: rgba(63, 185, 80, 0.1);
|
||||
border: 1px solid rgba(63, 185, 80, 0.25);
|
||||
cursor: pointer;
|
||||
position: relative;
|
||||
margin-right: 8px;
|
||||
transition: background 0.15s;
|
||||
}
|
||||
|
||||
.tee-shield:hover {
|
||||
background: rgba(63, 185, 80, 0.18);
|
||||
}
|
||||
|
||||
.tee-shield svg {
|
||||
flex-shrink: 0;
|
||||
}
|
||||
|
||||
#tee-shield-label {
|
||||
font-weight: 500;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.tee-popover {
|
||||
display: none;
|
||||
position: absolute;
|
||||
top: 100%;
|
||||
right: 0;
|
||||
margin-top: 8px;
|
||||
background: var(--bg-secondary);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: var(--radius);
|
||||
padding: 16px;
|
||||
min-width: 340px;
|
||||
max-width: 420px;
|
||||
z-index: 100;
|
||||
box-shadow: var(--shadow);
|
||||
}
|
||||
|
||||
.tee-popover.visible {
|
||||
display: block;
|
||||
}
|
||||
|
||||
.tee-popover-title {
|
||||
font-size: 13px;
|
||||
font-weight: 600;
|
||||
color: var(--text);
|
||||
margin-bottom: 12px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 6px;
|
||||
}
|
||||
|
||||
.tee-popover-title svg {
|
||||
color: var(--success);
|
||||
}
|
||||
|
||||
.tee-field {
|
||||
margin-bottom: 10px;
|
||||
}
|
||||
|
||||
.tee-field:last-child {
|
||||
margin-bottom: 0;
|
||||
}
|
||||
|
||||
.tee-field-label {
|
||||
font-size: 11px;
|
||||
font-weight: 500;
|
||||
color: var(--text-secondary);
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.5px;
|
||||
margin-bottom: 3px;
|
||||
}
|
||||
|
||||
.tee-field-value {
|
||||
font-size: 12px;
|
||||
font-family: "SFMono-Regular", Consolas, "Liberation Mono", Menlo, monospace;
|
||||
color: var(--text);
|
||||
word-break: break-all;
|
||||
background: var(--bg);
|
||||
padding: 4px 8px;
|
||||
border-radius: 4px;
|
||||
border: 1px solid var(--border);
|
||||
}
|
||||
|
||||
.tee-popover-actions {
|
||||
margin-top: 12px;
|
||||
display: flex;
|
||||
gap: 8px;
|
||||
}
|
||||
|
||||
.tee-btn-copy {
|
||||
padding: 4px 10px;
|
||||
background: none;
|
||||
border: 1px solid var(--border);
|
||||
border-radius: var(--radius);
|
||||
color: var(--text-secondary);
|
||||
cursor: pointer;
|
||||
font-size: 11px;
|
||||
transition: color 0.15s, border-color 0.15s;
|
||||
}
|
||||
|
||||
.tee-btn-copy:hover {
|
||||
color: var(--text);
|
||||
border-color: var(--text-secondary);
|
||||
}
|
||||
|
||||
.tee-popover-loading {
|
||||
font-size: 12px;
|
||||
color: var(--text-secondary);
|
||||
padding: 8px 0;
|
||||
}
|
||||
|
||||
/* Tab Panels */
|
||||
.tab-panel {
|
||||
display: none;
|
||||
|
||||
@@ -477,6 +477,7 @@ mod tests {
|
||||
workspace: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
|
||||
@@ -23,6 +23,10 @@ pub struct AgentConfig {
|
||||
pub max_cost_per_day_cents: Option<u64>,
|
||||
/// Maximum LLM/tool actions per hour. None = unlimited.
|
||||
pub max_actions_per_hour: Option<u64>,
|
||||
/// Maximum tool-call iterations per agentic loop invocation. Default 50.
|
||||
pub max_tool_iterations: usize,
|
||||
/// When true, skip tool approval checks entirely. For benchmarks/CI.
|
||||
pub auto_approve_tools: bool,
|
||||
}
|
||||
|
||||
impl AgentConfig {
|
||||
@@ -115,6 +119,22 @@ impl AgentConfig {
|
||||
key: "MAX_ACTIONS_PER_HOUR".to_string(),
|
||||
message: format!("must be a positive integer: {e}"),
|
||||
})?,
|
||||
max_tool_iterations: optional_env("AGENT_MAX_TOOL_ITERATIONS")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "AGENT_MAX_TOOL_ITERATIONS".to_string(),
|
||||
message: format!("must be a positive integer: {e}"),
|
||||
})?
|
||||
.unwrap_or(settings.agent.max_tool_iterations),
|
||||
auto_approve_tools: optional_env("AGENT_AUTO_APPROVE_TOOLS")?
|
||||
.map(|s| s.parse())
|
||||
.transpose()
|
||||
.map_err(|e| ConfigError::InvalidValue {
|
||||
key: "AGENT_AUTO_APPROVE_TOOLS".to_string(),
|
||||
message: format!("must be 'true' or 'false': {e}"),
|
||||
})?
|
||||
.unwrap_or(settings.agent.auto_approve_tools),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+4
-49
@@ -121,37 +121,7 @@ pub struct LlmConfig {
|
||||
pub tinfoil: Option<TinfoilConfig>,
|
||||
}
|
||||
|
||||
/// API mode for NEAR AI.
|
||||
///
|
||||
/// - `Responses` = **NEAR AI Chat** (`private.near.ai`, session token auth)
|
||||
/// - `ChatCompletions` = **NEAR AI Cloud** (`cloud-api.near.ai`, API key auth)
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum NearAiApiMode {
|
||||
/// NEAR AI Chat: Responses API with session token auth
|
||||
#[default]
|
||||
Responses,
|
||||
/// NEAR AI Cloud: Chat Completions API with API key auth
|
||||
ChatCompletions,
|
||||
}
|
||||
|
||||
impl std::str::FromStr for NearAiApiMode {
|
||||
type Err = String;
|
||||
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
match s.to_lowercase().as_str() {
|
||||
"responses" | "response" => Ok(Self::Responses),
|
||||
"chat_completions" | "chatcompletions" | "chat" | "completions" => {
|
||||
Ok(Self::ChatCompletions)
|
||||
}
|
||||
_ => Err(format!(
|
||||
"invalid API mode '{}', expected 'responses' or 'chat_completions'",
|
||||
s
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// NEAR AI configuration (shared by Chat and Cloud modes).
|
||||
/// NEAR AI configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NearAiConfig {
|
||||
/// Model to use (e.g., "claude-3-5-sonnet-20241022", "gpt-4o")
|
||||
@@ -160,16 +130,13 @@ pub struct NearAiConfig {
|
||||
/// Falls back to the main model if not set.
|
||||
pub cheap_model: Option<String>,
|
||||
/// Base URL for the NEAR AI API.
|
||||
/// Chat mode default: `https://private.near.ai`
|
||||
/// Cloud mode default: `https://cloud-api.near.ai`
|
||||
/// Default: `https://private.near.ai` (session token) or `https://cloud-api.near.ai` (API key)
|
||||
pub base_url: String,
|
||||
/// Base URL for auth/refresh endpoints (default: https://private.near.ai)
|
||||
pub auth_base_url: String,
|
||||
/// Path to session file (default: ~/.ironclaw/session.json)
|
||||
pub session_path: PathBuf,
|
||||
/// API mode: NEAR AI Chat (Responses) or NEAR AI Cloud (ChatCompletions)
|
||||
pub api_mode: NearAiApiMode,
|
||||
/// API key for NEAR AI Cloud (required for ChatCompletions mode)
|
||||
/// API key for NEAR AI Cloud. When set, uses API key auth; otherwise uses session token auth.
|
||||
pub api_key: Option<SecretString>,
|
||||
/// Optional fallback model for failover (default: None).
|
||||
/// When set, a secondary provider is created with this model and wrapped
|
||||
@@ -229,17 +196,6 @@ impl LlmConfig {
|
||||
// Resolve NEAR AI config only when backend is NearAi (or when explicitly configured)
|
||||
let nearai_api_key = optional_env("NEARAI_API_KEY")?.map(SecretString::from);
|
||||
|
||||
let api_mode = if let Some(mode_str) = optional_env("NEARAI_API_MODE")? {
|
||||
mode_str.parse().map_err(|e| ConfigError::InvalidValue {
|
||||
key: "NEARAI_API_MODE".to_string(),
|
||||
message: e,
|
||||
})?
|
||||
} else if nearai_api_key.is_some() {
|
||||
NearAiApiMode::ChatCompletions
|
||||
} else {
|
||||
NearAiApiMode::Responses
|
||||
};
|
||||
|
||||
let nearai = NearAiConfig {
|
||||
model: optional_env("NEARAI_MODEL")?
|
||||
.or_else(|| settings.selected_model.clone())
|
||||
@@ -249,7 +205,7 @@ impl LlmConfig {
|
||||
}),
|
||||
cheap_model: optional_env("NEARAI_CHEAP_MODEL")?,
|
||||
base_url: optional_env("NEARAI_BASE_URL")?.unwrap_or_else(|| {
|
||||
if api_mode == NearAiApiMode::ChatCompletions {
|
||||
if nearai_api_key.is_some() {
|
||||
"https://cloud-api.near.ai".to_string()
|
||||
} else {
|
||||
"https://private.near.ai".to_string()
|
||||
@@ -260,7 +216,6 @@ impl LlmConfig {
|
||||
session_path: optional_env("NEARAI_SESSION_PATH")?
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(default_session_path),
|
||||
api_mode,
|
||||
api_key: nearai_api_key,
|
||||
fallback_model: optional_env("NEARAI_FALLBACK_MODEL")?,
|
||||
max_retries: parse_optional_env("NEARAI_MAX_RETRIES", 3)?,
|
||||
|
||||
+1
-1
@@ -37,7 +37,7 @@ pub use self::embeddings::EmbeddingsConfig;
|
||||
pub use self::heartbeat::HeartbeatConfig;
|
||||
pub use self::hygiene::HygieneConfig;
|
||||
pub use self::llm::{
|
||||
AnthropicDirectConfig, LlmBackend, LlmConfig, NearAiApiMode, NearAiConfig, OllamaConfig,
|
||||
AnthropicDirectConfig, LlmBackend, LlmConfig, NearAiConfig, OllamaConfig,
|
||||
OpenAiCompatibleConfig, OpenAiDirectConfig, TinfoilConfig,
|
||||
};
|
||||
pub use self::routines::RoutineConfig;
|
||||
|
||||
@@ -296,14 +296,6 @@ impl LlmProvider for CircuitBreakerProvider {
|
||||
self.inner.set_model(model)
|
||||
}
|
||||
|
||||
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||
self.inner.seed_response_chain(thread_id, response_id)
|
||||
}
|
||||
|
||||
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||
self.inner.get_response_chain_id(thread_id)
|
||||
}
|
||||
|
||||
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||
self.inner.calculate_cost(input_tokens, output_tokens)
|
||||
}
|
||||
|
||||
@@ -359,15 +359,6 @@ impl LlmProvider for FailoverProvider {
|
||||
.await
|
||||
}
|
||||
|
||||
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||
self.providers[self.last_used.load(Ordering::Relaxed)]
|
||||
.seed_response_chain(thread_id, response_id);
|
||||
}
|
||||
|
||||
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||
self.providers[self.last_used.load(Ordering::Relaxed)].get_response_chain_id(thread_id)
|
||||
}
|
||||
|
||||
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||
self.providers[self.last_used.load(Ordering::Relaxed)]
|
||||
.calculate_cost(input_tokens, output_tokens)
|
||||
@@ -413,7 +404,6 @@ mod tests {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
}))),
|
||||
tool_complete_result: Mutex::new(Some(Ok(ToolCompletionResponse {
|
||||
content: Some(content.to_string()),
|
||||
@@ -421,7 +411,6 @@ mod tests {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
}))),
|
||||
}
|
||||
}
|
||||
@@ -803,7 +792,6 @@ mod tests {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -829,7 +817,6 @@ mod tests {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+21
-34
@@ -1,7 +1,7 @@
|
||||
//! LLM integration for the agent.
|
||||
//!
|
||||
//! Supports multiple backends:
|
||||
//! - **NEAR AI** (default): Session-based or API key auth via NEAR AI proxy
|
||||
//! - **NEAR AI** (default): Session token or API key auth via Chat Completions API
|
||||
//! - **OpenAI**: Direct API access with your own key
|
||||
//! - **Anthropic**: Direct API access with your own key
|
||||
//! - **Ollama**: Local model inference
|
||||
@@ -10,7 +10,6 @@
|
||||
pub mod circuit_breaker;
|
||||
pub mod costs;
|
||||
pub mod failover;
|
||||
mod nearai;
|
||||
mod nearai_chat;
|
||||
mod provider;
|
||||
mod reasoning;
|
||||
@@ -21,8 +20,7 @@ pub mod session;
|
||||
|
||||
pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider};
|
||||
pub use failover::{CooldownConfig, FailoverProvider};
|
||||
pub use nearai::{ModelInfo, NearAiProvider};
|
||||
pub use nearai_chat::NearAiChatProvider;
|
||||
pub use nearai_chat::{ModelInfo, NearAiChatProvider};
|
||||
pub use provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
||||
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse, ToolDefinition, ToolResult,
|
||||
@@ -41,7 +39,7 @@ use std::sync::Arc;
|
||||
use rig::client::CompletionClient;
|
||||
use secrecy::ExposeSecret;
|
||||
|
||||
use crate::config::{LlmBackend, LlmConfig, NearAiApiMode, NearAiConfig};
|
||||
use crate::config::{LlmBackend, LlmConfig, NearAiConfig};
|
||||
use crate::error::LlmError;
|
||||
|
||||
/// Create an LLM provider based on configuration.
|
||||
@@ -71,24 +69,18 @@ pub fn create_llm_provider_with_config(
|
||||
config: &NearAiConfig,
|
||||
session: Arc<SessionManager>,
|
||||
) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
match config.api_mode {
|
||||
NearAiApiMode::Responses => {
|
||||
tracing::info!(
|
||||
model = %config.model,
|
||||
base_url = %config.base_url,
|
||||
"Using NEAR AI Chat (Responses API, session token auth)"
|
||||
);
|
||||
Ok(Arc::new(NearAiProvider::new(config.clone(), session)?))
|
||||
}
|
||||
NearAiApiMode::ChatCompletions => {
|
||||
tracing::info!(
|
||||
model = %config.model,
|
||||
base_url = %config.base_url,
|
||||
"Using NEAR AI Cloud (Chat Completions API, API key auth)"
|
||||
);
|
||||
Ok(Arc::new(NearAiChatProvider::new(config.clone())?))
|
||||
}
|
||||
}
|
||||
let auth_mode = if config.api_key.is_some() {
|
||||
"API key"
|
||||
} else {
|
||||
"session token"
|
||||
};
|
||||
tracing::info!(
|
||||
model = %config.model,
|
||||
base_url = %config.base_url,
|
||||
auth = auth_mode,
|
||||
"Using NEAR AI (Chat Completions API)"
|
||||
);
|
||||
Ok(Arc::new(NearAiChatProvider::new(config.clone(), session)?))
|
||||
}
|
||||
|
||||
fn create_openai_provider(config: &LlmConfig) -> Result<Arc<dyn LlmProvider>, LlmError> {
|
||||
@@ -254,7 +246,7 @@ fn create_openai_compatible_provider(config: &LlmConfig) -> Result<Arc<dyn LlmPr
|
||||
/// Create a cheap/fast LLM provider for lightweight tasks (heartbeat, routing, evaluation).
|
||||
///
|
||||
/// Uses `NEARAI_CHEAP_MODEL` if set, otherwise falls back to the main provider.
|
||||
/// Currently only supports NEAR AI backends (Responses and ChatCompletions modes).
|
||||
/// Currently only supports NEAR AI backend.
|
||||
pub fn create_cheap_llm_provider(
|
||||
config: &LlmConfig,
|
||||
session: Arc<SessionManager>,
|
||||
@@ -275,20 +267,16 @@ pub fn create_cheap_llm_provider(
|
||||
let mut cheap_config = config.nearai.clone();
|
||||
cheap_config.model = cheap_model.clone();
|
||||
|
||||
tracing::info!("Cheap LLM provider: {}", cheap_model);
|
||||
|
||||
match cheap_config.api_mode {
|
||||
NearAiApiMode::Responses => Ok(Some(Arc::new(NearAiProvider::new(cheap_config, session)?))),
|
||||
NearAiApiMode::ChatCompletions => {
|
||||
Ok(Some(Arc::new(NearAiChatProvider::new(cheap_config)?)))
|
||||
}
|
||||
}
|
||||
Ok(Some(Arc::new(NearAiChatProvider::new(
|
||||
cheap_config,
|
||||
session,
|
||||
)?)))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::config::{LlmBackend, NearAiApiMode, NearAiConfig};
|
||||
use crate::config::{LlmBackend, NearAiConfig};
|
||||
use std::path::PathBuf;
|
||||
|
||||
fn test_nearai_config() -> NearAiConfig {
|
||||
@@ -298,7 +286,6 @@ mod tests {
|
||||
base_url: "https://api.near.ai".to_string(),
|
||||
auth_base_url: "https://private.near.ai".to_string(),
|
||||
session_path: PathBuf::from("/tmp/test-session.json"),
|
||||
api_mode: NearAiApiMode::Responses,
|
||||
api_key: None,
|
||||
fallback_model: None,
|
||||
max_retries: 3,
|
||||
|
||||
-1205
File diff suppressed because it is too large
Load Diff
+199
-61
@@ -1,8 +1,12 @@
|
||||
//! NEAR AI Cloud provider implementation (Chat Completions API).
|
||||
//! NEAR AI provider implementation (Chat Completions API).
|
||||
//!
|
||||
//! This provider uses the NEAR AI Cloud API (`cloud-api.near.ai`) which
|
||||
//! exposes an OpenAI-compatible chat completions endpoint with API key
|
||||
//! authentication.
|
||||
//! This provider uses the OpenAI-compatible Chat Completions endpoint with
|
||||
//! dual auth support:
|
||||
//! - **API key auth**: When `NEARAI_API_KEY` is set, uses Bearer API key
|
||||
//! - **Session token auth**: Otherwise, uses `SessionManager` for Bearer session token
|
||||
//! with automatic renewal on 401 errors
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use reqwest::Client;
|
||||
@@ -14,38 +18,51 @@ use serde::{Deserialize, Serialize};
|
||||
use crate::config::NearAiConfig;
|
||||
use crate::error::LlmError;
|
||||
use crate::llm::provider::{
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, ModelMetadata,
|
||||
Role, ToolCall, ToolCompletionRequest, ToolCompletionResponse,
|
||||
ChatMessage, CompletionRequest, CompletionResponse, FinishReason, LlmProvider, Role, ToolCall,
|
||||
ToolCompletionRequest, ToolCompletionResponse,
|
||||
};
|
||||
use crate::llm::session::SessionManager;
|
||||
|
||||
/// NEAR AI Cloud provider (Chat Completions API, API key auth).
|
||||
/// Information about an available model from NEAR AI API.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ModelInfo {
|
||||
/// Model identifier.
|
||||
#[serde(alias = "id", alias = "model")]
|
||||
pub name: String,
|
||||
/// Optional provider name.
|
||||
#[serde(default)]
|
||||
pub provider: Option<String>,
|
||||
}
|
||||
|
||||
/// NEAR AI provider (Chat Completions API, dual auth).
|
||||
pub struct NearAiChatProvider {
|
||||
client: Client,
|
||||
config: NearAiConfig,
|
||||
/// Session manager for session token auth (used when no API key is set).
|
||||
session: Arc<SessionManager>,
|
||||
active_model: std::sync::RwLock<String>,
|
||||
flatten_tool_messages: bool,
|
||||
}
|
||||
|
||||
impl NearAiChatProvider {
|
||||
/// Create a new NEAR AI Cloud provider with API key auth.
|
||||
/// Create a new NEAR AI Chat Completions provider.
|
||||
///
|
||||
/// Auth mode is determined by `config.api_key`:
|
||||
/// - If set, uses Bearer API key auth
|
||||
/// - If not set, uses session token auth via `SessionManager`
|
||||
///
|
||||
/// By default this enables tool-message flattening for compatibility with
|
||||
/// providers that reject `role: "tool"` messages.
|
||||
pub fn new(config: NearAiConfig) -> Result<Self, LlmError> {
|
||||
Self::new_with_flatten(config, true)
|
||||
pub fn new(config: NearAiConfig, session: Arc<SessionManager>) -> Result<Self, LlmError> {
|
||||
Self::new_with_flatten(config, session, true)
|
||||
}
|
||||
|
||||
/// Create a chat completions provider with configurable tool-message flattening.
|
||||
pub fn new_with_flatten(
|
||||
config: NearAiConfig,
|
||||
session: Arc<SessionManager>,
|
||||
flatten_tool_messages: bool,
|
||||
) -> Result<Self, LlmError> {
|
||||
if config.api_key.is_none() {
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let client = Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.build()
|
||||
@@ -58,6 +75,7 @@ impl NearAiChatProvider {
|
||||
Ok(Self {
|
||||
client,
|
||||
config,
|
||||
session,
|
||||
active_model,
|
||||
flatten_tool_messages,
|
||||
})
|
||||
@@ -74,23 +92,50 @@ impl NearAiChatProvider {
|
||||
}
|
||||
}
|
||||
|
||||
fn api_key(&self) -> String {
|
||||
self.config
|
||||
.api_key
|
||||
.as_ref()
|
||||
.map(|k| k.expose_secret().to_string())
|
||||
.unwrap_or_default()
|
||||
/// Returns true if using API key auth, false if session token auth.
|
||||
fn uses_api_key(&self) -> bool {
|
||||
self.config.api_key.is_some()
|
||||
}
|
||||
|
||||
/// Resolve the Bearer token for the current auth mode.
|
||||
async fn resolve_bearer_token(&self) -> Result<String, LlmError> {
|
||||
if let Some(ref api_key) = self.config.api_key {
|
||||
Ok(api_key.expose_secret().to_string())
|
||||
} else {
|
||||
let token = self.session.get_token().await?;
|
||||
Ok(token.expose_secret().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a single request to the chat completions API.
|
||||
///
|
||||
/// Does not retry internally — retries are handled by the external
|
||||
/// For session token auth, handles 401 by calling `session.handle_auth_failure()`
|
||||
/// and retrying once.
|
||||
///
|
||||
/// Does not retry on other errors — retries are handled by the external
|
||||
/// `RetryProvider` wrapper in the composition chain.
|
||||
async fn send_request<T: Serialize, R: for<'de> Deserialize<'de>>(
|
||||
&self,
|
||||
body: &T,
|
||||
) -> Result<R, LlmError> {
|
||||
match self.send_request_inner(body).await {
|
||||
Ok(result) => Ok(result),
|
||||
Err(LlmError::SessionExpired { .. }) if !self.uses_api_key() => {
|
||||
// Session expired, attempt renewal and retry once
|
||||
self.session.handle_auth_failure().await?;
|
||||
self.send_request_inner(body).await
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Inner request implementation (single attempt).
|
||||
async fn send_request_inner<T: Serialize, R: for<'de> Deserialize<'de>>(
|
||||
&self,
|
||||
body: &T,
|
||||
) -> Result<R, LlmError> {
|
||||
let url = self.api_url("chat/completions");
|
||||
let token = self.resolve_bearer_token().await?;
|
||||
|
||||
tracing::debug!("Sending request to NEAR AI Chat: {}", url);
|
||||
|
||||
@@ -103,7 +148,7 @@ impl NearAiChatProvider {
|
||||
let response = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key()))
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(body)
|
||||
.send()
|
||||
@@ -126,6 +171,17 @@ impl NearAiChatProvider {
|
||||
let status_code = status.as_u16();
|
||||
|
||||
if status_code == 401 {
|
||||
// For session token auth, distinguish session expired from plain auth failure
|
||||
if !self.uses_api_key() {
|
||||
let lower = response_text.to_lowercase();
|
||||
let is_session_expired = lower.contains("session")
|
||||
&& (lower.contains("expired") || lower.contains("invalid"));
|
||||
if is_session_expired {
|
||||
return Err(LlmError::SessionExpired {
|
||||
provider: "nearai_chat".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
return Err(LlmError::AuthFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
});
|
||||
@@ -154,14 +210,31 @@ impl NearAiChatProvider {
|
||||
})
|
||||
}
|
||||
|
||||
/// Fetch available models with full metadata from the `/v1/models` endpoint.
|
||||
async fn fetch_models(&self) -> Result<Vec<ApiModelEntry>, LlmError> {
|
||||
/// Fetch available models from the NEAR AI API.
|
||||
///
|
||||
/// Handles session renewal on 401 (same pattern as `send_request`).
|
||||
/// Supports multiple response formats: `{models: [...]}`, `{data: [...]}`, and plain array.
|
||||
pub async fn list_models_full(&self) -> Result<Vec<ModelInfo>, LlmError> {
|
||||
match self.list_models_inner().await {
|
||||
Ok(models) => Ok(models),
|
||||
Err(LlmError::SessionExpired { .. }) if !self.uses_api_key() => {
|
||||
self.session.handle_auth_failure().await?;
|
||||
self.list_models_inner().await
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
async fn list_models_inner(&self) -> Result<Vec<ModelInfo>, LlmError> {
|
||||
let url = self.api_url("models");
|
||||
let token = self.resolve_bearer_token().await?;
|
||||
|
||||
tracing::debug!("Fetching models from: {}", url);
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", self.api_key()))
|
||||
.header("Authorization", format!("Bearer {}", token))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| LlmError::RequestFailed {
|
||||
@@ -176,6 +249,11 @@ impl NearAiChatProvider {
|
||||
})?;
|
||||
|
||||
if !status.is_success() {
|
||||
if status.as_u16() == 401 && !self.uses_api_key() {
|
||||
return Err(LlmError::SessionExpired {
|
||||
provider: "nearai_chat".to_string(),
|
||||
});
|
||||
}
|
||||
let truncated = crate::agent::truncate_for_preview(&response_text, 512);
|
||||
return Err(LlmError::RequestFailed {
|
||||
provider: "nearai_chat".to_string(),
|
||||
@@ -183,29 +261,97 @@ impl NearAiChatProvider {
|
||||
});
|
||||
}
|
||||
|
||||
// Flexible model entry parsing -- handle various field names
|
||||
#[derive(Deserialize)]
|
||||
struct ModelsResponse {
|
||||
data: Vec<ApiModelEntry>,
|
||||
struct ModelMetadataInner {
|
||||
#[serde(default)]
|
||||
name: Option<String>,
|
||||
#[serde(default, alias = "modelName", alias = "model_name")]
|
||||
model_name: Option<String>,
|
||||
}
|
||||
|
||||
let resp: ModelsResponse =
|
||||
serde_json::from_str(&response_text).map_err(|e| LlmError::InvalidResponse {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!("JSON parse error: {}", e),
|
||||
})?;
|
||||
#[derive(Deserialize)]
|
||||
struct ModelEntry {
|
||||
#[serde(default)]
|
||||
name: Option<String>,
|
||||
#[serde(default)]
|
||||
id: Option<String>,
|
||||
#[serde(default)]
|
||||
model: Option<String>,
|
||||
#[serde(default, alias = "modelName", alias = "model_name")]
|
||||
model_name: Option<String>,
|
||||
#[serde(default, alias = "modelId", alias = "model_id")]
|
||||
model_id: Option<String>,
|
||||
#[serde(default)]
|
||||
metadata: Option<ModelMetadataInner>,
|
||||
}
|
||||
|
||||
Ok(resp.data)
|
||||
impl ModelEntry {
|
||||
fn get_name(&self) -> Option<String> {
|
||||
self.name
|
||||
.clone()
|
||||
.or_else(|| self.id.clone())
|
||||
.or_else(|| self.model.clone())
|
||||
.or_else(|| self.model_name.clone())
|
||||
.or_else(|| self.model_id.clone())
|
||||
.or_else(|| self.metadata.as_ref().and_then(|m| m.name.clone()))
|
||||
.or_else(|| self.metadata.as_ref().and_then(|m| m.model_name.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ModelsResponse {
|
||||
#[serde(default)]
|
||||
models: Option<Vec<ModelEntry>>,
|
||||
#[serde(default)]
|
||||
data: Option<Vec<ModelEntry>>,
|
||||
}
|
||||
|
||||
// Try {models: [...]} or {data: [...]} format
|
||||
if let Ok(resp) = serde_json::from_str::<ModelsResponse>(&response_text)
|
||||
&& let Some(entries) = resp.models.or(resp.data)
|
||||
{
|
||||
let models: Vec<ModelInfo> = entries
|
||||
.into_iter()
|
||||
.filter_map(|e| {
|
||||
e.get_name().map(|name| ModelInfo {
|
||||
name,
|
||||
provider: None,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
if !models.is_empty() {
|
||||
return Ok(models);
|
||||
}
|
||||
}
|
||||
|
||||
// Try direct array format
|
||||
if let Ok(entries) = serde_json::from_str::<Vec<ModelEntry>>(&response_text) {
|
||||
let models: Vec<ModelInfo> = entries
|
||||
.into_iter()
|
||||
.filter_map(|e| {
|
||||
e.get_name().map(|name| ModelInfo {
|
||||
name,
|
||||
provider: None,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
if !models.is_empty() {
|
||||
return Ok(models);
|
||||
}
|
||||
}
|
||||
|
||||
// Couldn't find model names in response
|
||||
Err(LlmError::InvalidResponse {
|
||||
provider: "nearai_chat".to_string(),
|
||||
reason: format!(
|
||||
"No model names found in response: {}",
|
||||
&response_text[..response_text.len().min(300)]
|
||||
),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Model entry as returned by the `/v1/models` API.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct ApiModelEntry {
|
||||
id: String,
|
||||
#[serde(default)]
|
||||
context_length: Option<u32>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmProvider for NearAiChatProvider {
|
||||
async fn complete(&self, req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
|
||||
@@ -252,7 +398,6 @@ impl LlmProvider for NearAiChatProvider {
|
||||
finish_reason,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -347,7 +492,6 @@ impl LlmProvider for NearAiChatProvider {
|
||||
finish_reason,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -361,18 +505,8 @@ impl LlmProvider for NearAiChatProvider {
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<String>, LlmError> {
|
||||
let models = self.fetch_models().await?;
|
||||
Ok(models.into_iter().map(|m| m.id).collect())
|
||||
}
|
||||
|
||||
async fn model_metadata(&self) -> Result<ModelMetadata, LlmError> {
|
||||
let active = self.active_model_name();
|
||||
let models = self.fetch_models().await?;
|
||||
let current = models.iter().find(|m| m.id == active);
|
||||
Ok(ModelMetadata {
|
||||
id: active,
|
||||
context_length: current.and_then(|m| m.context_length),
|
||||
})
|
||||
let models = self.list_models_full().await?;
|
||||
Ok(models.into_iter().map(|m| m.name).collect())
|
||||
}
|
||||
|
||||
fn active_model_name(&self) -> String {
|
||||
@@ -613,6 +747,7 @@ fn parse_usage(usage: Option<&ChatCompletionUsage>) -> (u32, u32) {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::llm::session::SessionConfig;
|
||||
|
||||
fn test_nearai_config(base_url: &str) -> NearAiConfig {
|
||||
NearAiConfig {
|
||||
@@ -620,7 +755,6 @@ mod tests {
|
||||
base_url: base_url.to_string(),
|
||||
auth_base_url: "https://private.near.ai".to_string(),
|
||||
session_path: std::path::PathBuf::from("/tmp/session.json"),
|
||||
api_mode: crate::config::NearAiApiMode::ChatCompletions,
|
||||
api_key: Some(secrecy::SecretString::from("test-key".to_string())),
|
||||
cheap_model: None,
|
||||
fallback_model: None,
|
||||
@@ -635,18 +769,22 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
fn test_session() -> Arc<SessionManager> {
|
||||
Arc::new(SessionManager::new(SessionConfig::default()))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_api_url_with_base_without_v1() {
|
||||
let mut cfg = test_nearai_config("http://127.0.0.1:8318");
|
||||
|
||||
let provider = NearAiChatProvider::new(cfg.clone()).expect("provider");
|
||||
let provider = NearAiChatProvider::new(cfg.clone(), test_session()).expect("provider");
|
||||
assert_eq!(
|
||||
provider.api_url("chat/completions"),
|
||||
"http://127.0.0.1:8318/v1/chat/completions"
|
||||
);
|
||||
|
||||
cfg.base_url = "http://127.0.0.1:8318/".to_string();
|
||||
let provider = NearAiChatProvider::new(cfg).expect("provider");
|
||||
let provider = NearAiChatProvider::new(cfg, test_session()).expect("provider");
|
||||
assert_eq!(
|
||||
provider.api_url("/chat/completions"),
|
||||
"http://127.0.0.1:8318/v1/chat/completions"
|
||||
@@ -657,7 +795,7 @@ mod tests {
|
||||
fn test_api_url_with_base_already_v1() {
|
||||
let cfg = test_nearai_config("http://127.0.0.1:8318/v1");
|
||||
|
||||
let provider = NearAiChatProvider::new(cfg).expect("provider");
|
||||
let provider = NearAiChatProvider::new(cfg, test_session()).expect("provider");
|
||||
assert_eq!(
|
||||
provider.api_url("chat/completions"),
|
||||
"http://127.0.0.1:8318/v1/chat/completions"
|
||||
|
||||
@@ -153,8 +153,6 @@ pub struct CompletionResponse {
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub finish_reason: FinishReason,
|
||||
/// Provider-specific response ID (e.g. for NEAR AI response chaining).
|
||||
pub response_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Why the completion finished.
|
||||
@@ -256,8 +254,6 @@ pub struct ToolCompletionResponse {
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub finish_reason: FinishReason,
|
||||
/// Provider-specific response ID (e.g. for NEAR AI response chaining).
|
||||
pub response_id: Option<String>,
|
||||
}
|
||||
|
||||
/// Metadata about a model returned by the provider's API.
|
||||
@@ -327,20 +323,6 @@ pub trait LlmProvider: Send + Sync {
|
||||
})
|
||||
}
|
||||
|
||||
/// Seed a response chain for a thread (e.g. restoring from DB).
|
||||
///
|
||||
/// Providers that support response chaining (e.g. NEAR AI `previous_response_id`)
|
||||
/// store this so subsequent calls send only delta messages.
|
||||
fn seed_response_chain(&self, _thread_id: &str, _response_id: String) {}
|
||||
|
||||
/// Get the last response chain ID for a thread.
|
||||
///
|
||||
/// Returns `None` if the provider doesn't support chaining or has no
|
||||
/// stored state for this thread.
|
||||
fn get_response_chain_id(&self, _thread_id: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Calculate cost for a completion.
|
||||
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||
let (input_cost, output_cost) = self.cost_per_token();
|
||||
|
||||
+687
-159
File diff suppressed because it is too large
Load Diff
@@ -228,14 +228,6 @@ impl LlmProvider for CachedProvider {
|
||||
fn set_model(&self, model: &str) -> Result<(), LlmError> {
|
||||
self.inner.set_model(model)
|
||||
}
|
||||
|
||||
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||
self.inner.seed_response_chain(thread_id, response_id);
|
||||
}
|
||||
|
||||
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||
self.inner.get_response_chain_id(thread_id)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -210,14 +210,6 @@ impl LlmProvider for RetryProvider {
|
||||
self.inner.set_model(model)
|
||||
}
|
||||
|
||||
fn seed_response_chain(&self, thread_id: &str, response_id: String) {
|
||||
self.inner.seed_response_chain(thread_id, response_id)
|
||||
}
|
||||
|
||||
fn get_response_chain_id(&self, thread_id: &str) -> Option<String> {
|
||||
self.inner.get_response_chain_id(thread_id)
|
||||
}
|
||||
|
||||
fn calculate_cost(&self, input_tokens: u32, output_tokens: u32) -> Decimal {
|
||||
self.inner.calculate_cost(input_tokens, output_tokens)
|
||||
}
|
||||
|
||||
@@ -445,7 +445,6 @@ where
|
||||
input_tokens: saturate_u32(response.usage.input_tokens),
|
||||
output_tokens: saturate_u32(response.usage.output_tokens),
|
||||
finish_reason: finish,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -511,7 +510,6 @@ where
|
||||
input_tokens: saturate_u32(response.usage.input_tokens),
|
||||
output_tokens: saturate_u32(response.usage.output_tokens),
|
||||
finish_reason: finish,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+19
-9
@@ -513,20 +513,30 @@ impl SessionManager {
|
||||
})? {
|
||||
value
|
||||
} else {
|
||||
tracing::warn!(
|
||||
"nearai.session_token missing; falling back to legacy nearai.session for backwards compatibility"
|
||||
);
|
||||
store
|
||||
// Try the legacy key. Only warn if it actually exists (real
|
||||
// backwards-compat migration). When neither key is present
|
||||
// (fresh install), just return the "No session in DB" error.
|
||||
let legacy = store
|
||||
.get_setting(&user_id, "nearai.session")
|
||||
.await
|
||||
.map_err(|e| LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: format!("DB query failed: {}", e),
|
||||
})?
|
||||
.ok_or(LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "No session in DB".to_string(),
|
||||
})?
|
||||
})?;
|
||||
match legacy {
|
||||
Some(value) => {
|
||||
tracing::warn!(
|
||||
"nearai.session_token missing; falling back to legacy nearai.session for backwards compatibility"
|
||||
);
|
||||
value
|
||||
}
|
||||
None => {
|
||||
return Err(LlmError::SessionRenewalFailed {
|
||||
provider: "nearai".to_string(),
|
||||
reason: "No session in DB".to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let session: SessionData =
|
||||
|
||||
+18
-31
@@ -3,7 +3,7 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::Parser;
|
||||
use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
use ironclaw::{
|
||||
agent::{Agent, AgentDeps, SessionManager},
|
||||
@@ -14,7 +14,7 @@ use ironclaw::{
|
||||
RegisteredEndpoint, SharedWasmChannel, WasmChannelLoader, WasmChannelRouter,
|
||||
WasmChannelRuntime, WasmChannelRuntimeConfig, create_wasm_channel_router,
|
||||
},
|
||||
web::log_layer::{LogBroadcaster, WebLogLayer},
|
||||
web::log_layer::LogBroadcaster,
|
||||
},
|
||||
cli::{
|
||||
Cli, Command, run_mcp_command, run_pairing_command, run_service_command,
|
||||
@@ -201,6 +201,9 @@ async fn main() -> anyhow::Result<()> {
|
||||
)
|
||||
.init();
|
||||
|
||||
let _ = dotenvy::dotenv();
|
||||
ironclaw::bootstrap::load_ironclaw_env();
|
||||
|
||||
return ironclaw::cli::run_doctor_command().await;
|
||||
}
|
||||
Some(Command::Status) => {
|
||||
@@ -210,6 +213,9 @@ async fn main() -> anyhow::Result<()> {
|
||||
)
|
||||
.init();
|
||||
|
||||
let _ = dotenvy::dotenv();
|
||||
ironclaw::bootstrap::load_ironclaw_env();
|
||||
|
||||
return run_status_command().await;
|
||||
}
|
||||
Some(Command::Worker {
|
||||
@@ -360,23 +366,14 @@ async fn main() -> anyhow::Result<()> {
|
||||
};
|
||||
let session = create_session_manager(session_config).await;
|
||||
|
||||
// Initialize tracing
|
||||
let env_filter = EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("ironclaw=info,tower_http=warn"));
|
||||
|
||||
// Create log broadcaster before tracing init so the WebLogLayer can capture all events.
|
||||
// This gets wired to the gateway's /api/logs/events SSE endpoint later.
|
||||
let log_broadcaster = Arc::new(LogBroadcaster::new());
|
||||
|
||||
tracing_subscriber::registry()
|
||||
.with(env_filter)
|
||||
.with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.with_target(false)
|
||||
.with_writer(ironclaw::tracing_fmt::TruncatingStderr::default()),
|
||||
)
|
||||
.with(WebLogLayer::new(Arc::clone(&log_broadcaster)))
|
||||
.init();
|
||||
// Initialize tracing with a reloadable EnvFilter so the gateway can switch
|
||||
// log levels (e.g. ironclaw=debug) at runtime without restarting.
|
||||
let log_level_handle =
|
||||
ironclaw::channels::web::log_layer::init_tracing(Arc::clone(&log_broadcaster));
|
||||
|
||||
// Create CLI channel
|
||||
let repl_channel = if let Some(ref msg) = cli.message {
|
||||
@@ -730,7 +727,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Initialize tool registry
|
||||
let tools = Arc::new(ToolRegistry::new());
|
||||
tools.register_builtin_tools();
|
||||
tracing::info!("Registered {} built-in tools", tools.count());
|
||||
|
||||
// Create embeddings provider if configured
|
||||
let embeddings: Option<Arc<dyn EmbeddingProvider>> = if config.embeddings.enabled {
|
||||
@@ -1024,11 +1020,12 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Set up orchestrator for sandboxed job execution
|
||||
// When allow_local_tools is false (default), the LLM uses create_job for FS/shell work.
|
||||
// When allow_local_tools is true, dev tools are also registered directly (current behavior).
|
||||
if config.agent.allow_local_tools {
|
||||
// register_builder_tool() already calls register_dev_tools() internally,
|
||||
// so only register them here when the builder didn't already do it.
|
||||
let builder_registered_dev_tools =
|
||||
config.builder.enabled && (config.agent.allow_local_tools || !config.sandbox.enabled);
|
||||
if config.agent.allow_local_tools && !builder_registered_dev_tools {
|
||||
tools.register_dev_tools();
|
||||
tracing::info!(
|
||||
"Local tools enabled (allow_local_tools=true), dev tools registered directly"
|
||||
);
|
||||
}
|
||||
|
||||
// Shared state for job events (used by both orchestrator and web gateway)
|
||||
@@ -1079,7 +1076,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
});
|
||||
|
||||
tracing::info!("Orchestrator API started on :50051, sandbox delegation enabled");
|
||||
if config.claude_code.enabled {
|
||||
tracing::info!(
|
||||
"Claude Code sandbox mode available (model: {}, max_turns: {})",
|
||||
@@ -1333,9 +1329,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Seed workspace with core identity files on first boot
|
||||
if let Some(ref ws) = workspace {
|
||||
match ws.seed_if_empty().await {
|
||||
Ok(count) if count > 0 => {
|
||||
tracing::info!("Workspace seeded with {} core files", count);
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to seed workspace: {}", e);
|
||||
@@ -1426,6 +1419,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
gw = gw.with_session_manager(Arc::clone(&session_manager));
|
||||
gw = gw.with_log_broadcaster(Arc::clone(&log_broadcaster));
|
||||
gw = gw.with_log_level_handle(Arc::clone(&log_level_handle));
|
||||
gw = gw.with_tool_registry(Arc::clone(&tools));
|
||||
if let Some(ref ext_mgr) = extension_manager {
|
||||
gw = gw.with_extension_manager(Arc::clone(ext_mgr));
|
||||
@@ -1464,11 +1458,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
gw.auth_token()
|
||||
));
|
||||
|
||||
tracing::info!(
|
||||
"Web gateway enabled on {}:{}",
|
||||
gw_config.host,
|
||||
gw_config.port
|
||||
);
|
||||
tracing::info!("Web UI: http://{}:{}/", gw_config.host, gw_config.port);
|
||||
|
||||
channel_names.push("gateway".to_string());
|
||||
@@ -1511,8 +1500,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
Some(session_manager),
|
||||
);
|
||||
|
||||
tracing::info!("Agent initialized, starting main loop...");
|
||||
|
||||
// Print boot screen for interactive CLI mode (not single-message mode).
|
||||
if config.channels.cli.enabled && cli.message.is_none() {
|
||||
let boot_info = ironclaw::boot_screen::BootInfo {
|
||||
|
||||
+15
-2
@@ -159,11 +159,13 @@ impl Default for Policy {
|
||||
PolicyAction::Warn,
|
||||
));
|
||||
|
||||
// Block shell command injection patterns
|
||||
// Block shell command injection patterns.
|
||||
// Only match actual dangerous command sequences, NOT backticked content
|
||||
// (backticks are standard markdown code formatting, not shell injection).
|
||||
policy.add_rule(PolicyRule::new(
|
||||
"shell_injection",
|
||||
"Potential shell command injection",
|
||||
r"(?i)(;\s*rm\s+-rf|;\s*curl\s+.*\|\s*sh|`.*`)",
|
||||
r"(?i)(;\s*rm\s+-rf|;\s*curl\s+.*\|\s*sh)",
|
||||
Severity::Critical,
|
||||
PolicyAction::Block,
|
||||
));
|
||||
@@ -233,6 +235,17 @@ mod tests {
|
||||
assert!(violations.iter().any(|r| r.action == PolicyAction::Warn));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_backticked_code_is_not_blocked() {
|
||||
let policy = Policy::default();
|
||||
// Markdown code snippets should never be blocked
|
||||
assert!(!policy.is_blocked("Use `print('hello')` to debug"));
|
||||
assert!(!policy.is_blocked("Run `pytest tests/` to check"));
|
||||
assert!(!policy.is_blocked("The error is in `foo.bar.baz`"));
|
||||
// Multi-backtick code fences should also pass
|
||||
assert!(!policy.is_blocked("```python\ndef foo():\n pass\n```"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_severity_ordering() {
|
||||
assert!(Severity::Critical > Severity::High);
|
||||
|
||||
@@ -302,6 +302,14 @@ pub struct AgentSettings {
|
||||
/// longer than this are pruned from memory.
|
||||
#[serde(default = "default_session_idle_timeout")]
|
||||
pub session_idle_timeout_secs: u64,
|
||||
|
||||
/// Maximum tool-call iterations per agentic loop invocation (default: 50).
|
||||
#[serde(default = "default_max_tool_iterations")]
|
||||
pub max_tool_iterations: usize,
|
||||
|
||||
/// When true, skip tool approval checks entirely. For benchmarks/CI.
|
||||
#[serde(default)]
|
||||
pub auto_approve_tools: bool,
|
||||
}
|
||||
|
||||
fn default_agent_name() -> String {
|
||||
@@ -332,6 +340,10 @@ fn default_max_repair_attempts() -> u32 {
|
||||
3
|
||||
}
|
||||
|
||||
fn default_max_tool_iterations() -> usize {
|
||||
50
|
||||
}
|
||||
|
||||
fn default_true() -> bool {
|
||||
true
|
||||
}
|
||||
@@ -347,6 +359,8 @@ impl Default for AgentSettings {
|
||||
repair_check_interval_secs: default_repair_interval(),
|
||||
max_repair_attempts: default_max_repair_attempts(),
|
||||
session_idle_timeout_secs: default_session_idle_timeout(),
|
||||
max_tool_iterations: default_max_tool_iterations(),
|
||||
auto_approve_tools: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+54
-1
@@ -826,6 +826,12 @@ impl SetupWizard {
|
||||
|
||||
self.session_manager = Some(session);
|
||||
|
||||
// Persist session token to the database so the runtime can load it
|
||||
// via `attach_store()` → `load_session_from_db()` without the
|
||||
// backwards-compat fallback. The session manager saved to disk but
|
||||
// doesn't have a DB store attached during onboarding.
|
||||
self.persist_session_to_db().await;
|
||||
|
||||
// If the user chose the API key path, NEARAI_API_KEY is now set
|
||||
// in the environment. Persist it to the encrypted secrets store
|
||||
// so inject_llm_keys_from_secrets() can load it on future runs.
|
||||
@@ -1160,7 +1166,6 @@ impl SetupWizard {
|
||||
base_url,
|
||||
auth_base_url,
|
||||
session_path: crate::llm::session::default_session_path(),
|
||||
api_mode: crate::config::NearAiApiMode::Responses,
|
||||
api_key: None,
|
||||
fallback_model: None,
|
||||
max_retries: 3,
|
||||
@@ -1861,6 +1866,54 @@ impl SetupWizard {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Persist the NEAR AI session token to the database.
|
||||
///
|
||||
/// The session manager writes to disk during `ensure_authenticated()` but
|
||||
/// doesn't have a DB store attached during onboarding. This reads the
|
||||
/// session file from disk and stores it under the `nearai.session_token`
|
||||
/// key so the runtime's `attach_store()` finds it without fallback.
|
||||
///
|
||||
/// Best-effort: silently ignores errors (no DB connection yet, no
|
||||
/// session file, etc.).
|
||||
async fn persist_session_to_db(&self) {
|
||||
let session_path = crate::llm::session::default_session_path();
|
||||
let data = match std::fs::read_to_string(&session_path) {
|
||||
Ok(d) if !d.trim().is_empty() => d,
|
||||
_ => return,
|
||||
};
|
||||
let value: serde_json::Value = match serde_json::from_str(&data) {
|
||||
Ok(v) => v,
|
||||
Err(_) => return,
|
||||
};
|
||||
|
||||
#[cfg(feature = "postgres")]
|
||||
if let Some(ref pool) = self.db_pool {
|
||||
let store = crate::history::Store::from_pool(pool.clone());
|
||||
if let Err(e) = store
|
||||
.set_setting("default", "nearai.session_token", &value)
|
||||
.await
|
||||
{
|
||||
tracing::debug!("Could not persist session token to postgres: {}", e);
|
||||
} else {
|
||||
tracing::debug!("Session token persisted to database");
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
if let Some(ref backend) = self.db_backend {
|
||||
use crate::db::SettingsStore as _;
|
||||
if let Err(e) = backend
|
||||
.set_setting("default", "nearai.session_token", &value)
|
||||
.await
|
||||
{
|
||||
tracing::debug!("Could not persist session token to libsql: {}", e);
|
||||
} else {
|
||||
tracing::debug!("Session token persisted to database");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Persist settings to DB and bootstrap .env after each step.
|
||||
///
|
||||
/// Silently ignores errors (e.g., DB not connected yet before step 1
|
||||
|
||||
@@ -197,7 +197,7 @@ impl SkillRegistry {
|
||||
let source = make_source(path.clone());
|
||||
match self.load_skill_md(&skill_md, trust, source).await {
|
||||
Ok((name, skill)) => {
|
||||
tracing::info!("Loaded skill: {}", name);
|
||||
tracing::debug!("Loaded skill: {}", name);
|
||||
results.push((name, skill));
|
||||
}
|
||||
Err(e) => {
|
||||
|
||||
@@ -168,7 +168,6 @@ impl LlmProvider for StubLlm {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -186,7 +185,6 @@ impl LlmProvider for StubLlm {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -227,7 +227,6 @@ impl WorkerHttpClient {
|
||||
input_tokens: proxy_resp.input_tokens,
|
||||
output_tokens: proxy_resp.output_tokens,
|
||||
finish_reason: parse_finish_reason(&proxy_resp.finish_reason),
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -255,7 +254,6 @@ impl WorkerHttpClient {
|
||||
input_tokens: proxy_resp.input_tokens,
|
||||
output_tokens: proxy_resp.output_tokens,
|
||||
finish_reason: parse_finish_reason(&proxy_resp.finish_reason),
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@ use ironclaw::{
|
||||
config::Config,
|
||||
history::Store,
|
||||
llm::{SessionConfig, create_llm_provider, create_session_manager},
|
||||
safety::SafetyLayer,
|
||||
workspace::Workspace,
|
||||
};
|
||||
|
||||
@@ -96,7 +97,8 @@ async fn test_heartbeat_end_to_end() {
|
||||
|
||||
let hb_config = ironclaw::agent::HeartbeatConfig::default();
|
||||
let hygiene_config = ironclaw::workspace::hygiene::HygieneConfig::default();
|
||||
let runner = HeartbeatRunner::new(hb_config, hygiene_config, workspace, llm);
|
||||
let safety = Arc::new(SafetyLayer::new(&config.safety));
|
||||
let runner = HeartbeatRunner::new(hb_config, hygiene_config, workspace, llm, safety);
|
||||
|
||||
let result = runner.check_heartbeat().await;
|
||||
|
||||
|
||||
@@ -71,7 +71,6 @@ impl LlmProvider for MockLlmProvider {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -97,7 +96,6 @@ impl LlmProvider for MockLlmProvider {
|
||||
input_tokens: 15,
|
||||
output_tokens: 8,
|
||||
finish_reason: FinishReason::ToolUse,
|
||||
response_id: None,
|
||||
})
|
||||
} else {
|
||||
Ok(ToolCompletionResponse {
|
||||
@@ -106,7 +104,6 @@ impl LlmProvider for MockLlmProvider {
|
||||
input_tokens: 10,
|
||||
output_tokens: 4,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -145,7 +142,6 @@ impl LlmProvider for FixedModelProvider {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -159,7 +155,6 @@ impl LlmProvider for FixedModelProvider {
|
||||
input_tokens: 10,
|
||||
output_tokens: 5,
|
||||
finish_reason: FinishReason::Stop,
|
||||
response_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -190,6 +185,7 @@ async fn start_test_server_with_provider(
|
||||
workspace: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
@@ -674,6 +670,7 @@ async fn test_no_llm_provider_returns_503() {
|
||||
workspace: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
|
||||
@@ -43,6 +43,7 @@ async fn start_test_server() -> (
|
||||
workspace: None,
|
||||
session_manager: None,
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: None,
|
||||
tool_registry: None,
|
||||
store: None,
|
||||
|
||||
Reference in New Issue
Block a user