diff --git a/Cargo.lock b/Cargo.lock index 533ba11..699db4e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -78,6 +78,27 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "dirs" +version = "5.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c45a9d03d6676652bcb5e724c7e988de1acad23a711b5217ab9cbecbec2225" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520f05a5cbd335fae5a99ff7a6ab8627577660ee5cfd6a94a6a929b52ff0321c" +dependencies = [ + "libc", + "option-ext", + "redox_users", + "windows-sys 0.48.0", +] + [[package]] name = "dyn-clone" version = "1.0.20" @@ -91,7 +112,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -106,6 +127,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -200,6 +227,17 @@ dependencies = [ "slab", ] +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "getrandom" version = "0.3.4" @@ -241,13 +279,16 @@ name = "harness-core" version = "0.1.0" dependencies = [ "async-trait", + "dirs", "futures", "globset", + "jsonc-parser", "rusqlite", "schemars", "serde", "serde_json", - "thiserror", + "tempfile", + "thiserror 2.0.18", "tokio", "tokio-util", "tracing", @@ -324,12 +365,30 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonc-parser" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d6d80e6d70e7911a29f3cf3f44f452df85d06f73572b494ca99a2cad3fcf8f4" +dependencies = [ + "serde_json", +] + [[package]] name = "libc" version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libredox" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" +dependencies = [ + "libc", +] + [[package]] name = "libsqlite3-sys" version = "0.30.1" @@ -341,6 +400,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "lock_api" version = "0.4.14" @@ -370,7 +435,7 @@ checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", "wasi", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -379,6 +444,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "option-ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" + [[package]] name = "parking_lot" version = "0.12.5" @@ -473,7 +544,7 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" dependencies = [ - "getrandom", + "getrandom 0.3.4", ] [[package]] @@ -485,6 +556,17 @@ dependencies = [ "bitflags", ] +[[package]] +name = "redox_users" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" +dependencies = [ + "getrandom 0.2.17", + "libredox", + "thiserror 1.0.69", +] + [[package]] name = "regex-automata" version = "0.4.14" @@ -516,6 +598,19 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -641,7 +736,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -655,13 +750,46 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.3.4", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -689,7 +817,7 @@ dependencies = [ "signal-hook-registry", "socket2", "tokio-macros", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -852,6 +980,15 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-sys" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-sys" version = "0.61.2" @@ -861,6 +998,63 @@ dependencies = [ "windows-link", ] +[[package]] +name = "windows-targets" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" + +[[package]] +name = "windows_i686_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" + +[[package]] +name = "windows_i686_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/Cargo.toml b/Cargo.toml index dff405c..53a5975 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,6 +57,9 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } tracing-appender = "0.2" dirs = "5" shell-words = "1" +tempfile = "3" +insta = "1" +wiremock = "0.6" [workspace.lints.rust] unsafe_code = "deny" diff --git a/crates/harness-core/Cargo.toml b/crates/harness-core/Cargo.toml index e9d9675..6f8bb52 100644 --- a/crates/harness-core/Cargo.toml +++ b/crates/harness-core/Cargo.toml @@ -14,12 +14,15 @@ serde_json = { workspace = true } schemars = { workspace = true } rusqlite = { workspace = true } globset = { workspace = true } +jsonc-parser = { workspace = true } +dirs = { workspace = true } ulid = { workspace = true } thiserror = { workspace = true } tracing = { workspace = true } [dev-dependencies] tokio = { workspace = true, features = ["test-util", "macros"] } +tempfile = { workspace = true } [lints] workspace = true diff --git a/crates/harness-core/src/config/load.rs b/crates/harness-core/src/config/load.rs new file mode 100644 index 0000000..d9988ed --- /dev/null +++ b/crates/harness-core/src/config/load.rs @@ -0,0 +1,269 @@ +use std::path::{Path, PathBuf}; + +use serde_json::Value; + +use super::schema::Config; + +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("reading {path}: {source}")] + Io { + path: PathBuf, + #[source] + source: std::io::Error, + }, + #[error("parsing {path}: {source}")] + Jsonc { + path: PathBuf, + #[source] + source: jsonc_parser::errors::ParseError, + }, + #[error("deserializing merged config: {0}")] + Serde(#[from] serde_json::Error), +} + +/// Env vars consulted directly (not via `{env:...}` interpolation) — highest precedence. +const ENV_MODEL: &str = "AI_HARNESS_MODEL"; +const ENV_CONFIG_PATH: &str = "AI_HARNESS_CONFIG"; +/// `(provider id, env var)` pairs used to fill in `providers..api_key` when unset. +const PROVIDER_ENV_KEYS: &[(&str, &str)] = &[ + ("anthropic", "ANTHROPIC_API_KEY"), + ("openai", "OPENAI_API_KEY"), +]; + +fn read_jsonc(path: &Path) -> Result { + let text = std::fs::read_to_string(path).map_err(|source| ConfigError::Io { + path: path.to_path_buf(), + source, + })?; + let value = + jsonc_parser::parse_to_serde_value(&text, &Default::default()).map_err(|source| { + ConfigError::Jsonc { + path: path.to_path_buf(), + source, + } + })?; + Ok(value.unwrap_or_else(|| Value::Object(Default::default()))) +} + +/// Objects deep-merge (overlay wins per key); anything else (including arrays) is replaced. +fn merge(base: &mut Value, overlay: Value) { + match (base, overlay) { + (Value::Object(base_map), Value::Object(overlay_map)) => { + for (key, value) in overlay_map { + match base_map.get_mut(&key) { + Some(existing) => merge(existing, value), + None => { + base_map.insert(key, value); + } + } + } + } + (base_slot, overlay_value) => { + *base_slot = overlay_value; + } + } +} + +/// Exact-string placeholders only (opencode-style): `"{env:VAR}"` and `"{file:path}"`. +/// `path` in `{file:...}` resolves relative to `base_dir`. +fn interpolate(value: &mut Value, base_dir: &Path) { + match value { + Value::String(s) => { + if let Some(var) = s.strip_prefix("{env:").and_then(|s| s.strip_suffix('}')) { + *s = std::env::var(var).unwrap_or_default(); + } else if let Some(rel) = s.strip_prefix("{file:").and_then(|s| s.strip_suffix('}')) { + *s = std::fs::read_to_string(base_dir.join(rel)).unwrap_or_default(); + } + } + Value::Array(items) => { + for item in items { + interpolate(item, base_dir); + } + } + Value::Object(map) => { + for (_, v) in map.iter_mut() { + interpolate(v, base_dir); + } + } + _ => {} + } +} + +fn global_config_path() -> Option { + dirs::config_dir().map(|d| d.join("ai-harness").join("config.json")) +} + +/// Ancestors of `cwd` that have a `.harness/config.json`, root-most first (lowest +/// precedence among the project chain; `cwd`'s own file, if any, applies last). +fn project_chain(cwd: &Path) -> Vec { + let mut found: Vec = cwd + .ancestors() + .map(|dir| dir.join(".harness").join("config.json")) + .filter(|p| p.is_file()) + .collect(); + found.reverse(); + found +} + +pub fn load(cwd: &Path) -> Result { + let mut merged = serde_json::to_value(Config::default())?; + + if let Some(path) = global_config_path() { + if path.is_file() { + merge(&mut merged, read_jsonc(&path)?); + } + } + + for path in project_chain(cwd) { + merge(&mut merged, read_jsonc(&path)?); + } + + if let Ok(path) = std::env::var(ENV_CONFIG_PATH) { + let path = PathBuf::from(path); + if path.is_file() { + merge(&mut merged, read_jsonc(&path)?); + } + } + + interpolate(&mut merged, cwd); + + let mut config: Config = serde_json::from_value(merged)?; + + if let Ok(model) = std::env::var(ENV_MODEL) { + config.model = Some(model); + } + for (provider_id, env_var) in PROVIDER_ENV_KEYS { + if let Ok(key) = std::env::var(env_var) { + let entry = config + .providers + .entry((*provider_id).to_string()) + .or_default(); + if entry.api_key.is_none() { + entry.api_key = Some(key); + } + } + } + + Ok(config) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Mutex; + + // Env var mutation races across tests in the same process; serialize them. + static ENV_LOCK: Mutex<()> = Mutex::new(()); + + #[test] + fn merge_deep_merges_objects_and_replaces_arrays() { + let mut base = serde_json::json!({ + "a": {"x": 1, "y": 2}, + "list": [1, 2, 3], + }); + let overlay = serde_json::json!({ + "a": {"y": 20, "z": 3}, + "list": [9], + }); + merge(&mut base, overlay); + assert_eq!( + base, + serde_json::json!({"a": {"x": 1, "y": 20, "z": 3}, "list": [9]}) + ); + } + + #[test] + fn interpolate_replaces_env_placeholder() { + let _lock = ENV_LOCK.lock().unwrap(); + std::env::set_var("HARNESS_TEST_VAR", "secret"); + let mut value = serde_json::json!({"api_key": "{env:HARNESS_TEST_VAR}"}); + interpolate(&mut value, Path::new(".")); + assert_eq!(value, serde_json::json!({"api_key": "secret"})); + std::env::remove_var("HARNESS_TEST_VAR"); + } + + #[test] + fn interpolate_replaces_file_placeholder() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("prompt.md"), "You are a reviewer.").unwrap(); + let mut value = serde_json::json!({"prompt": "{file:prompt.md}"}); + interpolate(&mut value, dir.path()); + assert_eq!(value, serde_json::json!({"prompt": "You are a reviewer."})); + } + + #[test] + fn load_applies_bundled_defaults_when_nothing_else_present() { + let _lock = ENV_LOCK.lock().unwrap(); + std::env::remove_var(ENV_MODEL); + std::env::remove_var(ENV_CONFIG_PATH); + let dir = tempfile::tempdir().unwrap(); + let config = load(dir.path()).unwrap(); + assert_eq!(config.orchestration.depth_limit, 3); + assert!(config.model.is_none()); + } + + #[test] + fn load_merges_project_chain_root_most_first() { + let _lock = ENV_LOCK.lock().unwrap(); + std::env::remove_var(ENV_MODEL); + std::env::remove_var(ENV_CONFIG_PATH); + + let root = tempfile::tempdir().unwrap(); + let child = root.path().join("child"); + std::fs::create_dir_all(child.join(".harness")).unwrap(); + std::fs::create_dir_all(root.path().join(".harness")).unwrap(); + + std::fs::write( + root.path().join(".harness/config.json"), + r#"{ "model": "root-model", "orchestration": { "depth_limit": 5 } }"#, + ) + .unwrap(); + std::fs::write( + child.join(".harness/config.json"), + r#"{ "model": "child-model" }"#, + ) + .unwrap(); + + let config = load(&child).unwrap(); + // child's config.json overrides root's `model` (applied later in the chain)... + assert_eq!(config.model.as_deref(), Some("child-model")); + // ...but root's untouched `orchestration.depth_limit` survives the merge. + assert_eq!(config.orchestration.depth_limit, 5); + } + + #[test] + fn env_model_override_wins_over_everything() { + let _lock = ENV_LOCK.lock().unwrap(); + std::env::remove_var(ENV_CONFIG_PATH); + std::env::set_var(ENV_MODEL, "env-model"); + let dir = tempfile::tempdir().unwrap(); + std::fs::create_dir_all(dir.path().join(".harness")).unwrap(); + std::fs::write( + dir.path().join(".harness/config.json"), + r#"{ "model": "project-model" }"#, + ) + .unwrap(); + let config = load(dir.path()).unwrap(); + assert_eq!(config.model.as_deref(), Some("env-model")); + std::env::remove_var(ENV_MODEL); + } + + #[test] + fn provider_env_key_fills_in_missing_api_key_only() { + let _lock = ENV_LOCK.lock().unwrap(); + std::env::remove_var(ENV_MODEL); + std::env::remove_var(ENV_CONFIG_PATH); + std::env::set_var("ANTHROPIC_API_KEY", "from-env"); + let dir = tempfile::tempdir().unwrap(); + let config = load(dir.path()).unwrap(); + assert_eq!( + config + .providers + .get("anthropic") + .and_then(|p| p.api_key.clone()), + Some("from-env".to_string()) + ); + std::env::remove_var("ANTHROPIC_API_KEY"); + } +} diff --git a/crates/harness-core/src/config/mod.rs b/crates/harness-core/src/config/mod.rs new file mode 100644 index 0000000..825e4e6 --- /dev/null +++ b/crates/harness-core/src/config/mod.rs @@ -0,0 +1,8 @@ +pub mod load; +pub mod schema; + +pub use load::{load, ConfigError}; +pub use schema::{ + AgentPatch, Config, LspServerConfig, McpServerConfig, OrchestrationConfig, ProviderConfig, + TuiConfig, +}; diff --git a/crates/harness-core/src/config/schema.rs b/crates/harness-core/src/config/schema.rs new file mode 100644 index 0000000..130a1e4 --- /dev/null +++ b/crates/harness-core/src/config/schema.rs @@ -0,0 +1,86 @@ +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; + +use crate::permission::Rule; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct ProviderConfig { + pub enabled: Option, + pub api_key: Option, + pub base_url: Option, +} + +/// Patch semantics over a loaded `AgentDef` (M4): only the fields set here override the +/// bundled/markdown definition. An unknown agent name with at least `model` or `prompt` +/// set creates a custom agent. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct AgentPatch { + pub disable: Option, + pub mode: Option, + pub model: Option, + pub temperature: Option, + pub prompt: Option, + pub tools: Option>, + pub permission: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct OrchestrationConfig { + pub depth_limit: u32, + pub job_board: bool, + pub max_reusable_per_agent: u32, + pub reminders: Option>, +} + +impl Default for OrchestrationConfig { + fn default() -> Self { + Self { + depth_limit: 3, + job_board: true, + max_reusable_per_agent: 2, + reminders: None, + } + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct McpServerConfig { + pub command: String, + pub args: Vec, + pub env: HashMap, + pub enabled: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct LspServerConfig { + pub command: String, + pub args: Vec, + pub extensions: Vec, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct TuiConfig { + pub theme: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct Config { + pub model: Option, + pub small_model: Option, + pub providers: HashMap, + pub agents: HashMap, + pub permissions: Vec, + pub orchestration: OrchestrationConfig, + pub mcp: HashMap, + pub lsp: HashMap, + pub tui: TuiConfig, + pub instructions: Vec, +} diff --git a/crates/harness-core/src/engine/doomloop.rs b/crates/harness-core/src/engine/doomloop.rs new file mode 100644 index 0000000..fd37557 --- /dev/null +++ b/crates/harness-core/src/engine/doomloop.rs @@ -0,0 +1,75 @@ +use std::collections::VecDeque; +use std::hash::{Hash, Hasher}; + +/// opencode's `DOOM_LOOP_THRESHOLD`: force a permission ask once the model repeats the +/// exact same tool call this many times in a row. +const DOOM_LOOP_THRESHOLD: usize = 3; + +fn fingerprint(tool_name: &str, input: &serde_json::Value) -> (String, u64) { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + input.to_string().hash(&mut hasher); + (tool_name.to_string(), hasher.finish()) +} + +/// Tracks the last `DOOM_LOOP_THRESHOLD` `(tool_name, hash(input))` pairs for one run. +pub struct DoomLoopGuard { + history: VecDeque<(String, u64)>, +} + +impl DoomLoopGuard { + pub fn new() -> Self { + Self { + history: VecDeque::with_capacity(DOOM_LOOP_THRESHOLD), + } + } + + /// Records this call and returns `true` if it's the `DOOM_LOOP_THRESHOLD`-th identical + /// call in a row (the caller should force a permission ask regardless of the ruleset). + pub fn record_and_check(&mut self, tool_name: &str, input: &serde_json::Value) -> bool { + let fp = fingerprint(tool_name, input); + self.history.push_back(fp.clone()); + if self.history.len() > DOOM_LOOP_THRESHOLD { + self.history.pop_front(); + } + self.history.len() == DOOM_LOOP_THRESHOLD && self.history.iter().all(|h| *h == fp) + } +} + +impl Default for DoomLoopGuard { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn triggers_after_three_identical_calls() { + let mut guard = DoomLoopGuard::new(); + let input = serde_json::json!({"command": "ls"}); + assert!(!guard.record_and_check("bash", &input)); + assert!(!guard.record_and_check("bash", &input)); + assert!(guard.record_and_check("bash", &input)); + } + + #[test] + fn different_input_resets_the_streak() { + let mut guard = DoomLoopGuard::new(); + assert!(!guard.record_and_check("bash", &serde_json::json!({"command": "ls"}))); + assert!(!guard.record_and_check("bash", &serde_json::json!({"command": "pwd"}))); + assert!(!guard.record_and_check("bash", &serde_json::json!({"command": "ls"}))); + } + + #[test] + fn keeps_triggering_while_the_streak_continues() { + let mut guard = DoomLoopGuard::new(); + let input = serde_json::json!({"command": "ls"}); + for _ in 0..2 { + guard.record_and_check("bash", &input); + } + assert!(guard.record_and_check("bash", &input)); + assert!(guard.record_and_check("bash", &input)); + } +} diff --git a/crates/harness-core/src/engine/loop.rs b/crates/harness-core/src/engine/loop.rs new file mode 100644 index 0000000..e495bba --- /dev/null +++ b/crates/harness-core/src/engine/loop.rs @@ -0,0 +1,529 @@ +use std::sync::Arc; + +use crate::event::RunOutcome; +use crate::llm::{ + FinishReason, Initiator, LlmRequest, Provider, ProviderError, Role as WireRole, ToolSchema, + WireContent, WireMessage, +}; +use crate::store::Store; +use crate::types::{Message, Part, PartBody, Role, SessionId, ToolState}; + +use super::doomloop::DoomLoopGuard; +use super::processor::{process_step, StepContext, StepResult}; +use super::retry; +use super::system; + +pub struct RunConfig { + pub agent_name: String, + pub agent_prompt: String, + pub model: crate::types::ModelRef, + pub temperature: Option, + pub max_steps: u32, + pub instructions: Vec, +} + +/// opencode's exit condition: keep stepping while the last assistant turn asked for more +/// tool calls; stop once it produced a final answer (or a fresh user message is waiting). +async fn should_continue(store: &Store, session_id: &SessionId) -> Result { + let messages = store + .messages(session_id.clone()) + .await + .map_err(|e| ProviderError::Decode(e.to_string()))?; + Ok(match messages.last() { + None => false, + Some(msg) if msg.role == Role::User => true, + Some(msg) => matches!(msg.finished, Some(FinishReason::ToolCalls)), + }) +} + +fn convert_message(message: &Message, parts: &[Part]) -> Vec { + match message.role { + Role::User => { + let content: Vec = parts + .iter() + .filter_map(|p| match &p.body { + PartBody::Text { text, .. } => Some(WireContent::Text { text: text.clone() }), + _ => None, + }) + .collect(); + if content.is_empty() { + vec![] + } else { + vec![WireMessage { + role: WireRole::User, + content, + }] + } + } + Role::Assistant => { + let mut assistant_content = Vec::new(); + let mut tool_results = Vec::new(); + for part in parts { + match &part.body { + PartBody::Text { text, .. } => { + assistant_content.push(WireContent::Text { text: text.clone() }) + } + PartBody::Tool { + call_id, + name, + state, + } => match state { + ToolState::Completed { input, output, .. } => { + assistant_content.push(WireContent::ToolCall { + call_id: call_id.clone(), + name: name.clone(), + input: input.clone(), + }); + tool_results.push(WireContent::ToolResult { + call_id: call_id.clone(), + output: output.clone(), + is_error: false, + }); + } + ToolState::Error { input, error } => { + assistant_content.push(WireContent::ToolCall { + call_id: call_id.clone(), + name: name.clone(), + input: input.clone(), + }); + tool_results.push(WireContent::ToolResult { + call_id: call_id.clone(), + output: error.clone(), + is_error: true, + }); + } + ToolState::Running { input, .. } => { + assistant_content.push(WireContent::ToolCall { + call_id: call_id.clone(), + name: name.clone(), + input: input.clone(), + }); + } + ToolState::Pending { .. } => {} + }, + _ => {} + } + } + let mut out = Vec::new(); + if !assistant_content.is_empty() { + out.push(WireMessage { + role: WireRole::Assistant, + content: assistant_content, + }); + } + if !tool_results.is_empty() { + out.push(WireMessage { + role: WireRole::Tool, + content: tool_results, + }); + } + out + } + } +} + +/// The outer loop: one call per session turn until the model stops asking for tool calls. +/// `now_fn` supplies `created_at`/`StepContext::now` stamps (kept out of the loop body so +/// tests can drive deterministic timestamps). +pub async fn run_session( + provider: Arc, + mut ctx: StepContext, + run_config: &RunConfig, + now_fn: impl Fn() -> i64, +) -> RunOutcome { + let mut doomloop = DoomLoopGuard::new(); + let mut steps = 0u32; + + loop { + match should_continue(&ctx.store, &ctx.session_id).await { + Ok(true) => {} + Ok(false) => return RunOutcome::Stopped, + Err(e) => { + return RunOutcome::Errored { + message: e.to_string(), + } + } + } + if steps >= run_config.max_steps { + return RunOutcome::Stopped; + } + steps += 1; + ctx.now = now_fn(); + + let messages = match ctx.store.messages(ctx.session_id.clone()).await { + Ok(m) => m, + Err(e) => { + return RunOutcome::Errored { + message: e.to_string(), + } + } + }; + let mut wire_messages = Vec::new(); + for message in &messages { + let parts = match ctx.store.parts(message.id.clone()).await { + Ok(p) => p, + Err(e) => { + return RunOutcome::Errored { + message: e.to_string(), + } + } + }; + wire_messages.extend(convert_message(message, &parts)); + } + + let system_blocks = system::assemble( + system::env_header(&ctx.cwd), + &run_config.agent_prompt, + &run_config.instructions, + ); + let tools: Vec = ctx + .tools + .all() + .iter() + .map(|t| ToolSchema { + name: t.name().to_string(), + description: t.description().to_string(), + parameters: t.parameters(), + }) + .collect(); + + let req = LlmRequest { + model: run_config.model.model_id.clone(), + system: system_blocks, + messages: wire_messages, + tools, + temperature: run_config.temperature, + max_tokens: None, + reasoning: None, + initiator: Initiator::User, + }; + + let stream_result = retry::with_retry(&ctx.cancel, || { + let provider = provider.clone(); + let req = req.clone(); + let cancel = ctx.cancel.child_token(); + async move { provider.stream(req, cancel).await } + }) + .await; + + let stream = match stream_result { + Ok(s) => s, + Err(ProviderError::Cancelled) => return RunOutcome::Aborted, + Err(e) => { + return RunOutcome::Errored { + message: e.to_string(), + } + } + }; + + let step = process_step( + stream, + &ctx, + run_config.model.clone(), + &run_config.agent_name, + &mut doomloop, + ) + .await; + + match step { + Ok(outcome) if outcome.aborted => return RunOutcome::Aborted, + Ok(outcome) => match outcome.result { + StepResult::Continue => continue, + StepResult::Stop => return RunOutcome::Stopped, + StepResult::Compact => return RunOutcome::Stopped, // stub until M6 + }, + Err(step_err) if matches!(step_err.source, ProviderError::Cancelled) => { + return RunOutcome::Aborted; + } + Err(step_err) => { + return RunOutcome::Errored { + message: step_err.source.to_string(), + }; + } + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::VecDeque; + use std::sync::Mutex as StdMutex; + + use async_trait::async_trait; + use tokio_util::sync::CancellationToken; + + use super::*; + use crate::event::EventBus; + use crate::llm::{LlmEvent, LlmEventStream}; + use crate::permission::{spawn_auto_approve, PermissionService}; + use crate::store::Store; + use crate::tool::{Tool, ToolCtx, ToolError, ToolOutput, ToolRegistry}; + use crate::types::{ModelInfo, ModelRef, Session, TokenUsage}; + + struct MockProvider { + steps: StdMutex>>>, + } + + impl MockProvider { + fn scripted(steps: Vec>>) -> Self { + Self { + steps: StdMutex::new(steps.into_iter().collect()), + } + } + } + + #[async_trait] + impl Provider for MockProvider { + fn id(&self) -> &str { + "mock" + } + async fn list_models(&self) -> Result, ProviderError> { + Ok(vec![]) + } + async fn stream( + &self, + _req: LlmRequest, + _cancel: CancellationToken, + ) -> Result { + let events = self + .steps + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| ProviderError::Decode("no more scripted steps".into()))?; + Ok(Box::pin(futures::stream::iter(events))) + } + } + + struct MockReadTool; + + #[async_trait] + impl Tool for MockReadTool { + fn name(&self) -> &str { + "read" + } + fn description(&self) -> &str { + "reads a file" + } + fn parameters(&self) -> serde_json::Value { + serde_json::json!({"type": "object", "properties": {"file_path": {"type": "string"}}}) + } + async fn execute( + &self, + _input: serde_json::Value, + _ctx: ToolCtx, + ) -> Result { + Ok(ToolOutput::new("read foo.txt", "mock file content")) + } + } + + fn usage(input: u64, output: u64) -> TokenUsage { + TokenUsage { + input, + output, + ..Default::default() + } + } + + async fn make_ctx( + store: Store, + bus: EventBus, + session_id: SessionId, + cwd: std::path::PathBuf, + ) -> StepContext { + let mut tools = ToolRegistry::new(); + tools.register(Arc::new(MockReadTool)); + let permissions = Arc::new(PermissionService::new(bus.clone())); + spawn_auto_approve(bus.clone(), permissions.clone()); + StepContext { + store, + bus, + tools, + permissions, + static_rules: Vec::new(), + extra_rules: Arc::new(std::sync::Mutex::new(Vec::new())), + session_id, + cwd: cwd.clone(), + data_dir: cwd.join("tool-output"), + cancel: CancellationToken::new(), + now: 1, + } + } + + #[tokio::test] + async fn text_then_tool_call_then_final_text_persists_correctly() { + let store = Store::open_in_memory().unwrap(); + let bus = EventBus::new(); + let model = ModelRef::new("mock", "mock-model"); + let session = Session::new_root("orchestrator", model.clone(), 1); + let session_id = session.id.clone(); + store.upsert_session(session).await.unwrap(); + + let user_message = Message::new_user(session_id.clone(), 1); + store.upsert_message(user_message.clone()).await.unwrap(); + store + .upsert_part(Part { + id: crate::types::PartId::new(), + message_id: user_message.id.clone(), + session_id: session_id.clone(), + idx: 0, + body: PartBody::Text { + text: "please read foo.txt".into(), + synthetic: false, + }, + }) + .await + .unwrap(); + + let provider = MockProvider::scripted(vec![ + vec![ + Ok(LlmEvent::TextStart { id: "t1".into() }), + Ok(LlmEvent::TextDelta { + id: "t1".into(), + text: "Let me check the file.".into(), + }), + Ok(LlmEvent::TextEnd { id: "t1".into() }), + Ok(LlmEvent::ToolCall { + call_id: "call_1".into(), + name: "read".into(), + input: serde_json::json!({"file_path": "foo.txt"}), + }), + Ok(LlmEvent::Finish { + reason: FinishReason::ToolCalls, + usage: usage(10, 5), + }), + ], + vec![ + Ok(LlmEvent::TextStart { id: "t2".into() }), + Ok(LlmEvent::TextDelta { + id: "t2".into(), + text: "The file contains: mock file content".into(), + }), + Ok(LlmEvent::TextEnd { id: "t2".into() }), + Ok(LlmEvent::Finish { + reason: FinishReason::Stop, + usage: usage(20, 8), + }), + ], + ]); + + let cwd = tempfile::tempdir().unwrap(); + let ctx = make_ctx( + store.clone(), + bus, + session_id.clone(), + cwd.path().to_path_buf(), + ) + .await; + let run_config = RunConfig { + agent_name: "orchestrator".into(), + agent_prompt: "You are a helpful assistant.".into(), + model, + temperature: None, + max_steps: 10, + instructions: Vec::new(), + }; + + let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await; + assert!(matches!(outcome, RunOutcome::Stopped)); + + let messages = store.messages(session_id.clone()).await.unwrap(); + assert_eq!( + messages.len(), + 3, + "user + tool-call assistant turn + final assistant turn" + ); + assert_eq!(messages[0].role, Role::User); + assert_eq!(messages[1].role, Role::Assistant); + assert_eq!(messages[1].finished, Some(FinishReason::ToolCalls)); + assert_eq!(messages[2].role, Role::Assistant); + assert_eq!(messages[2].finished, Some(FinishReason::Stop)); + + let first_parts = store.parts(messages[1].id.clone()).await.unwrap(); + let tool_part = first_parts + .iter() + .find(|p| matches!(p.body, PartBody::Tool { .. })) + .expect("tool part persisted"); + match &tool_part.body { + PartBody::Tool { name, state, .. } => { + assert_eq!(name, "read"); + match state { + ToolState::Completed { output, .. } => assert_eq!(output, "mock file content"), + other => panic!("expected Completed tool state, got {other:?}"), + } + } + _ => unreachable!(), + } + let first_text = first_parts + .iter() + .find_map(|p| match &p.body { + PartBody::Text { text, .. } => Some(text.clone()), + _ => None, + }) + .expect("text part persisted"); + assert_eq!(first_text, "Let me check the file."); + + let final_parts = store.parts(messages[2].id.clone()).await.unwrap(); + let final_text = final_parts + .iter() + .find_map(|p| match &p.body { + PartBody::Text { text, .. } => Some(text.clone()), + _ => None, + }) + .expect("final text part persisted"); + assert_eq!(final_text, "The file contains: mock file content"); + } + + #[tokio::test] + async fn provider_error_on_first_event_is_reported_and_run_errors() { + let store = Store::open_in_memory().unwrap(); + let bus = EventBus::new(); + let model = ModelRef::new("mock", "mock-model"); + let session = Session::new_root("orchestrator", model.clone(), 1); + let session_id = session.id.clone(); + store.upsert_session(session).await.unwrap(); + + let user_message = Message::new_user(session_id.clone(), 1); + store.upsert_message(user_message.clone()).await.unwrap(); + store + .upsert_part(Part { + id: crate::types::PartId::new(), + message_id: user_message.id.clone(), + session_id: session_id.clone(), + idx: 0, + body: PartBody::Text { + text: "hi".into(), + synthetic: false, + }, + }) + .await + .unwrap(); + + // ContextOverflow is never retried, so this should surface immediately as Errored. + let provider = MockProvider::scripted(vec![vec![Err(ProviderError::ContextOverflow)]]); + + let cwd = tempfile::tempdir().unwrap(); + let ctx = make_ctx( + store.clone(), + bus, + session_id.clone(), + cwd.path().to_path_buf(), + ) + .await; + let run_config = RunConfig { + agent_name: "orchestrator".into(), + agent_prompt: "You are a helpful assistant.".into(), + model, + temperature: None, + max_steps: 10, + instructions: Vec::new(), + }; + + let outcome = run_session(Arc::new(provider), ctx, &run_config, || 2).await; + assert!(matches!(outcome, RunOutcome::Errored { .. })); + + // No assistant message should have been created — the error hit before any event. + let messages = store.messages(session_id).await.unwrap(); + assert_eq!(messages.len(), 1); + } +} diff --git a/crates/harness-core/src/engine/mod.rs b/crates/harness-core/src/engine/mod.rs new file mode 100644 index 0000000..c3d6c91 --- /dev/null +++ b/crates/harness-core/src/engine/mod.rs @@ -0,0 +1,10 @@ +pub mod doomloop; +pub mod processor; +pub mod retry; +#[path = "loop.rs"] +pub mod session_loop; +pub mod system; + +pub use doomloop::DoomLoopGuard; +pub use processor::{process_step, StepContext, StepError, StepOutcome, StepResult}; +pub use session_loop::{run_session, RunConfig}; diff --git a/crates/harness-core/src/engine/processor.rs b/crates/harness-core/src/engine/processor.rs new file mode 100644 index 0000000..b771a25 --- /dev/null +++ b/crates/harness-core/src/engine/processor.rs @@ -0,0 +1,652 @@ +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use futures::StreamExt; +use tokio_util::sync::CancellationToken; + +use crate::event::{AppEvent, EventBus}; +use crate::llm::{FinishReason, LlmEvent, LlmEventStream, ProviderError}; +use crate::permission::{PermissionService, Ruleset}; +use crate::store::Store; +use crate::tool::{MetadataSink, PermissionHandle, Tool, ToolCtx, ToolError, ToolRegistry}; +use crate::types::{Message, MessageId, Part, PartBody, PartId, SessionId, TokenUsage, ToolState}; + +use super::doomloop::DoomLoopGuard; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StepResult { + Continue, + Stop, + Compact, +} + +pub struct StepOutcome { + pub result: StepResult, + pub message_id: Option, + pub usage: TokenUsage, + pub aborted: bool, +} + +pub struct StepError { + pub source: ProviderError, + /// If `true`, the assistant message/parts were already persisted for this step — + /// the outer loop must surface this as a terminal error, not retry. + pub any_persisted: bool, +} + +/// Everything a step needs that stays constant across the run (cheap to construct per step; +/// owns clones/handles, not the session itself, so the loop keeps ownership of `Session`). +pub struct StepContext { + pub store: Store, + pub bus: EventBus, + pub tools: ToolRegistry, + pub permissions: Arc, + pub static_rules: Ruleset, + pub extra_rules: Arc>, + pub session_id: SessionId, + pub cwd: PathBuf, + /// Session's `tool-output` spill directory (see `tool::truncate`). + pub data_dir: PathBuf, + pub cancel: CancellationToken, + /// Wall-clock for `created_at` stamps — passed in so tests stay deterministic. + pub now: i64, +} + +struct FlushTracker { + last_flush: Instant, + bytes_since_flush: usize, +} + +impl FlushTracker { + fn new() -> Self { + Self { + last_flush: Instant::now(), + bytes_since_flush: 0, + } + } + + /// docs/02-engine.md: flush to store every 50ms or 2KB of buffered delta text. + fn should_flush(&mut self, new_bytes: usize) -> bool { + self.bytes_since_flush += new_bytes; + if self.bytes_since_flush >= 2048 || self.last_flush.elapsed() >= Duration::from_millis(50) + { + self.bytes_since_flush = 0; + self.last_flush = Instant::now(); + true + } else { + false + } + } +} + +struct Run<'a> { + ctx: &'a StepContext, + assistant: Option, + next_idx: u32, + active_text: Option, + active_reasoning: Option, + pending_tools: HashMap, + text_buf: HashMap, + reasoning_buf: HashMap, + tool_input_buf: HashMap, + flushers: HashMap, + parts_by_id: HashMap, +} + +impl<'a> Run<'a> { + fn new(ctx: &'a StepContext) -> Self { + Self { + ctx, + assistant: None, + next_idx: 0, + active_text: None, + active_reasoning: None, + pending_tools: HashMap::new(), + text_buf: HashMap::new(), + reasoning_buf: HashMap::new(), + tool_input_buf: HashMap::new(), + flushers: HashMap::new(), + parts_by_id: HashMap::new(), + } + } + + fn message_id(&self) -> MessageId { + self.assistant + .as_ref() + .expect("assistant created") + .id + .clone() + } + + async fn ensure_assistant( + &mut self, + model: crate::types::ModelRef, + agent: &str, + ) -> Result<(), ProviderError> { + if self.assistant.is_some() { + return Ok(()); + } + let message = + Message::new_assistant(self.ctx.session_id.clone(), model, agent, self.ctx.now); + self.ctx + .store + .upsert_message(message.clone()) + .await + .map_err(|e| ProviderError::Decode(e.to_string()))?; + self.ctx.bus.publish(AppEvent::MessageCreated { + message: message.clone(), + }); + + let step_start = Part { + id: PartId::new(), + message_id: message.id.clone(), + session_id: self.ctx.session_id.clone(), + idx: self.next_idx, + body: PartBody::StepStart, + }; + self.next_idx += 1; + self.persist_part(step_start).await?; + + self.assistant = Some(message); + Ok(()) + } + + async fn persist_part(&mut self, part: Part) -> Result<(), ProviderError> { + self.ctx + .store + .upsert_part(part.clone()) + .await + .map_err(|e| ProviderError::Decode(e.to_string()))?; + self.ctx + .bus + .publish(AppEvent::PartUpdated { part: part.clone() }); + self.parts_by_id.insert(part.id.clone(), part); + Ok(()) + } + + async fn on_text_start(&mut self, id: String) -> Result<(), ProviderError> { + let part_id = PartId::from(id); + let part = Part { + id: part_id.clone(), + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: self.next_idx, + body: PartBody::Text { + text: String::new(), + synthetic: false, + }, + }; + self.next_idx += 1; + self.text_buf.insert(part_id.clone(), String::new()); + self.flushers.insert(part_id.clone(), FlushTracker::new()); + self.active_text = Some(part_id); + self.persist_part(part).await + } + + async fn on_text_delta(&mut self, id: String, text: String) -> Result<(), ProviderError> { + let part_id = PartId::from(id); + let buf = self.text_buf.entry(part_id.clone()).or_default(); + buf.push_str(&text); + let full_text = buf.clone(); + + self.ctx.bus.publish(AppEvent::PartDelta { + part_id: part_id.clone(), + message_id: self.message_id(), + delta: text.clone(), + }); + + let should_flush = self + .flushers + .entry(part_id.clone()) + .or_insert_with(FlushTracker::new) + .should_flush(text.len()); + if should_flush { + let part = Part { + id: part_id, + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: 0, // overwritten below from cached idx + body: PartBody::Text { + text: full_text, + synthetic: false, + }, + }; + self.persist_existing(part).await?; + } + Ok(()) + } + + async fn on_text_end(&mut self, id: String) -> Result<(), ProviderError> { + let part_id = PartId::from(id); + let full_text = self.text_buf.get(&part_id).cloned().unwrap_or_default(); + let part = Part { + id: part_id, + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: 0, + body: PartBody::Text { + text: full_text, + synthetic: false, + }, + }; + self.active_text = None; + self.persist_existing(part).await + } + + async fn on_reasoning_start(&mut self, id: String) -> Result<(), ProviderError> { + let part_id = PartId::from(id); + let part = Part { + id: part_id.clone(), + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: self.next_idx, + body: PartBody::Reasoning { + text: String::new(), + signature: None, + }, + }; + self.next_idx += 1; + self.reasoning_buf.insert(part_id.clone(), String::new()); + self.flushers.insert(part_id.clone(), FlushTracker::new()); + self.active_reasoning = Some(part_id); + self.persist_part(part).await + } + + async fn on_reasoning_delta(&mut self, id: String, text: String) -> Result<(), ProviderError> { + let part_id = PartId::from(id); + let buf = self.reasoning_buf.entry(part_id.clone()).or_default(); + buf.push_str(&text); + let full_text = buf.clone(); + + self.ctx.bus.publish(AppEvent::PartDelta { + part_id: part_id.clone(), + message_id: self.message_id(), + delta: text.clone(), + }); + + let should_flush = self + .flushers + .entry(part_id.clone()) + .or_insert_with(FlushTracker::new) + .should_flush(text.len()); + if should_flush { + let part = Part { + id: part_id, + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: 0, + body: PartBody::Reasoning { + text: full_text, + signature: None, + }, + }; + self.persist_existing(part).await?; + } + Ok(()) + } + + async fn on_reasoning_end( + &mut self, + id: String, + signature: Option, + ) -> Result<(), ProviderError> { + let part_id = PartId::from(id); + let full_text = self + .reasoning_buf + .get(&part_id) + .cloned() + .unwrap_or_default(); + let part = Part { + id: part_id, + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: 0, + body: PartBody::Reasoning { + text: full_text, + signature, + }, + }; + self.active_reasoning = None; + self.persist_existing(part).await + } + + /// Re-persists a part that already exists, preserving its original `idx`. + async fn persist_existing(&mut self, mut part: Part) -> Result<(), ProviderError> { + if let Some(existing) = self.parts_by_id.get(&part.id) { + part.idx = existing.idx; + } + self.persist_part(part).await + } + + async fn on_tool_input_start( + &mut self, + call_id: String, + name: String, + ) -> Result<(), ProviderError> { + let part_id = PartId::new(); + self.pending_tools.insert(call_id.clone(), part_id.clone()); + self.tool_input_buf.insert(call_id.clone(), String::new()); + let part = Part { + id: part_id, + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: self.next_idx, + body: PartBody::Tool { + call_id, + name, + state: ToolState::Pending { + partial_input: String::new(), + }, + }, + }; + self.next_idx += 1; + self.persist_part(part).await + } + + async fn on_tool_input_delta( + &mut self, + call_id: String, + json: String, + ) -> Result<(), ProviderError> { + let buf = self.tool_input_buf.entry(call_id.clone()).or_default(); + buf.push_str(&json); + if let Some(part_id) = self.pending_tools.get(&call_id).cloned() { + self.ctx.bus.publish(AppEvent::PartDelta { + part_id, + message_id: self.message_id(), + delta: json, + }); + } + Ok(()) + } + + async fn on_tool_call( + &mut self, + call_id: String, + name: String, + input: serde_json::Value, + doomloop: &mut DoomLoopGuard, + ) -> Result<(), ProviderError> { + let part_id = self.pending_tools.remove(&call_id).unwrap_or_default(); + let running = Part { + id: part_id.clone(), + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: self.next_idx_for(&part_id), + body: PartBody::Tool { + call_id: call_id.clone(), + name: name.clone(), + state: ToolState::Running { + input: input.clone(), + title: None, + metadata: serde_json::Value::Null, + }, + }, + }; + self.persist_existing(running).await?; + + let mut doom_blocked = None; + if doomloop.record_and_check(&name, &input) { + let cancel = self.ctx.cancel.child_token(); + if let Err(e) = self + .ctx + .permissions + .force_ask( + &self.ctx.session_id, + crate::permission::AskInput { + permission: "doom_loop".into(), + pattern: name.clone(), + always_pattern: name.clone(), + metadata: serde_json::json!({"tool": name, "input": input}), + }, + &cancel, + ) + .await + { + doom_blocked = Some(format!("model appears stuck repeating {name}: {e}")); + } + } + + let tool = self.ctx.tools.get(&name); + let (title, output, metadata, is_error) = if let Some(reason) = doom_blocked { + ("error".to_string(), reason, serde_json::Value::Null, true) + } else { + match tool { + None => ( + "error".to_string(), + format!("unknown tool: {name}"), + serde_json::Value::Null, + true, + ), + Some(tool) => self.run_tool(tool, call_id.clone(), input.clone()).await, + } + }; + + let final_state = if is_error { + ToolState::Error { + input, + error: output.clone(), + } + } else { + ToolState::Completed { + input, + title: title.clone(), + output: output.clone(), + metadata, + duration_ms: 0, + } + }; + let final_part = Part { + id: part_id, + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: 0, + body: PartBody::Tool { + call_id, + name, + state: final_state, + }, + }; + self.persist_existing(final_part).await + } + + fn next_idx_for(&mut self, part_id: &PartId) -> u32 { + if let Some(existing) = self.parts_by_id.get(part_id) { + existing.idx + } else { + let idx = self.next_idx; + self.next_idx += 1; + idx + } + } + + async fn run_tool( + &self, + tool: Arc, + call_id: String, + input: serde_json::Value, + ) -> (String, String, serde_json::Value, bool) { + let call_cancel = self.ctx.cancel.child_token(); + let (metadata_sink, _metadata_rx) = MetadataSink::channel(); + let ask = PermissionHandle::new( + self.ctx.permissions.clone(), + self.ctx.session_id.clone(), + self.ctx.static_rules.clone(), + self.ctx.extra_rules.clone(), + call_cancel.clone(), + ); + let tool_ctx = ToolCtx { + session_id: self.ctx.session_id.clone(), + message_id: self.message_id(), + call_id: call_id.clone(), + cwd: self.ctx.cwd.clone(), + data_dir: self.ctx.data_dir.clone(), + cancel: call_cancel.clone(), + ask, + metadata: metadata_sink, + }; + + let result = tokio::select! { + res = tool.execute(input, tool_ctx) => res, + _ = call_cancel.cancelled() => Err(ToolError::Cancelled), + }; + + match result { + Ok(output) => (output.title, output.output, output.metadata, false), + Err(err) => ( + "error".to_string(), + err.to_string(), + serde_json::Value::Null, + true, + ), + } + } + + async fn on_finish( + &mut self, + reason: FinishReason, + usage: TokenUsage, + ) -> Result { + let part = Part { + id: PartId::new(), + message_id: self.message_id(), + session_id: self.ctx.session_id.clone(), + idx: self.next_idx, + body: PartBody::StepFinish { + usage, + cost: 0.0, + reason: reason.clone(), + }, + }; + self.next_idx += 1; + self.persist_part(part).await?; + + let result = match reason { + FinishReason::ToolCalls => StepResult::Continue, + _ => StepResult::Stop, + }; + + if let Some(message) = self.assistant.as_mut() { + message.finished = Some(reason); + message.usage = usage; + self.ctx + .store + .upsert_message(message.clone()) + .await + .map_err(|e| ProviderError::Decode(e.to_string()))?; + self.ctx.bus.publish(AppEvent::MessageUpdated { + message: message.clone(), + }); + } + Ok(result) + } + + async fn mark_errored(&mut self, source: &ProviderError) { + if let Some(message) = self.assistant.as_mut() { + message.error = Some(crate::types::MessageError::Provider(source.to_string())); + if self.ctx.store.upsert_message(message.clone()).await.is_ok() { + self.ctx.bus.publish(AppEvent::MessageUpdated { + message: message.clone(), + }); + } + } + } +} + +/// Consumes one provider stream end-to-end: persists parts/messages as events arrive and +/// executes tool calls inline, in stream order (opencode's processor.ts semantics). +pub async fn process_step( + mut stream: LlmEventStream, + ctx: &StepContext, + model: crate::types::ModelRef, + agent: &str, + doomloop: &mut DoomLoopGuard, +) -> Result { + let mut run = Run::new(ctx); + let mut usage = TokenUsage::default(); + let mut result = StepResult::Stop; + + loop { + let cancelled = ctx.cancel.is_cancelled(); + if cancelled { + run.mark_errored(&ProviderError::Cancelled).await; + return Ok(StepOutcome { + result: StepResult::Stop, + message_id: run.assistant.as_ref().map(|m| m.id.clone()), + usage, + aborted: true, + }); + } + + let item = stream.next().await; + let Some(item) = item else { break }; + + let event = match item { + Ok(event) => event, + Err(source) => { + let any_persisted = run.assistant.is_some(); + run.mark_errored(&source).await; + return Err(StepError { + source, + any_persisted, + }); + } + }; + + run.ensure_assistant(model.clone(), agent) + .await + .map_err(|source| StepError { + source, + any_persisted: false, + })?; + + let outcome = match event { + LlmEvent::TextStart { id } => run.on_text_start(id).await, + LlmEvent::TextDelta { id, text } => run.on_text_delta(id, text).await, + LlmEvent::TextEnd { id } => run.on_text_end(id).await, + LlmEvent::ReasoningStart { id } => run.on_reasoning_start(id).await, + LlmEvent::ReasoningDelta { id, text } => run.on_reasoning_delta(id, text).await, + LlmEvent::ReasoningEnd { id, signature } => run.on_reasoning_end(id, signature).await, + LlmEvent::ToolInputStart { call_id, name } => { + run.on_tool_input_start(call_id, name).await + } + LlmEvent::ToolInputDelta { call_id, json } => { + run.on_tool_input_delta(call_id, json).await + } + LlmEvent::ToolCall { + call_id, + name, + input, + } => run.on_tool_call(call_id, name, input, doomloop).await, + LlmEvent::Finish { + reason, + usage: finish_usage, + } => { + usage = finish_usage; + match run.on_finish(reason, finish_usage).await { + Ok(r) => { + result = r; + Ok(()) + } + Err(e) => Err(e), + } + } + }; + + if let Err(source) = outcome { + return Err(StepError { + source, + any_persisted: true, + }); + } + } + + Ok(StepOutcome { + result, + message_id: run.assistant.as_ref().map(|m| m.id.clone()), + usage, + aborted: false, + }) +} diff --git a/crates/harness-core/src/engine/retry.rs b/crates/harness-core/src/engine/retry.rs new file mode 100644 index 0000000..4a38111 --- /dev/null +++ b/crates/harness-core/src/engine/retry.rs @@ -0,0 +1,113 @@ +use std::future::Future; +use std::time::Duration; + +use tokio_util::sync::CancellationToken; + +use crate::llm::ProviderError; + +const MAX_ATTEMPTS: u32 = 8; + +fn delay_for(err: &ProviderError, attempt: u32) -> Duration { + if let ProviderError::RateLimited { + retry_after: Some(d), + } = err + { + return *d; + } + let secs = 2u64.saturating_mul(1u64 << attempt.min(4)); + Duration::from_secs(secs.min(30)) +} + +/// Wraps only the *start* of a step (acquiring the provider stream) in retry — once any +/// event has been consumed from the stream we surface the error instead (opencode behavior). +/// +/// `ContextOverflow` is never retried; the caller reacts with `StepResult::Compact`. +pub async fn with_retry( + cancel: &CancellationToken, + mut attempt_fn: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempt = 0u32; + loop { + match attempt_fn().await { + Ok(v) => return Ok(v), + Err(e) if e.is_retryable() && attempt + 1 < MAX_ATTEMPTS => { + let delay = delay_for(&e, attempt); + attempt += 1; + tokio::select! { + _ = tokio::time::sleep(delay) => {} + _ = cancel.cancelled() => return Err(ProviderError::Cancelled), + } + } + Err(e) => return Err(e), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicU32, Ordering}; + + #[tokio::test(start_paused = true)] + async fn retries_rate_limited_until_success() { + let calls = AtomicU32::new(0); + let cancel = CancellationToken::new(); + let result = with_retry(&cancel, || { + let n = calls.fetch_add(1, Ordering::SeqCst); + async move { + if n < 2 { + Err(ProviderError::RateLimited { retry_after: None }) + } else { + Ok(42) + } + } + }) + .await; + assert_eq!(result.unwrap(), 42); + assert_eq!(calls.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn does_not_retry_context_overflow() { + let calls = AtomicU32::new(0); + let cancel = CancellationToken::new(); + let result: Result<(), _> = with_retry(&cancel, || { + calls.fetch_add(1, Ordering::SeqCst); + async move { Err(ProviderError::ContextOverflow) } + }) + .await; + assert!(matches!(result, Err(ProviderError::ContextOverflow))); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test(start_paused = true)] + async fn gives_up_after_max_attempts() { + let calls = AtomicU32::new(0); + let cancel = CancellationToken::new(); + let result: Result<(), _> = with_retry(&cancel, || { + calls.fetch_add(1, Ordering::SeqCst); + async move { Err(ProviderError::Overloaded) } + }) + .await; + assert!(result.is_err()); + assert_eq!(calls.load(Ordering::SeqCst), MAX_ATTEMPTS); + } + + #[tokio::test] + async fn cancellation_aborts_the_wait_between_retries() { + let calls = AtomicU32::new(0); + let cancel = CancellationToken::new(); + cancel.cancel(); + let result: Result<(), _> = with_retry(&cancel, || { + calls.fetch_add(1, Ordering::SeqCst); + async move { Err(ProviderError::Overloaded) } + }) + .await; + assert!(matches!(result, Err(ProviderError::Cancelled))); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } +} diff --git a/crates/harness-core/src/engine/system.rs b/crates/harness-core/src/engine/system.rs new file mode 100644 index 0000000..22b0349 --- /dev/null +++ b/crates/harness-core/src/engine/system.rs @@ -0,0 +1,56 @@ +use std::path::Path; + +/// cwd, platform, date, best-effort git status. Kept as its own block (not merged into the +/// agent prompt) because Anthropic cache breakpoints need the system prompt structured as +/// separate blocks (see `LlmRequest.system: Vec`). +pub fn env_header(cwd: &Path) -> String { + let git_status = std::process::Command::new("git") + .arg("status") + .arg("--short") + .current_dir(cwd) + .output() + .ok() + .filter(|o| o.status.success()) + .map(|o| String::from_utf8_lossy(&o.stdout).trim().to_string()) + .filter(|s| !s.is_empty()); + + let mut header = format!( + "Working directory: {}\nPlatform: {}\n", + cwd.display(), + std::env::consts::OS + ); + if let Some(status) = git_status { + header.push_str(&format!("Git status:\n{status}\n")); + } + header +} + +/// Ordered system prompt blocks: environment header, agent prompt, then project +/// instructions (e.g. AGENTS.md contents). Order matches `02-engine.md`. +pub fn assemble(env_header: String, agent_prompt: &str, instructions: &[String]) -> Vec { + let mut blocks = vec![env_header, agent_prompt.to_string()]; + blocks.extend(instructions.iter().cloned()); + blocks +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn env_header_includes_cwd_and_platform() { + let header = env_header(Path::new(".")); + assert!(header.contains("Working directory:")); + assert!(header.contains(std::env::consts::OS)); + } + + #[test] + fn assemble_orders_env_agent_then_instructions() { + let blocks = assemble( + "ENV".to_string(), + "AGENT", + &["AGENTS.md contents".to_string()], + ); + assert_eq!(blocks, vec!["ENV", "AGENT", "AGENTS.md contents"]); + } +} diff --git a/crates/harness-core/src/event.rs b/crates/harness-core/src/event.rs index f687aba..87316f0 100644 --- a/crates/harness-core/src/event.rs +++ b/crates/harness-core/src/event.rs @@ -24,7 +24,8 @@ pub struct PermissionRequest { pub id: String, pub session_id: SessionId, pub permission: String, - pub patterns: Vec, + pub pattern: String, + pub always_pattern: String, pub metadata: serde_json::Value, } diff --git a/crates/harness-core/src/lib.rs b/crates/harness-core/src/lib.rs index 60206c0..f224e12 100644 --- a/crates/harness-core/src/lib.rs +++ b/crates/harness-core/src/lib.rs @@ -1,7 +1,10 @@ +pub mod config; +pub mod engine; pub mod event; pub mod llm; pub mod permission; pub mod store; +pub mod tool; pub mod types; pub use event::{AppEvent, EventBus}; diff --git a/crates/harness-core/src/llm.rs b/crates/harness-core/src/llm.rs index 0061f7e..3727510 100644 --- a/crates/harness-core/src/llm.rs +++ b/crates/harness-core/src/llm.rs @@ -1,7 +1,13 @@ +use std::pin::Pin; +use std::time::Duration; + +use async_trait::async_trait; +use futures::Stream; use serde::{Deserialize, Serialize}; +use tokio_util::sync::CancellationToken; + +use crate::types::{ModelInfo, TokenUsage}; -// Full LlmEvent/LlmRequest/Provider trait land in M1 alongside harness-providers. -// FinishReason lives here (not types/) because it's part of the provider-facing vocabulary. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub enum FinishReason { Stop, @@ -10,3 +16,185 @@ pub enum FinishReason { ContentFilter, Unknown(String), } + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub enum LlmEvent { + TextStart { + id: String, + }, + TextDelta { + id: String, + text: String, + }, + TextEnd { + id: String, + }, + ReasoningStart { + id: String, + }, + ReasoningDelta { + id: String, + text: String, + }, + ReasoningEnd { + id: String, + signature: Option, + }, + ToolInputStart { + call_id: String, + name: String, + }, + ToolInputDelta { + call_id: String, + json: String, + }, + ToolCall { + call_id: String, + name: String, + input: serde_json::Value, + }, + Finish { + reason: FinishReason, + usage: TokenUsage, + }, +} + +#[derive(Debug, Clone)] +pub enum ProviderError { + RateLimited { + retry_after: Option, + }, + Overloaded, + /// Never retried — the engine reacts with `StepResult::Compact`. + ContextOverflow, + Auth(String), + Http { + status: u16, + body: String, + }, + Network(String), + Decode(String), + Cancelled, +} + +impl std::fmt::Display for ProviderError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ProviderError::RateLimited { retry_after } => { + write!(f, "rate limited (retry_after={retry_after:?})") + } + ProviderError::Overloaded => write!(f, "provider overloaded"), + ProviderError::ContextOverflow => write!(f, "context window exceeded"), + ProviderError::Auth(msg) => write!(f, "auth error: {msg}"), + ProviderError::Http { status, body } => write!(f, "http {status}: {body}"), + ProviderError::Network(msg) => write!(f, "network error: {msg}"), + ProviderError::Decode(msg) => write!(f, "decode error: {msg}"), + ProviderError::Cancelled => write!(f, "cancelled"), + } + } +} + +impl std::error::Error for ProviderError {} + +impl ProviderError { + /// Retry policy classification (`engine/retry.rs`): 5xx counts as `Http`, so check status. + pub fn is_retryable(&self) -> bool { + matches!( + self, + ProviderError::RateLimited { .. } + | ProviderError::Overloaded + | ProviderError::Network(_) + ) || matches!(self, ProviderError::Http { status, .. } if *status >= 500) + } +} + +pub type LlmEventStream = Pin> + Send>>; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Role { + System, + User, + Assistant, + Tool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum WireContent { + Text { + text: String, + }, + ToolCall { + call_id: String, + name: String, + input: serde_json::Value, + }, + ToolResult { + call_id: String, + output: String, + is_error: bool, + }, + Image { + mime_type: String, + data: String, + }, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct WireMessage { + pub role: Role, + pub content: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolSchema { + pub name: String, + pub description: String, + pub parameters: serde_json::Value, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ReasoningEffort { + Low, + Medium, + High, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ReasoningOpts { + pub effort: Option, + /// Anthropic extended-thinking token budget. + pub budget_tokens: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Initiator { + User, + Agent, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LlmRequest { + pub model: String, + pub system: Vec, + pub messages: Vec, + pub tools: Vec, + pub temperature: Option, + pub max_tokens: Option, + pub reasoning: Option, + pub initiator: Initiator, +} + +#[async_trait] +pub trait Provider: Send + Sync { + fn id(&self) -> &str; + async fn list_models(&self) -> Result, ProviderError>; + async fn stream( + &self, + req: LlmRequest, + cancel: CancellationToken, + ) -> Result; +} diff --git a/crates/harness-core/src/permission/mod.rs b/crates/harness-core/src/permission/mod.rs index cb36611..17f4b2d 100644 --- a/crates/harness-core/src/permission/mod.rs +++ b/crates/harness-core/src/permission/mod.rs @@ -1,3 +1,7 @@ pub mod rule; +pub mod service; pub use rule::{evaluate, Action, Rule, Ruleset}; +pub use service::{ + spawn_auto_approve, AskDecision, AskError, AskInput, PermissionReply, PermissionService, +}; diff --git a/crates/harness-core/src/permission/service.rs b/crates/harness-core/src/permission/service.rs new file mode 100644 index 0000000..f4dc957 --- /dev/null +++ b/crates/harness-core/src/permission/service.rs @@ -0,0 +1,359 @@ +use std::collections::HashMap; +use std::sync::Mutex; + +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; +use ulid::Ulid; + +use crate::event::{AppEvent, EventBus, PermissionRequest}; +use crate::types::SessionId; + +use super::rule::{evaluate, Action, Rule, Ruleset}; + +pub struct AskInput { + pub permission: String, + /// Pattern evaluated against the ruleset stack, and granted on a `Once` reply. + pub pattern: String, + /// Coarser pattern granted (as an `Allow` rule on `session.extra_rules`) on an `Always` reply. + pub always_pattern: String, + pub metadata: serde_json::Value, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PermissionReply { + Once, + Always, + Reject, +} + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum AskError { + #[error("permission denied by ruleset")] + Denied, + #[error("permission rejected by user")] + Rejected, + #[error("cancelled")] + Cancelled, +} + +#[derive(Debug)] +pub enum AskDecision { + /// Ruleset allowed it outright, or the user replied `Once`. + Allowed, + /// User replied `Always` — caller persists this rule onto `session.extra_rules`. + AllowedAlways(Rule), +} + +/// Port of opencode's permission/index.ts ask flow: evaluate the ruleset stack first; +/// only fall through to a blocking, event-published ask when the verdict is `Ask`. +pub struct PermissionService { + bus: EventBus, + pending: Mutex>>, +} + +impl PermissionService { + pub fn new(bus: EventBus) -> Self { + Self { + bus, + pending: Mutex::new(HashMap::new()), + } + } + + pub async fn ask( + &self, + session_id: &SessionId, + stack: &[&Ruleset], + input: AskInput, + cancel: &CancellationToken, + ) -> Result { + match evaluate(stack, &input.permission, &input.pattern) { + Action::Allow => Ok(AskDecision::Allowed), + Action::Deny => Err(AskError::Denied), + Action::Ask => self.ask_user(session_id, input, cancel).await, + } + } + + /// Bypasses ruleset evaluation entirely — used by the doom-loop guard, which must ask + /// regardless of any `Allow` rule. + pub async fn force_ask( + &self, + session_id: &SessionId, + input: AskInput, + cancel: &CancellationToken, + ) -> Result { + self.ask_user(session_id, input, cancel).await + } + + async fn ask_user( + &self, + session_id: &SessionId, + input: AskInput, + cancel: &CancellationToken, + ) -> Result { + let id = Ulid::new().to_string(); + let (tx, rx) = oneshot::channel(); + self.pending.lock().unwrap().insert(id.clone(), tx); + + self.bus.publish(AppEvent::PermissionAsked { + request: PermissionRequest { + id: id.clone(), + session_id: session_id.clone(), + permission: input.permission.clone(), + pattern: input.pattern.clone(), + always_pattern: input.always_pattern.clone(), + metadata: input.metadata.clone(), + }, + }); + + let reply = tokio::select! { + reply = rx => reply.map_err(|_| AskError::Cancelled)?, + _ = cancel.cancelled() => { + self.pending.lock().unwrap().remove(&id); + return Err(AskError::Cancelled); + } + }; + + self.bus.publish(AppEvent::PermissionResolved { id }); + + match reply { + PermissionReply::Once => Ok(AskDecision::Allowed), + PermissionReply::Always => Ok(AskDecision::AllowedAlways(Rule { + permission: input.permission, + pattern: input.always_pattern, + action: Action::Allow, + })), + PermissionReply::Reject => Err(AskError::Rejected), + } + } + + /// Resolve a pending ask (called by a frontend: TUI keypress, auto-approve stub, ...). + /// Returns `false` if `id` had already been resolved or never existed. + pub fn reply(&self, id: &str, reply: PermissionReply) -> bool { + match self.pending.lock().unwrap().remove(id) { + Some(tx) => tx.send(reply).is_ok(), + None => false, + } + } +} + +/// Auto-approves every ask as `Once` — the M1..M2 stand-in for a real TUI permission modal. +pub fn spawn_auto_approve(bus: EventBus, service: std::sync::Arc) { + let mut rx = bus.subscribe(); + tokio::spawn(async move { + while let Ok(event) = rx.recv().await { + if let AppEvent::PermissionAsked { request } = event { + service.reply(&request.id, PermissionReply::Once); + } + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rule(permission: &str, pattern: &str, action: Action) -> Rule { + Rule { + permission: permission.to_string(), + pattern: pattern.to_string(), + action, + } + } + + #[tokio::test] + async fn allow_rule_short_circuits_without_asking() { + let service = PermissionService::new(EventBus::new()); + let rules: Ruleset = vec![rule("bash", "ls*", Action::Allow)]; + let outcome = service + .ask( + &SessionId::new(), + &[&rules], + AskInput { + permission: "bash".into(), + pattern: "ls -la".into(), + always_pattern: "ls *".into(), + metadata: serde_json::Value::Null, + }, + &CancellationToken::new(), + ) + .await + .unwrap(); + assert!(matches!(outcome, AskDecision::Allowed)); + } + + #[tokio::test] + async fn deny_rule_short_circuits_to_denied_error() { + let service = PermissionService::new(EventBus::new()); + let rules: Ruleset = vec![rule("edit", "*.lock", Action::Deny)]; + let err = service + .ask( + &SessionId::new(), + &[&rules], + AskInput { + permission: "edit".into(), + pattern: "Cargo.lock".into(), + always_pattern: "Cargo.lock".into(), + metadata: serde_json::Value::Null, + }, + &CancellationToken::new(), + ) + .await + .unwrap_err(); + assert_eq!(err, AskError::Denied); + } + + #[tokio::test] + async fn ask_publishes_event_and_waits_for_reply() { + let bus = EventBus::new(); + let mut sub = bus.subscribe(); + let service = std::sync::Arc::new(PermissionService::new(bus)); + let svc = service.clone(); + + let handle = tokio::spawn(async move { + svc.ask( + &SessionId::new(), + &[], + AskInput { + permission: "bash".into(), + pattern: "rm -rf /".into(), + always_pattern: "rm *".into(), + metadata: serde_json::Value::Null, + }, + &CancellationToken::new(), + ) + .await + }); + + let event = sub.recv().await.unwrap(); + let id = match event { + AppEvent::PermissionAsked { request } => request.id, + _ => panic!("expected PermissionAsked"), + }; + assert!(service.reply(&id, PermissionReply::Once)); + + let outcome = handle.await.unwrap().unwrap(); + assert!(matches!(outcome, AskDecision::Allowed)); + } + + #[tokio::test] + async fn always_reply_yields_rule_to_persist() { + let bus = EventBus::new(); + let mut sub = bus.subscribe(); + let service = std::sync::Arc::new(PermissionService::new(bus)); + let svc = service.clone(); + + let handle = tokio::spawn(async move { + svc.ask( + &SessionId::new(), + &[], + AskInput { + permission: "bash".into(), + pattern: "git push origin main".into(), + always_pattern: "git push*".into(), + metadata: serde_json::Value::Null, + }, + &CancellationToken::new(), + ) + .await + }); + + let request = match sub.recv().await.unwrap() { + AppEvent::PermissionAsked { request } => request, + _ => panic!("expected PermissionAsked"), + }; + service.reply(&request.id, PermissionReply::Always); + + let outcome = handle.await.unwrap().unwrap(); + match outcome { + AskDecision::AllowedAlways(rule) => { + assert_eq!(rule.permission, "bash"); + assert_eq!(rule.pattern, "git push*"); + assert_eq!(rule.action, Action::Allow); + } + _ => panic!("expected AllowedAlways"), + } + } + + #[tokio::test] + async fn reject_reply_yields_rejected_error() { + let bus = EventBus::new(); + let mut sub = bus.subscribe(); + let service = std::sync::Arc::new(PermissionService::new(bus)); + let svc = service.clone(); + + let handle = tokio::spawn(async move { + svc.ask( + &SessionId::new(), + &[], + AskInput { + permission: "bash".into(), + pattern: "rm -rf /".into(), + always_pattern: "rm *".into(), + metadata: serde_json::Value::Null, + }, + &CancellationToken::new(), + ) + .await + }); + + let request = match sub.recv().await.unwrap() { + AppEvent::PermissionAsked { request } => request, + _ => panic!("expected PermissionAsked"), + }; + service.reply(&request.id, PermissionReply::Reject); + + let err = handle.await.unwrap().unwrap_err(); + assert_eq!(err, AskError::Rejected); + } + + #[tokio::test] + async fn cancellation_aborts_a_pending_ask() { + let bus = EventBus::new(); + let service = std::sync::Arc::new(PermissionService::new(bus)); + let cancel = CancellationToken::new(); + let cancel2 = cancel.clone(); + let svc = service.clone(); + + let handle = tokio::spawn(async move { + svc.ask( + &SessionId::new(), + &[], + AskInput { + permission: "bash".into(), + pattern: "rm -rf /".into(), + always_pattern: "rm *".into(), + metadata: serde_json::Value::Null, + }, + &cancel2, + ) + .await + }); + + cancel.cancel(); + let err = handle.await.unwrap().unwrap_err(); + assert_eq!(err, AskError::Cancelled); + } + + #[tokio::test] + async fn auto_approve_resolves_asks_as_once() { + let bus = EventBus::new(); + let service = std::sync::Arc::new(PermissionService::new(bus.clone())); + spawn_auto_approve(bus, service.clone()); + + let outcome = service + .ask( + &SessionId::new(), + &[], + AskInput { + permission: "bash".into(), + pattern: "ls".into(), + always_pattern: "ls *".into(), + metadata: serde_json::Value::Null, + }, + &CancellationToken::new(), + ) + .await + .unwrap(); + assert!(matches!(outcome, AskDecision::Allowed)); + } +} diff --git a/crates/harness-core/src/tool/mod.rs b/crates/harness-core/src/tool/mod.rs new file mode 100644 index 0000000..9c06e08 --- /dev/null +++ b/crates/harness-core/src/tool/mod.rs @@ -0,0 +1,295 @@ +pub mod truncate; + +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use tokio_util::sync::CancellationToken; + +use crate::permission::{AskDecision, AskError, AskInput, PermissionService, Ruleset}; +use crate::types::{MessageId, SessionId}; + +#[derive(Debug, thiserror::Error)] +pub enum ToolError { + #[error("permission denied")] + Denied, + #[error("rejected by user")] + Rejected, + #[error("cancelled")] + Cancelled, + #[error("invalid input: {0}")] + Invalid(String), + #[error("{0}")] + Other(String), +} + +impl From for ToolError { + fn from(err: AskError) -> Self { + match err { + AskError::Denied => ToolError::Denied, + AskError::Rejected => ToolError::Rejected, + AskError::Cancelled => ToolError::Cancelled, + } + } +} + +/// Live metadata updates for a running tool call — consumed by the processor to +/// upsert the part's `ToolState::Running.metadata` and emit `AppEvent::PartUpdated`. +#[derive(Clone)] +pub struct MetadataSink(tokio::sync::mpsc::UnboundedSender); + +impl MetadataSink { + pub fn channel() -> ( + Self, + tokio::sync::mpsc::UnboundedReceiver, + ) { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + (Self(tx), rx) + } + + pub fn set(&self, value: serde_json::Value) { + let _ = self.0.send(value); + } +} + +/// Bound to one tool call: session + a config/agent ruleset snapshot + the session's +/// live `extra_rules` (mutated in place when the user replies "Always"). +pub struct PermissionHandle { + service: Arc, + session_id: SessionId, + static_rules: Ruleset, + extra_rules: Arc>, + cancel: CancellationToken, +} + +impl PermissionHandle { + pub fn new( + service: Arc, + session_id: SessionId, + static_rules: Ruleset, + extra_rules: Arc>, + cancel: CancellationToken, + ) -> Self { + Self { + service, + session_id, + static_rules, + extra_rules, + cancel, + } + } + + pub async fn ask( + &self, + permission: impl Into, + pattern: impl Into, + always_pattern: impl Into, + metadata: serde_json::Value, + ) -> Result<(), ToolError> { + let extra_snapshot = self.extra_rules.lock().unwrap().clone(); + let stack: [&Ruleset; 2] = [&self.static_rules, &extra_snapshot]; + let decision = self + .service + .ask( + &self.session_id, + &stack, + AskInput { + permission: permission.into(), + pattern: pattern.into(), + always_pattern: always_pattern.into(), + metadata, + }, + &self.cancel, + ) + .await?; + if let AskDecision::AllowedAlways(rule) = decision { + self.extra_rules.lock().unwrap().push(rule); + } + Ok(()) + } +} + +pub struct ToolCtx { + pub session_id: SessionId, + pub message_id: MessageId, + pub call_id: String, + pub cwd: PathBuf, + /// Per-session directory tool output spills to when it exceeds `truncate::MAX_CHARS`. + pub data_dir: PathBuf, + pub cancel: CancellationToken, + pub ask: PermissionHandle, + pub metadata: MetadataSink, +} + +pub struct ToolOutput { + pub title: String, + pub output: String, + pub metadata: serde_json::Value, +} + +impl ToolOutput { + pub fn new(title: impl Into, output: impl Into) -> Self { + Self { + title: title.into(), + output: output.into(), + metadata: serde_json::Value::Null, + } + } +} + +#[async_trait] +pub trait Tool: Send + Sync { + fn name(&self) -> &str; + fn description(&self) -> &str; + fn parameters(&self) -> serde_json::Value; + async fn execute( + &self, + input: serde_json::Value, + ctx: ToolCtx, + ) -> Result; +} + +/// Also exposed for tests/tool-error metadata (opencode's `invalid` pattern: echo the schema back). +pub fn invalid_input(tool: &dyn Tool, message: impl std::fmt::Display) -> ToolError { + ToolError::Invalid(format!("{message}\nExpected schema: {}", tool.parameters())) +} + +#[derive(Default, Clone)] +pub struct ToolRegistry { + tools: HashMap>, +} + +impl ToolRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn register(&mut self, tool: Arc) { + self.tools.insert(tool.name().to_string(), tool); + } + + pub fn get(&self, name: &str) -> Option> { + self.tools.get(name).cloned() + } + + pub fn all(&self) -> Vec> { + self.tools.values().cloned().collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::event::EventBus; + use crate::permission::Rule; + + struct Echo; + + #[async_trait] + impl Tool for Echo { + fn name(&self) -> &str { + "echo" + } + fn description(&self) -> &str { + "echoes input" + } + fn parameters(&self) -> serde_json::Value { + serde_json::json!({"type": "object"}) + } + async fn execute( + &self, + input: serde_json::Value, + _ctx: ToolCtx, + ) -> Result { + Ok(ToolOutput::new("echo", input.to_string())) + } + } + + #[test] + fn registry_registers_and_looks_up_by_name() { + let mut registry = ToolRegistry::new(); + registry.register(Arc::new(Echo)); + assert!(registry.get("echo").is_some()); + assert!(registry.get("missing").is_none()); + assert_eq!(registry.all().len(), 1); + } + + #[tokio::test] + async fn permission_handle_allows_when_ruleset_grants() { + let service = Arc::new(PermissionService::new(EventBus::new())); + let handle = PermissionHandle::new( + service, + SessionId::new(), + vec![Rule { + permission: "bash".into(), + pattern: "ls*".into(), + action: crate::permission::Action::Allow, + }], + Arc::new(Mutex::new(Vec::new())), + CancellationToken::new(), + ); + handle + .ask("bash", "ls -la", "ls *", serde_json::Value::Null) + .await + .unwrap(); + } + + #[tokio::test] + async fn permission_handle_denies_and_maps_to_tool_error() { + let service = Arc::new(PermissionService::new(EventBus::new())); + let handle = PermissionHandle::new( + service, + SessionId::new(), + vec![Rule { + permission: "bash".into(), + pattern: "rm*".into(), + action: crate::permission::Action::Deny, + }], + Arc::new(Mutex::new(Vec::new())), + CancellationToken::new(), + ); + let err = handle + .ask("bash", "rm -rf /", "rm *", serde_json::Value::Null) + .await + .unwrap_err(); + assert!(matches!(err, ToolError::Denied)); + } + + #[tokio::test] + async fn permission_handle_always_mutates_shared_extra_rules() { + let bus = EventBus::new(); + let mut sub = bus.subscribe(); + let service = Arc::new(PermissionService::new(bus)); + let extra = Arc::new(Mutex::new(Vec::new())); + let handle = PermissionHandle::new( + service.clone(), + SessionId::new(), + Vec::new(), + extra.clone(), + CancellationToken::new(), + ); + + let ask = tokio::spawn(async move { + handle + .ask( + "bash", + "git push origin", + "git push*", + serde_json::Value::Null, + ) + .await + }); + + let request = match sub.recv().await.unwrap() { + crate::event::AppEvent::PermissionAsked { request } => request, + _ => panic!("expected PermissionAsked"), + }; + service.reply(&request.id, crate::permission::PermissionReply::Always); + ask.await.unwrap().unwrap(); + + let rules = extra.lock().unwrap(); + assert_eq!(rules.len(), 1); + assert_eq!(rules[0].pattern, "git push*"); + } +} diff --git a/crates/harness-core/src/tool/truncate.rs b/crates/harness-core/src/tool/truncate.rs new file mode 100644 index 0000000..c4a93b9 --- /dev/null +++ b/crates/harness-core/src/tool/truncate.rs @@ -0,0 +1,75 @@ +use std::path::Path; + +/// docs/05-tools.md: cap tool output at 30,000 chars, keep head + tail, spill the full +/// text to `/tool-output/.txt` so a subagent can Grep/Read it later. +pub const MAX_CHARS: usize = 30_000; +const HEAD_CHARS: usize = MAX_CHARS / 2; +const TAIL_CHARS: usize = MAX_CHARS - HEAD_CHARS; + +pub struct Truncated { + pub text: String, + pub truncated: bool, +} + +/// `data_dir` is the session's tool-output directory; `call_id` names the spill file. +pub fn truncate(text: &str, data_dir: &Path, call_id: &str) -> std::io::Result { + if text.chars().count() <= MAX_CHARS { + return Ok(Truncated { + text: text.to_string(), + truncated: false, + }); + } + + let chars: Vec = text.chars().collect(); + let head: String = chars[..HEAD_CHARS].iter().collect(); + let tail: String = chars[chars.len() - TAIL_CHARS..].iter().collect(); + let hidden = chars.len() - HEAD_CHARS - TAIL_CHARS; + + std::fs::create_dir_all(data_dir)?; + let spill_path = data_dir.join(format!("{call_id}.txt")); + std::fs::write(&spill_path, text)?; + + let message = format!( + "{head}\n... [{hidden} chars truncated; full output: {} ] ...\n{tail}", + spill_path.display() + ); + Ok(Truncated { + text: message, + truncated: true, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn short_output_is_unchanged() { + let dir = tempfile::tempdir().unwrap(); + let result = truncate("hello world", dir.path(), "call_1").unwrap(); + assert!(!result.truncated); + assert_eq!(result.text, "hello world"); + assert!(!dir.path().join("call_1.txt").exists()); + } + + #[test] + fn long_output_is_truncated_and_spilled() { + let dir = tempfile::tempdir().unwrap(); + let long = "a".repeat(MAX_CHARS + 1000); + let result = truncate(&long, dir.path(), "call_2").unwrap(); + assert!(result.truncated); + assert!(result.text.contains("chars truncated")); + assert!(result.text.len() < long.len()); + + let spilled = std::fs::read_to_string(dir.path().join("call_2.txt")).unwrap(); + assert_eq!(spilled, long); + } + + #[test] + fn boundary_exactly_at_max_is_not_truncated() { + let dir = tempfile::tempdir().unwrap(); + let text = "a".repeat(MAX_CHARS); + let result = truncate(&text, dir.path(), "call_3").unwrap(); + assert!(!result.truncated); + } +}