Compare commits
5 Commits
420f00494f
...
d83f8c84b7
| Author | SHA1 | Date | |
|---|---|---|---|
| d83f8c84b7 | |||
| 9ed278bcb5 | |||
| 4118279da1 | |||
| a5af873924 | |||
| 5ac646b3c6 |
Generated
+1375
-11
File diff suppressed because it is too large
Load Diff
@@ -57,6 +57,9 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
|
||||
tracing-appender = "0.2"
|
||||
dirs = "5"
|
||||
shell-words = "1"
|
||||
libc = "0.2"
|
||||
tempfile = "3"
|
||||
bytes = "1"
|
||||
|
||||
[workspace.lints.rust]
|
||||
unsafe_code = "deny"
|
||||
|
||||
@@ -10,6 +10,13 @@ harness-providers = { workspace = true }
|
||||
harness-tools = { workspace = true }
|
||||
harness-mcp = { workspace = true }
|
||||
harness-lsp = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
dirs = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -1 +1,219 @@
|
||||
// Composition root: builds the Engine and registers providers/tools/mcp/lsp. Lands in M1+.
|
||||
//! Composition root: builds the pieces `harness-tui` (and later `harness-server`) need —
|
||||
//! config, store, event bus, permission service, tool registry, provider registry — and
|
||||
//! exposes a small headless API (`run_prompt`) used by the `harness run -p` debug command.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use harness_core::config::{self, Config, ConfigError};
|
||||
use harness_core::engine::{run_session, RunConfig, StepContext};
|
||||
use harness_core::event::{EventBus, RunOutcome};
|
||||
use harness_core::permission::{spawn_auto_approve, PermissionService};
|
||||
use harness_core::store::{Store, StoreError};
|
||||
use harness_core::tool::ToolRegistry;
|
||||
use harness_core::types::{Message, ModelRef, Part, PartBody, PartId, Session, SessionId};
|
||||
use harness_providers::{AnthropicProvider, ProviderRegistry};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
/// Placeholder until M4's markdown agent registry lands (`assets/agents/orchestrator.md`).
|
||||
pub const DEFAULT_AGENT_PROMPT: &str =
|
||||
"You are a helpful coding assistant with access to tools for reading, editing, running \
|
||||
commands, and searching code. Use them as needed to satisfy the user's request, then \
|
||||
reply with a concise final answer.";
|
||||
|
||||
const DEFAULT_MAX_STEPS: u32 = 50;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum AppError {
|
||||
#[error(transparent)]
|
||||
Config(#[from] ConfigError),
|
||||
#[error(transparent)]
|
||||
Store(#[from] StoreError),
|
||||
#[error("model ref must be \"provider/model\", got {0:?}")]
|
||||
InvalidModelRef(String),
|
||||
#[error("no provider registered for {0:?} (missing API key?)")]
|
||||
UnknownProvider(String),
|
||||
}
|
||||
|
||||
fn now_ms() -> i64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis() as i64
|
||||
}
|
||||
|
||||
pub struct App {
|
||||
pub config: Config,
|
||||
pub store: Store,
|
||||
pub bus: EventBus,
|
||||
pub permissions: Arc<PermissionService>,
|
||||
pub tools: ToolRegistry,
|
||||
pub providers: ProviderRegistry,
|
||||
pub cwd: PathBuf,
|
||||
data_dir: PathBuf,
|
||||
}
|
||||
|
||||
impl App {
|
||||
/// In-memory store, auto-approve permission frontend — the M1 headless configuration.
|
||||
/// A persistent SQLite-backed store and a real TUI permission modal land in M2.
|
||||
pub fn init(cwd: PathBuf) -> Result<Self, AppError> {
|
||||
let config = config::load(&cwd)?;
|
||||
let bus = EventBus::new();
|
||||
let store = Store::open_in_memory()?;
|
||||
let permissions = Arc::new(PermissionService::new(bus.clone()));
|
||||
spawn_auto_approve(bus.clone(), permissions.clone());
|
||||
|
||||
let mut tools = ToolRegistry::new();
|
||||
harness_tools::register_builtins(&mut tools);
|
||||
|
||||
let mut providers = ProviderRegistry::new();
|
||||
if let Some(key) = config
|
||||
.providers
|
||||
.get("anthropic")
|
||||
.and_then(|p| p.api_key.clone())
|
||||
{
|
||||
providers.register(Arc::new(AnthropicProvider::new(key)));
|
||||
}
|
||||
|
||||
let data_dir = dirs::data_dir()
|
||||
.unwrap_or_else(std::env::temp_dir)
|
||||
.join("ai-harness")
|
||||
.join("tool-output");
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
store,
|
||||
bus,
|
||||
permissions,
|
||||
tools,
|
||||
providers,
|
||||
cwd,
|
||||
data_dir,
|
||||
})
|
||||
}
|
||||
|
||||
/// Runs a single headless turn: creates a root session, appends `prompt` as the user
|
||||
/// message, and drives the engine loop to completion. Returns the outcome plus the
|
||||
/// session id so the caller can fetch the transcript via `final_text`.
|
||||
pub async fn run_prompt(
|
||||
&self,
|
||||
prompt: String,
|
||||
model_ref: &str,
|
||||
) -> Result<(RunOutcome, SessionId), AppError> {
|
||||
let (provider_id, model_id) = model_ref
|
||||
.split_once('/')
|
||||
.ok_or_else(|| AppError::InvalidModelRef(model_ref.to_string()))?;
|
||||
let provider = self
|
||||
.providers
|
||||
.get(provider_id)
|
||||
.ok_or_else(|| AppError::UnknownProvider(provider_id.to_string()))?;
|
||||
let model = ModelRef::new(provider_id, model_id);
|
||||
|
||||
let now = now_ms();
|
||||
let session = Session::new_root("orchestrator", model.clone(), now);
|
||||
let session_id = session.id.clone();
|
||||
self.store.upsert_session(session).await?;
|
||||
|
||||
let user_message = Message::new_user(session_id.clone(), now);
|
||||
self.store.upsert_message(user_message.clone()).await?;
|
||||
self.store
|
||||
.upsert_part(Part {
|
||||
id: PartId::new(),
|
||||
message_id: user_message.id,
|
||||
session_id: session_id.clone(),
|
||||
idx: 0,
|
||||
body: PartBody::Text {
|
||||
text: prompt,
|
||||
synthetic: false,
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
let ctx = StepContext {
|
||||
store: self.store.clone(),
|
||||
bus: self.bus.clone(),
|
||||
tools: self.tools.clone(),
|
||||
permissions: self.permissions.clone(),
|
||||
static_rules: self.config.permissions.clone(),
|
||||
extra_rules: Arc::new(Mutex::new(Vec::new())),
|
||||
session_id: session_id.clone(),
|
||||
cwd: self.cwd.clone(),
|
||||
data_dir: self.data_dir.join(session_id.to_string()),
|
||||
cancel: CancellationToken::new(),
|
||||
now,
|
||||
};
|
||||
let run_config = RunConfig {
|
||||
agent_name: "orchestrator".to_string(),
|
||||
agent_prompt: DEFAULT_AGENT_PROMPT.to_string(),
|
||||
model,
|
||||
temperature: None,
|
||||
max_steps: DEFAULT_MAX_STEPS,
|
||||
instructions: self.config.instructions.clone(),
|
||||
};
|
||||
|
||||
let outcome = run_session(provider, ctx, &run_config, now_ms).await;
|
||||
Ok((outcome, session_id))
|
||||
}
|
||||
|
||||
/// Concatenates the `Text` parts of the last message in the session — the final
|
||||
/// assistant reply for a `harness run` invocation to print.
|
||||
pub async fn final_text(&self, session_id: &SessionId) -> Result<String, AppError> {
|
||||
let messages = self.store.messages(session_id.clone()).await?;
|
||||
let Some(last) = messages.last() else {
|
||||
return Ok(String::new());
|
||||
};
|
||||
let parts = self.store.parts(last.id.clone()).await?;
|
||||
let text = parts
|
||||
.iter()
|
||||
.filter_map(|p| match &p.body {
|
||||
PartBody::Text { text, .. } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
Ok(text)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn init_loads_defaults_with_no_providers_configured() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::env::remove_var("ANTHROPIC_API_KEY");
|
||||
let app = App::init(dir.path().to_path_buf()).unwrap();
|
||||
assert!(app.providers.get("anthropic").is_none());
|
||||
assert_eq!(app.tools.all().len(), 6);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_prompt_errors_on_unregistered_provider() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::env::remove_var("ANTHROPIC_API_KEY");
|
||||
let app = App::init(dir.path().to_path_buf()).unwrap();
|
||||
let err = app
|
||||
.run_prompt("hi".into(), "anthropic/claude-sonnet-4-5")
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, AppError::UnknownProvider(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_prompt_errors_on_malformed_model_ref() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let app = App::init(dir.path().to_path_buf()).unwrap();
|
||||
let err = app.run_prompt("hi".into(), "no-slash").await.unwrap_err();
|
||||
assert!(matches!(err, AppError::InvalidModelRef(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn final_text_is_empty_for_unknown_session() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let app = App::init(dir.path().to_path_buf()).unwrap();
|
||||
let text = app.final_text(&SessionId::new()).await.unwrap();
|
||||
assert_eq!(text, "");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.<id>.api_key` when unset.
|
||||
const PROVIDER_ENV_KEYS: &[(&str, &str)] = &[
|
||||
("anthropic", "ANTHROPIC_API_KEY"),
|
||||
("openai", "OPENAI_API_KEY"),
|
||||
];
|
||||
|
||||
fn read_jsonc(path: &Path) -> Result<Value, ConfigError> {
|
||||
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<PathBuf> {
|
||||
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<PathBuf> {
|
||||
let mut found: Vec<PathBuf> = 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<Config, ConfigError> {
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
};
|
||||
@@ -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<bool>,
|
||||
pub api_key: Option<String>,
|
||||
pub base_url: Option<String>,
|
||||
}
|
||||
|
||||
/// 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<bool>,
|
||||
pub mode: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub temperature: Option<f32>,
|
||||
pub prompt: Option<String>,
|
||||
pub tools: Option<HashMap<String, bool>>,
|
||||
pub permission: Option<Vec<Rule>>,
|
||||
}
|
||||
|
||||
#[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<HashMap<String, String>>,
|
||||
}
|
||||
|
||||
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<String>,
|
||||
pub env: HashMap<String, String>,
|
||||
pub enabled: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct LspServerConfig {
|
||||
pub command: String,
|
||||
pub args: Vec<String>,
|
||||
pub extensions: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct TuiConfig {
|
||||
pub theme: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
#[serde(default)]
|
||||
pub struct Config {
|
||||
pub model: Option<String>,
|
||||
pub small_model: Option<String>,
|
||||
pub providers: HashMap<String, ProviderConfig>,
|
||||
pub agents: HashMap<String, AgentPatch>,
|
||||
pub permissions: Vec<Rule>,
|
||||
pub orchestration: OrchestrationConfig,
|
||||
pub mcp: HashMap<String, McpServerConfig>,
|
||||
pub lsp: HashMap<String, LspServerConfig>,
|
||||
pub tui: TuiConfig,
|
||||
pub instructions: Vec<String>,
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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<f32>,
|
||||
pub max_steps: u32,
|
||||
pub instructions: Vec<String>,
|
||||
}
|
||||
|
||||
/// 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<bool, ProviderError> {
|
||||
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<WireMessage> {
|
||||
match message.role {
|
||||
Role::User => {
|
||||
let content: Vec<WireContent> = 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<dyn Provider>,
|
||||
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<ToolSchema> = 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<VecDeque<Vec<Result<LlmEvent, ProviderError>>>>,
|
||||
}
|
||||
|
||||
impl MockProvider {
|
||||
fn scripted(steps: Vec<Vec<Result<LlmEvent, ProviderError>>>) -> 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<Vec<ModelInfo>, ProviderError> {
|
||||
Ok(vec![])
|
||||
}
|
||||
async fn stream(
|
||||
&self,
|
||||
_req: LlmRequest,
|
||||
_cancel: CancellationToken,
|
||||
) -> Result<LlmEventStream, ProviderError> {
|
||||
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<ToolOutput, ToolError> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -0,0 +1,657 @@
|
||||
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<MessageId>,
|
||||
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<PermissionService>,
|
||||
pub static_rules: Ruleset,
|
||||
pub extra_rules: Arc<Mutex<Ruleset>>,
|
||||
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<Message>,
|
||||
next_idx: u32,
|
||||
active_text: Option<PartId>,
|
||||
active_reasoning: Option<PartId>,
|
||||
pending_tools: HashMap<String, PartId>,
|
||||
text_buf: HashMap<PartId, String>,
|
||||
reasoning_buf: HashMap<PartId, String>,
|
||||
tool_input_buf: HashMap<String, String>,
|
||||
flushers: HashMap<PartId, FlushTracker>,
|
||||
parts_by_id: HashMap<PartId, Part>,
|
||||
}
|
||||
|
||||
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<String>,
|
||||
) -> 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,
|
||||
}
|
||||
};
|
||||
|
||||
// docs/05-tools.md: cap tool output at 30k chars regardless of which tool produced it.
|
||||
let output = crate::tool::truncate::truncate(&output, &self.ctx.data_dir, &call_id)
|
||||
.map(|t| t.text)
|
||||
.unwrap_or(output);
|
||||
|
||||
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<dyn Tool>,
|
||||
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<StepResult, ProviderError> {
|
||||
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<StepOutcome, StepError> {
|
||||
let mut run = Run::new(ctx);
|
||||
let mut usage = TokenUsage::default();
|
||||
let mut result = StepResult::Stop;
|
||||
|
||||
loop {
|
||||
let item = tokio::select! {
|
||||
item = stream.next() => item,
|
||||
_ = ctx.cancel.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 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,
|
||||
})
|
||||
}
|
||||
@@ -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<F, Fut, T>(
|
||||
cancel: &CancellationToken,
|
||||
mut attempt_fn: F,
|
||||
) -> Result<T, ProviderError>
|
||||
where
|
||||
F: FnMut() -> Fut,
|
||||
Fut: Future<Output = Result<T, ProviderError>>,
|
||||
{
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -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<String>`).
|
||||
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<String> {
|
||||
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"]);
|
||||
}
|
||||
}
|
||||
@@ -24,7 +24,8 @@ pub struct PermissionRequest {
|
||||
pub id: String,
|
||||
pub session_id: SessionId,
|
||||
pub permission: String,
|
||||
pub patterns: Vec<String>,
|
||||
pub pattern: String,
|
||||
pub always_pattern: String,
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -1,7 +1,13 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
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<String>,
|
||||
},
|
||||
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<Duration>,
|
||||
},
|
||||
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<Box<dyn Stream<Item = Result<LlmEvent, ProviderError>> + 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<WireContent>,
|
||||
}
|
||||
|
||||
#[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<ReasoningEffort>,
|
||||
/// Anthropic extended-thinking token budget.
|
||||
pub budget_tokens: Option<u32>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
pub messages: Vec<WireMessage>,
|
||||
pub tools: Vec<ToolSchema>,
|
||||
pub temperature: Option<f32>,
|
||||
pub max_tokens: Option<u32>,
|
||||
pub reasoning: Option<ReasoningOpts>,
|
||||
pub initiator: Initiator,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait Provider: Send + Sync {
|
||||
fn id(&self) -> &str;
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError>;
|
||||
async fn stream(
|
||||
&self,
|
||||
req: LlmRequest,
|
||||
cancel: CancellationToken,
|
||||
) -> Result<LlmEventStream, ProviderError>;
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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<HashMap<String, oneshot::Sender<PermissionReply>>>,
|
||||
}
|
||||
|
||||
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<AskDecision, AskError> {
|
||||
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<AskDecision, AskError> {
|
||||
self.ask_user(session_id, input, cancel).await
|
||||
}
|
||||
|
||||
async fn ask_user(
|
||||
&self,
|
||||
session_id: &SessionId,
|
||||
input: AskInput,
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<AskDecision, AskError> {
|
||||
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<PermissionService>) {
|
||||
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));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,296 @@
|
||||
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<AskError> 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<serde_json::Value>);
|
||||
|
||||
impl MetadataSink {
|
||||
pub fn channel() -> (
|
||||
Self,
|
||||
tokio::sync::mpsc::UnboundedReceiver<serde_json::Value>,
|
||||
) {
|
||||
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<PermissionService>,
|
||||
session_id: SessionId,
|
||||
static_rules: Ruleset,
|
||||
extra_rules: Arc<Mutex<Ruleset>>,
|
||||
cancel: CancellationToken,
|
||||
}
|
||||
|
||||
impl PermissionHandle {
|
||||
pub fn new(
|
||||
service: Arc<PermissionService>,
|
||||
session_id: SessionId,
|
||||
static_rules: Ruleset,
|
||||
extra_rules: Arc<Mutex<Ruleset>>,
|
||||
cancel: CancellationToken,
|
||||
) -> Self {
|
||||
Self {
|
||||
service,
|
||||
session_id,
|
||||
static_rules,
|
||||
extra_rules,
|
||||
cancel,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn ask(
|
||||
&self,
|
||||
permission: impl Into<String>,
|
||||
pattern: impl Into<String>,
|
||||
always_pattern: impl Into<String>,
|
||||
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,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ToolOutput {
|
||||
pub title: String,
|
||||
pub output: String,
|
||||
pub metadata: serde_json::Value,
|
||||
}
|
||||
|
||||
impl ToolOutput {
|
||||
pub fn new(title: impl Into<String>, output: impl Into<String>) -> 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<ToolOutput, ToolError>;
|
||||
}
|
||||
|
||||
/// 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<String, Arc<dyn Tool>>,
|
||||
}
|
||||
|
||||
impl ToolRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn register(&mut self, tool: Arc<dyn Tool>) {
|
||||
self.tools.insert(tool.name().to_string(), tool);
|
||||
}
|
||||
|
||||
pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
|
||||
self.tools.get(name).cloned()
|
||||
}
|
||||
|
||||
pub fn all(&self) -> Vec<Arc<dyn Tool>> {
|
||||
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<ToolOutput, ToolError> {
|
||||
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*");
|
||||
}
|
||||
}
|
||||
@@ -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 `<data_dir>/tool-output/<call_id>.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<Truncated> {
|
||||
if text.chars().count() <= MAX_CHARS {
|
||||
return Ok(Truncated {
|
||||
text: text.to_string(),
|
||||
truncated: false,
|
||||
});
|
||||
}
|
||||
|
||||
let chars: Vec<char> = 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);
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,21 @@ license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
harness-core = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
async-stream = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
reqwest = { workspace = true }
|
||||
eventsource-stream = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["test-util", "macros"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use harness_core::llm::{LlmEventStream, LlmRequest, Provider, ProviderError};
|
||||
use harness_core::types::ModelInfo;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::codec::anthropic::{build_request, decode, initiator_header};
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
|
||||
const ANTHROPIC_VERSION: &str = "2023-06-01";
|
||||
const ANTHROPIC_BETA: &str =
|
||||
"interleaved-thinking-2025-05-14,fine-grained-tool-streaming-2025-05-14";
|
||||
|
||||
pub struct AnthropicProvider {
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
client: reqwest::Client,
|
||||
}
|
||||
|
||||
impl AnthropicProvider {
|
||||
pub fn new(api_key: impl Into<String>) -> Self {
|
||||
Self {
|
||||
api_key: api_key.into(),
|
||||
base_url: DEFAULT_BASE_URL.to_string(),
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn with_base_url(api_key: impl Into<String>, base_url: impl Into<String>) -> Self {
|
||||
Self {
|
||||
api_key: api_key.into(),
|
||||
base_url: base_url.into(),
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_error(
|
||||
status: reqwest::StatusCode,
|
||||
body: String,
|
||||
retry_after: Option<Duration>,
|
||||
) -> ProviderError {
|
||||
match status.as_u16() {
|
||||
401 | 403 => ProviderError::Auth(body),
|
||||
429 => ProviderError::RateLimited { retry_after },
|
||||
529 => ProviderError::Overloaded,
|
||||
s if (500..600).contains(&s) => ProviderError::Overloaded,
|
||||
s => ProviderError::Http { status: s, body },
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for AnthropicProvider {
|
||||
fn id(&self) -> &str {
|
||||
"anthropic"
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
||||
// Full models.dev metadata integration lands in M3; callers fall back to
|
||||
// config-specified model strings until then.
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
req: LlmRequest,
|
||||
cancel: CancellationToken,
|
||||
) -> Result<LlmEventStream, ProviderError> {
|
||||
let body = build_request(&req);
|
||||
let url = format!("{}/v1/messages", self.base_url);
|
||||
let initiator = initiator_header(req.initiator);
|
||||
|
||||
let send = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", ANTHROPIC_VERSION)
|
||||
.header("anthropic-beta", ANTHROPIC_BETA)
|
||||
.header("x-initiator", initiator)
|
||||
.json(&body)
|
||||
.send();
|
||||
|
||||
let response = tokio::select! {
|
||||
result = send => result.map_err(|e| ProviderError::Network(e.to_string()))?,
|
||||
_ = cancel.cancelled() => return Err(ProviderError::Cancelled),
|
||||
};
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let retry_after = response
|
||||
.headers()
|
||||
.get("retry-after")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.map(Duration::from_secs);
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
return Err(Self::classify_error(status, body, retry_after));
|
||||
}
|
||||
|
||||
Ok(decode(response.bytes_stream()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn classifies_status_codes() {
|
||||
assert!(matches!(
|
||||
AnthropicProvider::classify_error(
|
||||
reqwest::StatusCode::TOO_MANY_REQUESTS,
|
||||
String::new(),
|
||||
None
|
||||
),
|
||||
ProviderError::RateLimited { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
AnthropicProvider::classify_error(
|
||||
reqwest::StatusCode::UNAUTHORIZED,
|
||||
String::new(),
|
||||
None
|
||||
),
|
||||
ProviderError::Auth(_)
|
||||
));
|
||||
assert!(matches!(
|
||||
AnthropicProvider::classify_error(
|
||||
reqwest::StatusCode::SERVICE_UNAVAILABLE,
|
||||
String::new(),
|
||||
None
|
||||
),
|
||||
ProviderError::Overloaded
|
||||
));
|
||||
assert!(matches!(
|
||||
AnthropicProvider::classify_error(
|
||||
reqwest::StatusCode::BAD_REQUEST,
|
||||
String::new(),
|
||||
None
|
||||
),
|
||||
ProviderError::Http { status: 400, .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn id_is_anthropic() {
|
||||
let provider = AnthropicProvider::new("test-key");
|
||||
assert_eq!(provider.id(), "anthropic");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn constructs_with_custom_base_url() {
|
||||
let provider = AnthropicProvider::with_base_url("test-key", "http://localhost:1234");
|
||||
assert_eq!(provider.base_url, "http://localhost:1234");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,499 @@
|
||||
//! Request builder + SSE decoder for Anthropic's `/v1/messages` streaming API.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_stream::try_stream;
|
||||
use eventsource_stream::Eventsource;
|
||||
use futures::Stream;
|
||||
use harness_core::llm::{
|
||||
FinishReason, Initiator, LlmEvent, LlmEventStream, LlmRequest, ProviderError, Role, WireContent,
|
||||
};
|
||||
use harness_core::types::TokenUsage;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
const DEFAULT_MAX_TOKENS: u32 = 8192;
|
||||
/// opencode's `applyCaching`: cache breakpoints on the first 2 system blocks + last 2
|
||||
/// non-system messages.
|
||||
const CACHE_BREAKPOINTS: usize = 2;
|
||||
|
||||
fn cache_control() -> Value {
|
||||
json!({"type": "ephemeral"})
|
||||
}
|
||||
|
||||
fn build_system(system: &[String]) -> Vec<Value> {
|
||||
system
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, block)| {
|
||||
let mut b = json!({"type": "text", "text": block});
|
||||
if i < CACHE_BREAKPOINTS {
|
||||
b["cache_control"] = cache_control();
|
||||
}
|
||||
b
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn wire_role_to_anthropic(role: Role) -> &'static str {
|
||||
match role {
|
||||
Role::User | Role::Tool => "user",
|
||||
Role::Assistant => "assistant",
|
||||
Role::System => "user", // system blocks are carried separately in `system`, not here
|
||||
}
|
||||
}
|
||||
|
||||
fn build_messages(messages: &[harness_core::llm::WireMessage]) -> Vec<Value> {
|
||||
let mut built: Vec<Value> = messages
|
||||
.iter()
|
||||
.map(|m| {
|
||||
let content: Vec<Value> = m
|
||||
.content
|
||||
.iter()
|
||||
.map(|c| match c {
|
||||
WireContent::Text { text } => json!({"type": "text", "text": text}),
|
||||
WireContent::ToolCall { call_id, name, input } => {
|
||||
json!({"type": "tool_use", "id": call_id, "name": name, "input": input})
|
||||
}
|
||||
WireContent::ToolResult { call_id, output, is_error } => {
|
||||
json!({"type": "tool_result", "tool_use_id": call_id, "content": output, "is_error": is_error})
|
||||
}
|
||||
WireContent::Image { mime_type, data } => {
|
||||
json!({"type": "image", "source": {"type": "base64", "media_type": mime_type, "data": data}})
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
json!({"role": wire_role_to_anthropic(m.role), "content": content})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let n = built.len();
|
||||
for msg in built.iter_mut().skip(n.saturating_sub(CACHE_BREAKPOINTS)) {
|
||||
if let Some(content) = msg["content"].as_array_mut() {
|
||||
if let Some(last) = content.last_mut() {
|
||||
last["cache_control"] = cache_control();
|
||||
}
|
||||
}
|
||||
}
|
||||
built
|
||||
}
|
||||
|
||||
pub fn build_request(req: &LlmRequest) -> Value {
|
||||
let mut body = json!({
|
||||
"model": req.model,
|
||||
"system": build_system(&req.system),
|
||||
"messages": build_messages(&req.messages),
|
||||
"max_tokens": req.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS),
|
||||
"stream": true,
|
||||
});
|
||||
|
||||
if !req.tools.is_empty() {
|
||||
body["tools"] = Value::Array(
|
||||
req.tools
|
||||
.iter()
|
||||
.map(|t| json!({"name": t.name, "description": t.description, "input_schema": t.parameters}))
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
if let Some(temp) = req.temperature {
|
||||
body["temperature"] = json!(temp);
|
||||
}
|
||||
if let Some(reasoning) = &req.reasoning {
|
||||
if let Some(budget) = reasoning.budget_tokens {
|
||||
body["thinking"] = json!({"type": "enabled", "budget_tokens": budget});
|
||||
}
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
pub fn initiator_header(initiator: Initiator) -> &'static str {
|
||||
match initiator {
|
||||
Initiator::User => "user",
|
||||
Initiator::Agent => "agent",
|
||||
}
|
||||
}
|
||||
|
||||
fn map_stop_reason(reason: Option<&str>) -> FinishReason {
|
||||
match reason {
|
||||
Some("end_turn") | Some("stop_sequence") => FinishReason::Stop,
|
||||
Some("tool_use") => FinishReason::ToolCalls,
|
||||
Some("max_tokens") => FinishReason::Length,
|
||||
Some(other) => FinishReason::Unknown(other.to_string()),
|
||||
None => FinishReason::Unknown("none".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct BlockState {
|
||||
kind: BlockKind,
|
||||
call_id: String,
|
||||
name: String,
|
||||
partial_json: String,
|
||||
signature: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Default, PartialEq, Eq)]
|
||||
enum BlockKind {
|
||||
#[default]
|
||||
Text,
|
||||
Thinking,
|
||||
ToolUse,
|
||||
}
|
||||
|
||||
/// Decodes a raw SSE byte stream into our normalized `LlmEvent` stream. Errors from the
|
||||
/// underlying HTTP stream and any Anthropic `error` event both surface as `Err`.
|
||||
pub fn decode<S, E>(byte_stream: S) -> LlmEventStream
|
||||
where
|
||||
S: Stream<Item = Result<bytes::Bytes, E>> + Send + 'static,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
let events = byte_stream.eventsource();
|
||||
Box::pin(try_stream! {
|
||||
futures::pin_mut!(events);
|
||||
let mut blocks: HashMap<u64, BlockState> = HashMap::new();
|
||||
let mut usage = TokenUsage::default();
|
||||
|
||||
while let Some(item) = futures::StreamExt::next(&mut events).await {
|
||||
let event = item.map_err(|e| ProviderError::Decode(e.to_string()))?;
|
||||
if event.data.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let value: Value = serde_json::from_str(&event.data)
|
||||
.map_err(|e| ProviderError::Decode(format!("{e}: {}", event.data)))?;
|
||||
let kind = value["type"].as_str().unwrap_or_default();
|
||||
|
||||
match kind {
|
||||
"message_start" => {
|
||||
let u = &value["message"]["usage"];
|
||||
usage.input = u["input_tokens"].as_u64().unwrap_or(0);
|
||||
usage.cache_write = u["cache_creation_input_tokens"].as_u64().unwrap_or(0);
|
||||
usage.cache_read = u["cache_read_input_tokens"].as_u64().unwrap_or(0);
|
||||
}
|
||||
"content_block_start" => {
|
||||
let index = value["index"].as_u64().unwrap_or(0);
|
||||
let block = &value["content_block"];
|
||||
match block["type"].as_str().unwrap_or_default() {
|
||||
"text" => {
|
||||
blocks.insert(index, BlockState { kind: BlockKind::Text, ..Default::default() });
|
||||
yield LlmEvent::TextStart { id: index.to_string() };
|
||||
}
|
||||
"thinking" => {
|
||||
blocks.insert(index, BlockState { kind: BlockKind::Thinking, ..Default::default() });
|
||||
yield LlmEvent::ReasoningStart { id: index.to_string() };
|
||||
}
|
||||
"tool_use" => {
|
||||
let call_id = block["id"].as_str().unwrap_or_default().to_string();
|
||||
let name = block["name"].as_str().unwrap_or_default().to_string();
|
||||
blocks.insert(index, BlockState {
|
||||
kind: BlockKind::ToolUse,
|
||||
call_id: call_id.clone(),
|
||||
name: name.clone(),
|
||||
..Default::default()
|
||||
});
|
||||
yield LlmEvent::ToolInputStart { call_id, name };
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
"content_block_delta" => {
|
||||
let index = value["index"].as_u64().unwrap_or(0);
|
||||
let delta = &value["delta"];
|
||||
match delta["type"].as_str().unwrap_or_default() {
|
||||
"text_delta" => {
|
||||
let text = delta["text"].as_str().unwrap_or_default().to_string();
|
||||
yield LlmEvent::TextDelta { id: index.to_string(), text };
|
||||
}
|
||||
"thinking_delta" => {
|
||||
let text = delta["thinking"].as_str().unwrap_or_default().to_string();
|
||||
yield LlmEvent::ReasoningDelta { id: index.to_string(), text };
|
||||
}
|
||||
"signature_delta" => {
|
||||
if let Some(state) = blocks.get_mut(&index) {
|
||||
state.signature = Some(delta["signature"].as_str().unwrap_or_default().to_string());
|
||||
}
|
||||
}
|
||||
"input_json_delta" => {
|
||||
let partial = delta["partial_json"].as_str().unwrap_or_default();
|
||||
if let Some(state) = blocks.get_mut(&index) {
|
||||
state.partial_json.push_str(partial);
|
||||
yield LlmEvent::ToolInputDelta { call_id: state.call_id.clone(), json: partial.to_string() };
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
"content_block_stop" => {
|
||||
let index = value["index"].as_u64().unwrap_or(0);
|
||||
if let Some(state) = blocks.remove(&index) {
|
||||
match state.kind {
|
||||
BlockKind::Text => yield LlmEvent::TextEnd { id: index.to_string() },
|
||||
BlockKind::Thinking => {
|
||||
yield LlmEvent::ReasoningEnd { id: index.to_string(), signature: state.signature };
|
||||
}
|
||||
BlockKind::ToolUse => {
|
||||
let input: Value = if state.partial_json.trim().is_empty() {
|
||||
json!({})
|
||||
} else {
|
||||
serde_json::from_str(&state.partial_json).unwrap_or(Value::Null)
|
||||
};
|
||||
yield LlmEvent::ToolCall { call_id: state.call_id, name: state.name, input };
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"message_delta" => {
|
||||
if let Some(out) = value["usage"]["output_tokens"].as_u64() {
|
||||
usage.output = out;
|
||||
}
|
||||
let stop_reason = value["delta"]["stop_reason"].as_str();
|
||||
yield LlmEvent::Finish { reason: map_stop_reason(stop_reason), usage };
|
||||
}
|
||||
"error" => {
|
||||
let message = value["error"]["message"].as_str().unwrap_or("unknown error").to_string();
|
||||
let err_type = value["error"]["type"].as_str().unwrap_or("");
|
||||
Err(match err_type {
|
||||
"overloaded_error" => ProviderError::Overloaded,
|
||||
"rate_limit_error" => ProviderError::RateLimited { retry_after: None },
|
||||
"authentication_error" | "permission_error" => ProviderError::Auth(message),
|
||||
_ => ProviderError::Http { status: 0, body: message },
|
||||
})?;
|
||||
}
|
||||
_ => {} // ping, message_stop: nothing to emit
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use futures::StreamExt;
|
||||
use harness_core::llm::{ReasoningOpts, ToolSchema, WireMessage};
|
||||
|
||||
fn sse_stream(raw: &'static str) -> LlmEventStream {
|
||||
let chunks: Vec<Result<bytes::Bytes, std::io::Error>> =
|
||||
vec![Ok(bytes::Bytes::from_static(raw.as_bytes()))];
|
||||
decode(futures::stream::iter(chunks))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_text_only_response() {
|
||||
let raw = concat!(
|
||||
"event: message_start\n",
|
||||
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":10}}}\n\n",
|
||||
"event: content_block_start\n",
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
|
||||
"event: content_block_delta\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n",
|
||||
"event: content_block_stop\n",
|
||||
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
||||
"event: message_delta\n",
|
||||
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\n",
|
||||
"event: message_stop\n",
|
||||
"data: {\"type\":\"message_stop\"}\n\n",
|
||||
);
|
||||
let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
LlmEvent::TextStart { id: "0".into() },
|
||||
LlmEvent::TextDelta {
|
||||
id: "0".into(),
|
||||
text: "Hello".into()
|
||||
},
|
||||
LlmEvent::TextEnd { id: "0".into() },
|
||||
LlmEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
usage: TokenUsage {
|
||||
input: 10,
|
||||
output: 5,
|
||||
..Default::default()
|
||||
},
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_tool_call_with_streamed_json_input() {
|
||||
let raw = concat!(
|
||||
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_1\",\"name\":\"read\"}}\n\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"file\\\"\"}}\n\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\":\\\"a.txt\\\"}\"}}\n\n",
|
||||
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
||||
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":8}}\n\n",
|
||||
);
|
||||
let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
LlmEvent::ToolInputStart {
|
||||
call_id: "call_1".into(),
|
||||
name: "read".into()
|
||||
},
|
||||
LlmEvent::ToolInputDelta {
|
||||
call_id: "call_1".into(),
|
||||
json: "{\"file\"".into()
|
||||
},
|
||||
LlmEvent::ToolInputDelta {
|
||||
call_id: "call_1".into(),
|
||||
json: ":\"a.txt\"}".into()
|
||||
},
|
||||
LlmEvent::ToolCall {
|
||||
call_id: "call_1".into(),
|
||||
name: "read".into(),
|
||||
input: json!({"file": "a.txt"}),
|
||||
},
|
||||
LlmEvent::Finish {
|
||||
reason: FinishReason::ToolCalls,
|
||||
usage: TokenUsage {
|
||||
input: 1,
|
||||
output: 8,
|
||||
..Default::default()
|
||||
},
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_thinking_block_with_signature() {
|
||||
let raw = concat!(
|
||||
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\"}}\n\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"pondering\"}}\n\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"sig123\"}}\n\n",
|
||||
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
||||
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":2}}\n\n",
|
||||
);
|
||||
let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
LlmEvent::ReasoningStart { id: "0".into() },
|
||||
LlmEvent::ReasoningDelta {
|
||||
id: "0".into(),
|
||||
text: "pondering".into()
|
||||
},
|
||||
LlmEvent::ReasoningEnd {
|
||||
id: "0".into(),
|
||||
signature: Some("sig123".into())
|
||||
},
|
||||
LlmEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
usage: TokenUsage {
|
||||
input: 1,
|
||||
output: 2,
|
||||
..Default::default()
|
||||
},
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mid_stream_error_event_surfaces_as_err() {
|
||||
let raw = concat!(
|
||||
"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
|
||||
"data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"overloaded\"}}\n\n",
|
||||
);
|
||||
let events: Vec<Result<LlmEvent, ProviderError>> = sse_stream(raw).collect().await;
|
||||
assert!(matches!(
|
||||
events.last(),
|
||||
Some(Err(ProviderError::Overloaded))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_applies_cache_control_to_first_two_system_blocks() {
|
||||
let req = LlmRequest {
|
||||
model: "claude-sonnet".into(),
|
||||
system: vec!["env".into(), "agent".into(), "instructions".into()],
|
||||
messages: vec![],
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
reasoning: None,
|
||||
initiator: Initiator::User,
|
||||
};
|
||||
let body = build_request(&req);
|
||||
let system = body["system"].as_array().unwrap();
|
||||
assert_eq!(system.len(), 3);
|
||||
assert!(system[0]["cache_control"].is_object());
|
||||
assert!(system[1]["cache_control"].is_object());
|
||||
assert!(system[2].get("cache_control").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_applies_cache_control_to_last_two_messages() {
|
||||
let messages = vec![
|
||||
WireMessage {
|
||||
role: Role::User,
|
||||
content: vec![WireContent::Text { text: "1".into() }],
|
||||
},
|
||||
WireMessage {
|
||||
role: Role::Assistant,
|
||||
content: vec![WireContent::Text { text: "2".into() }],
|
||||
},
|
||||
WireMessage {
|
||||
role: Role::User,
|
||||
content: vec![WireContent::Text { text: "3".into() }],
|
||||
},
|
||||
];
|
||||
let req = LlmRequest {
|
||||
model: "claude-sonnet".into(),
|
||||
system: vec![],
|
||||
messages,
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
reasoning: None,
|
||||
initiator: Initiator::User,
|
||||
};
|
||||
let body = build_request(&req);
|
||||
let msgs = body["messages"].as_array().unwrap();
|
||||
assert!(msgs[0]["content"][0].get("cache_control").is_none());
|
||||
assert!(msgs[1]["content"][0]["cache_control"].is_object());
|
||||
assert!(msgs[2]["content"][0]["cache_control"].is_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_maps_tool_schema_to_input_schema_key() {
|
||||
let req = LlmRequest {
|
||||
model: "claude-sonnet".into(),
|
||||
system: vec![],
|
||||
messages: vec![],
|
||||
tools: vec![ToolSchema {
|
||||
name: "read".into(),
|
||||
description: "reads a file".into(),
|
||||
parameters: json!({"type": "object"}),
|
||||
}],
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
reasoning: None,
|
||||
initiator: Initiator::User,
|
||||
};
|
||||
let body = build_request(&req);
|
||||
assert_eq!(body["tools"][0]["name"], "read");
|
||||
assert_eq!(body["tools"][0]["input_schema"], json!({"type": "object"}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_includes_thinking_budget_when_reasoning_set() {
|
||||
let req = LlmRequest {
|
||||
model: "claude-sonnet".into(),
|
||||
system: vec![],
|
||||
messages: vec![],
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
reasoning: Some(ReasoningOpts {
|
||||
effort: None,
|
||||
budget_tokens: Some(2048),
|
||||
}),
|
||||
initiator: Initiator::User,
|
||||
};
|
||||
let body = build_request(&req);
|
||||
assert_eq!(body["thinking"]["budget_tokens"], 2048);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod anthropic;
|
||||
@@ -1 +1,6 @@
|
||||
// Provider implementations (Copilot, OpenAI, Anthropic) land here in M1/M3.
|
||||
pub mod anthropic;
|
||||
pub mod codec;
|
||||
pub mod registry;
|
||||
|
||||
pub use anthropic::AnthropicProvider;
|
||||
pub use registry::ProviderRegistry;
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use harness_core::llm::Provider;
|
||||
|
||||
#[derive(Default, Clone)]
|
||||
pub struct ProviderRegistry {
|
||||
providers: HashMap<String, Arc<dyn Provider>>,
|
||||
}
|
||||
|
||||
impl ProviderRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn register(&mut self, provider: Arc<dyn Provider>) {
|
||||
self.providers.insert(provider.id().to_string(), provider);
|
||||
}
|
||||
|
||||
pub fn get(&self, id: &str) -> Option<Arc<dyn Provider>> {
|
||||
self.providers.get(id).cloned()
|
||||
}
|
||||
|
||||
/// Resolves a `"provider/model"` string to its provider and bare model id.
|
||||
pub fn resolve(&self, model_ref: &str) -> Option<(Arc<dyn Provider>, String)> {
|
||||
let (provider_id, model_id) = model_ref.split_once('/')?;
|
||||
self.get(provider_id).map(|p| (p, model_id.to_string()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::anthropic::AnthropicProvider;
|
||||
|
||||
#[test]
|
||||
fn resolves_provider_slash_model() {
|
||||
let mut registry = ProviderRegistry::new();
|
||||
registry.register(Arc::new(AnthropicProvider::new("key")));
|
||||
let (provider, model) = registry.resolve("anthropic/claude-sonnet-4-5").unwrap();
|
||||
assert_eq!(provider.id(), "anthropic");
|
||||
assert_eq!(model, "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_provider_resolves_to_none() {
|
||||
let registry = ProviderRegistry::new();
|
||||
assert!(registry.resolve("openai/gpt-4o").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_model_ref_resolves_to_none() {
|
||||
let mut registry = ProviderRegistry::new();
|
||||
registry.register(Arc::new(AnthropicProvider::new("key")));
|
||||
assert!(registry.resolve("no-slash-here").is_none());
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,27 @@ license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
harness-core = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
schemars = { workspace = true }
|
||||
ignore = { workspace = true }
|
||||
globset = { workspace = true }
|
||||
grep-searcher = { workspace = true }
|
||||
grep-regex = { workspace = true }
|
||||
grep-matcher = { workspace = true }
|
||||
similar = { workspace = true }
|
||||
shell-words = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
libc = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
tokio = { workspace = true, features = ["test-util", "macros"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -0,0 +1,215 @@
|
||||
use std::process::Stdio;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
|
||||
const DEFAULT_TIMEOUT_MS: u64 = 2 * 60 * 1000;
|
||||
const MAX_TIMEOUT_MS: u64 = 10 * 60 * 1000;
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
struct BashParams {
|
||||
command: String,
|
||||
timeout_ms: Option<u64>,
|
||||
cwd: Option<String>,
|
||||
#[allow(dead_code)]
|
||||
description: String,
|
||||
}
|
||||
|
||||
pub struct BashTool;
|
||||
|
||||
/// The "always allow" pattern is coarser than the exact command: opencode/docs/05-tools.md
|
||||
/// grants `<first word> *` (e.g. `git *`), not the literal command string.
|
||||
fn always_pattern(command: &str) -> String {
|
||||
shell_words::split(command)
|
||||
.ok()
|
||||
.and_then(|words| words.into_iter().next())
|
||||
.map(|first| format!("{first} *"))
|
||||
.unwrap_or_else(|| "*".to_string())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[allow(unsafe_code)]
|
||||
fn kill_group(pid: u32) {
|
||||
// Negative pid signals the whole process group (spawned with `process_group(0)`).
|
||||
// Safety: `kill` is a plain libc syscall; passing a negative pid targets the group,
|
||||
// which is exactly the process tree we spawned and want to tear down.
|
||||
unsafe {
|
||||
libc::kill(-(pid as i32), libc::SIGKILL);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn kill_group(_pid: u32) {}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for BashTool {
|
||||
fn name(&self) -> &str {
|
||||
"bash"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Runs a shell command and returns its combined stdout/stderr."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::to_value(schemars::schema_for!(BashParams)).unwrap()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: BashParams =
|
||||
serde_json::from_value(input).map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
let timeout_ms = params
|
||||
.timeout_ms
|
||||
.unwrap_or(DEFAULT_TIMEOUT_MS)
|
||||
.min(MAX_TIMEOUT_MS);
|
||||
|
||||
ctx.ask
|
||||
.ask(
|
||||
"bash",
|
||||
params.command.clone(),
|
||||
always_pattern(¶ms.command),
|
||||
serde_json::json!({"command": params.command}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let cwd = params
|
||||
.cwd
|
||||
.as_ref()
|
||||
.map(|c| crate::paths::resolve(&ctx.cwd, c))
|
||||
.unwrap_or_else(|| ctx.cwd.clone());
|
||||
|
||||
let mut cmd = tokio::process::Command::new("sh");
|
||||
cmd.arg("-c")
|
||||
.arg(¶ms.command)
|
||||
.current_dir(&cwd)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
#[cfg(unix)]
|
||||
cmd.process_group(0);
|
||||
|
||||
let child = cmd
|
||||
.spawn()
|
||||
.map_err(|e| ToolError::Other(format!("spawn failed: {e}")))?;
|
||||
let pid = child.id();
|
||||
|
||||
tokio::select! {
|
||||
result = child.wait_with_output() => {
|
||||
let output = result.map_err(|e| ToolError::Other(format!("wait failed: {e}")))?;
|
||||
let mut combined = String::from_utf8_lossy(&output.stdout).into_owned();
|
||||
combined.push_str(&String::from_utf8_lossy(&output.stderr));
|
||||
let title = format!("{} (exit {})", params.command, output.status.code().unwrap_or(-1));
|
||||
Ok(ToolOutput::new(title, combined))
|
||||
}
|
||||
_ = tokio::time::sleep(Duration::from_millis(timeout_ms)) => {
|
||||
if let Some(pid) = pid { kill_group(pid); }
|
||||
Err(ToolError::Other(format!("command timed out after {timeout_ms}ms")))
|
||||
}
|
||||
_ = ctx.cancel.cancelled() => {
|
||||
if let Some(pid) = pid { kill_group(pid); }
|
||||
Err(ToolError::Cancelled)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use harness_core::event::EventBus;
|
||||
use harness_core::permission::{spawn_auto_approve, PermissionService};
|
||||
use harness_core::tool::{MetadataSink, PermissionHandle};
|
||||
use harness_core::types::SessionId;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
fn ctx(cwd: std::path::PathBuf) -> ToolCtx {
|
||||
let bus = EventBus::new();
|
||||
let service = Arc::new(PermissionService::new(bus.clone()));
|
||||
spawn_auto_approve(bus, service.clone());
|
||||
let (metadata, _rx) = MetadataSink::channel();
|
||||
ToolCtx {
|
||||
session_id: SessionId::new(),
|
||||
message_id: harness_core::types::MessageId::new(),
|
||||
call_id: "call_1".into(),
|
||||
data_dir: cwd.join("tool-output"),
|
||||
cwd,
|
||||
cancel: CancellationToken::new(),
|
||||
ask: PermissionHandle::new(
|
||||
service,
|
||||
SessionId::new(),
|
||||
Vec::new(),
|
||||
Arc::new(Mutex::new(Vec::new())),
|
||||
CancellationToken::new(),
|
||||
),
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn always_pattern_uses_first_word() {
|
||||
assert_eq!(always_pattern("git push origin main"), "git *");
|
||||
assert_eq!(always_pattern("ls -la"), "ls *");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn runs_command_and_captures_stdout() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let output = BashTool
|
||||
.execute(
|
||||
serde_json::json!({"command": "echo hello", "description": "say hi"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.output.contains("hello"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn captures_stderr_too() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let output = BashTool
|
||||
.execute(
|
||||
serde_json::json!({"command": "echo err 1>&2", "description": "stderr"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.output.contains("err"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn times_out_long_running_commands() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let err = BashTool
|
||||
.execute(
|
||||
serde_json::json!({"command": "sleep 5", "timeout_ms": 50, "description": "slow"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, ToolError::Other(_)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn respects_cwd_override() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let output = BashTool
|
||||
.execute(
|
||||
serde_json::json!({"command": "pwd", "description": "where"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output
|
||||
.output
|
||||
.trim()
|
||||
.ends_with(dir.path().file_name().unwrap().to_str().unwrap()));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,467 @@
|
||||
mod replacers;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex as StdMutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
use similar::{ChangeTag, TextDiff};
|
||||
use tokio::sync::Mutex as AsyncMutex;
|
||||
|
||||
use crate::paths;
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
struct EditParams {
|
||||
file_path: String,
|
||||
old_string: String,
|
||||
new_string: String,
|
||||
replace_all: Option<bool>,
|
||||
}
|
||||
|
||||
/// Per-file `tokio::Mutex` lock map so concurrent edits to different files proceed in
|
||||
/// parallel while edits to the *same* file serialize (docs/05-tools.md).
|
||||
pub struct EditTool {
|
||||
locks: StdMutex<HashMap<PathBuf, Arc<AsyncMutex<()>>>>,
|
||||
}
|
||||
|
||||
impl EditTool {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
locks: StdMutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn lock_for(&self, path: &Path) -> Arc<AsyncMutex<()>> {
|
||||
self.locks
|
||||
.lock()
|
||||
.unwrap()
|
||||
.entry(path.to_path_buf())
|
||||
.or_insert_with(|| Arc::new(AsyncMutex::new(())))
|
||||
.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for EditTool {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_line_endings(text: &str) -> String {
|
||||
text.replace("\r\n", "\n")
|
||||
}
|
||||
|
||||
fn detect_line_ending(text: &str) -> &'static str {
|
||||
if text.contains("\r\n") {
|
||||
"\r\n"
|
||||
} else {
|
||||
"\n"
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_line_ending(text: &str, ending: &str) -> String {
|
||||
if ending == "\n" {
|
||||
text.to_string()
|
||||
} else {
|
||||
text.replace('\n', "\r\n")
|
||||
}
|
||||
}
|
||||
|
||||
const UTF8_BOM: [u8; 3] = [0xEF, 0xBB, 0xBF];
|
||||
|
||||
fn strip_bom(bytes: &[u8]) -> (bool, &[u8]) {
|
||||
if bytes.starts_with(&UTF8_BOM) {
|
||||
(true, &bytes[3..])
|
||||
} else {
|
||||
(false, bytes)
|
||||
}
|
||||
}
|
||||
|
||||
fn unified_diff(before: &str, after: &str, file_path: &str) -> String {
|
||||
TextDiff::from_lines(before, after)
|
||||
.unified_diff()
|
||||
.header(&format!("a/{file_path}"), &format!("b/{file_path}"))
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn diff_stats(before: &str, after: &str) -> (usize, usize) {
|
||||
let diff = TextDiff::from_lines(before, after);
|
||||
let mut added = 0;
|
||||
let mut removed = 0;
|
||||
for change in diff.iter_all_changes() {
|
||||
match change.tag() {
|
||||
ChangeTag::Insert => added += 1,
|
||||
ChangeTag::Delete => removed += 1,
|
||||
ChangeTag::Equal => {}
|
||||
}
|
||||
}
|
||||
(added, removed)
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for EditTool {
|
||||
fn name(&self) -> &str {
|
||||
"edit"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Replaces an exact span of text in a file (fuzzy-matched against whitespace/indentation drift)."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::to_value(schemars::schema_for!(EditParams)).unwrap()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: EditParams =
|
||||
serde_json::from_value(input).map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
|
||||
if params.old_string == params.new_string {
|
||||
return Err(ToolError::Invalid(
|
||||
"No changes to apply: oldString and newString are identical.".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let path = paths::resolve(&ctx.cwd, ¶ms.file_path);
|
||||
let pattern = paths::relative_pattern(&ctx.cwd, &path);
|
||||
let lock = self.lock_for(&path);
|
||||
let _guard = lock.lock().await;
|
||||
|
||||
let metadata = tokio::fs::metadata(&path).await.ok();
|
||||
if metadata.as_ref().is_some_and(|m| m.is_dir()) {
|
||||
return Err(ToolError::Invalid(format!(
|
||||
"Path is a directory, not a file: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
|
||||
let (content_old, content_new, had_bom) = if params.old_string.is_empty() {
|
||||
if metadata.is_some() {
|
||||
return Err(ToolError::Invalid(
|
||||
"oldString cannot be empty when editing an existing file. Provide the exact text to replace, or use write for an intentional full-file replacement.".into(),
|
||||
));
|
||||
}
|
||||
(String::new(), params.new_string.clone(), false)
|
||||
} else {
|
||||
if metadata.is_none() {
|
||||
return Err(ToolError::Invalid(format!(
|
||||
"File {} not found",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
let bytes = tokio::fs::read(&path)
|
||||
.await
|
||||
.map_err(|e| ToolError::Other(format!("{}: {e}", path.display())))?;
|
||||
let (had_bom, stripped) = strip_bom(&bytes);
|
||||
let text = String::from_utf8_lossy(stripped).into_owned();
|
||||
|
||||
let ending = detect_line_ending(&text);
|
||||
let old = convert_line_ending(&normalize_line_endings(¶ms.old_string), ending);
|
||||
let new_ = convert_line_ending(&normalize_line_endings(¶ms.new_string), ending);
|
||||
let replaced =
|
||||
replacers::replace(&text, &old, &new_, params.replace_all.unwrap_or(false))
|
||||
.map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
(text, replaced, had_bom)
|
||||
};
|
||||
|
||||
let diff = unified_diff(&content_old, &content_new, ¶ms.file_path);
|
||||
ctx.ask
|
||||
.ask(
|
||||
"edit",
|
||||
pattern.clone(),
|
||||
"*",
|
||||
serde_json::json!({"file_path": params.file_path, "diff": diff}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
if let Some(parent) = path.parent() {
|
||||
tokio::fs::create_dir_all(parent)
|
||||
.await
|
||||
.map_err(|e| ToolError::Other(format!("{}: {e}", parent.display())))?;
|
||||
}
|
||||
let mut out_bytes = Vec::with_capacity(content_new.len() + 3);
|
||||
if had_bom {
|
||||
out_bytes.extend_from_slice(&UTF8_BOM);
|
||||
}
|
||||
out_bytes.extend_from_slice(content_new.as_bytes());
|
||||
tokio::fs::write(&path, &out_bytes)
|
||||
.await
|
||||
.map_err(|e| ToolError::Other(format!("{}: {e}", path.display())))?;
|
||||
|
||||
let (added, removed) = diff_stats(&content_old, &content_new);
|
||||
let mut output = ToolOutput::new(pattern, "Edit applied successfully.".to_string());
|
||||
output.metadata = serde_json::json!({"diff": diff, "added": added, "removed": removed});
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use harness_core::event::EventBus;
|
||||
use harness_core::permission::{spawn_auto_approve, PermissionService};
|
||||
use harness_core::tool::{MetadataSink, PermissionHandle};
|
||||
use harness_core::types::SessionId;
|
||||
use std::sync::Mutex;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
fn ctx(cwd: PathBuf) -> ToolCtx {
|
||||
let bus = EventBus::new();
|
||||
let service = Arc::new(PermissionService::new(bus.clone()));
|
||||
spawn_auto_approve(bus, service.clone());
|
||||
let (metadata, _rx) = MetadataSink::channel();
|
||||
ToolCtx {
|
||||
session_id: SessionId::new(),
|
||||
message_id: harness_core::types::MessageId::new(),
|
||||
call_id: "call_1".into(),
|
||||
data_dir: cwd.join("tool-output"),
|
||||
cwd,
|
||||
cancel: CancellationToken::new(),
|
||||
ask: PermissionHandle::new(
|
||||
service,
|
||||
SessionId::new(),
|
||||
Vec::new(),
|
||||
Arc::new(Mutex::new(Vec::new())),
|
||||
CancellationToken::new(),
|
||||
),
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creates_new_file_when_old_string_is_empty() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let result = EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "newfile.txt", "old_string": "", "new_string": "new content"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.metadata["diff"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("new content"));
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(dir.path().join("newfile.txt")).unwrap(),
|
||||
"new content"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_empty_old_string_on_existing_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("existing.txt"), "original").unwrap();
|
||||
let err = EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "existing.txt", "old_string": "", "new_string": "x"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
matches!(&err, ToolError::Invalid(msg) if msg.contains("oldString cannot be empty"))
|
||||
);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(dir.path().join("existing.txt")).unwrap(),
|
||||
"original"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creates_new_file_with_nested_directories() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "nested/dir/file.txt", "old_string": "", "new_string": "nested file"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(dir.path().join("nested/dir/file.txt")).unwrap(),
|
||||
"nested file"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replaces_text_in_existing_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("existing.txt"), "old content here").unwrap();
|
||||
let result = EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "existing.txt", "old_string": "old content", "new_string": "new content"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.output.contains("Edit applied successfully"));
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(dir.path().join("existing.txt")).unwrap(),
|
||||
"new content here"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn preserves_bom_and_only_edits_visible_content() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut original = UTF8_BOM.to_vec();
|
||||
original.extend_from_slice(b"using System;\nclass Test {}\n");
|
||||
std::fs::write(dir.path().join("existing.cs"), &original).unwrap();
|
||||
|
||||
let result = EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "existing.cs", "old_string": "using System;", "new_string": "using Up;"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.metadata["diff"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("-using System;"));
|
||||
assert!(result.metadata["diff"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.contains("+using Up;"));
|
||||
|
||||
let content = std::fs::read(dir.path().join("existing.cs")).unwrap();
|
||||
assert_eq!(&content[..3], &UTF8_BOM);
|
||||
assert_eq!(&content[3..], b"using Up;\nclass Test {}\n");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn errors_when_file_does_not_exist() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let err = EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "nonexistent.txt", "old_string": "old", "new_string": "new"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(&err, ToolError::Invalid(msg) if msg.contains("not found")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn errors_when_old_equals_new() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("file.txt"), "content").unwrap();
|
||||
let err = EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "file.txt", "old_string": "same", "new_string": "same"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(&err, ToolError::Invalid(msg) if msg.contains("identical")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn errors_when_path_is_a_directory() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::create_dir(dir.path().join("adir")).unwrap();
|
||||
let err = EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "adir", "old_string": "old", "new_string": "new"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(&err, ToolError::Invalid(msg) if msg.contains("directory")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replaces_all_occurrences_with_replace_all_option() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("file.txt"), "foo bar foo baz foo").unwrap();
|
||||
EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "file.txt", "old_string": "foo", "new_string": "qux", "replace_all": true}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(dir.path().join("file.txt")).unwrap(),
|
||||
"qux bar qux baz qux"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handles_crlf_line_endings() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("file.txt"), "line1\r\nold\r\nline3").unwrap();
|
||||
EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "file.txt", "old_string": "old", "new_string": "new"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(dir.path().join("file.txt")).unwrap(),
|
||||
"line1\r\nnew\r\nline3"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tracks_file_diff_statistics() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("file.txt"), "line1\nline2\nline3").unwrap();
|
||||
let result = EditTool::new()
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "file.txt", "old_string": "line2", "new_string": "new line a\nnew line b"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(result.metadata["added"].as_u64().unwrap() > 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_edits_to_the_same_file_serialize_without_losing_either_change() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(
|
||||
dir.path().join("file.txt"),
|
||||
"top = 0\nmiddle = keep\nbottom = 0\n",
|
||||
)
|
||||
.unwrap();
|
||||
let tool = Arc::new(EditTool::new());
|
||||
|
||||
let t1 = tool.clone();
|
||||
let cwd1 = dir.path().to_path_buf();
|
||||
let h1 = tokio::spawn(async move {
|
||||
t1.execute(
|
||||
serde_json::json!({"file_path": "file.txt", "old_string": "top = 0", "new_string": "top = 1"}),
|
||||
ctx(cwd1),
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
let t2 = tool.clone();
|
||||
let cwd2 = dir.path().to_path_buf();
|
||||
let h2 = tokio::spawn(async move {
|
||||
t2.execute(
|
||||
serde_json::json!({"file_path": "file.txt", "old_string": "bottom = 0", "new_string": "bottom = 2"}),
|
||||
ctx(cwd2),
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
h1.await.unwrap().unwrap();
|
||||
h2.await.unwrap().unwrap();
|
||||
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(dir.path().join("file.txt")).unwrap(),
|
||||
"top = 1\nmiddle = keep\nbottom = 2\n"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,696 @@
|
||||
//! Verbatim port of opencode's replacer chain (`packages/opencode/src/tool/edit.ts`).
|
||||
//! Each replacer returns candidate matches in the same order the TS generator yields them;
|
||||
//! `replace()` walks replacers in a fixed sequence and takes the first candidate that
|
||||
//! resolves unambiguously (or throws — the disproportionate-match guard is a hard stop,
|
||||
//! not a "try the next candidate").
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
|
||||
pub enum EditError {
|
||||
#[error("No changes to apply: oldString and newString are identical.")]
|
||||
Identical,
|
||||
#[error("oldString cannot be empty when editing an existing file. Provide the exact text to replace, or use write for an intentional full-file replacement.")]
|
||||
EmptyOldString,
|
||||
#[error("Could not find oldString in the file. It must match exactly, including whitespace, indentation, and line endings.")]
|
||||
NotFound,
|
||||
#[error("Found multiple matches for oldString. Provide more surrounding context to make the match unique.")]
|
||||
MultipleMatches,
|
||||
#[error("Refusing replacement because the matched span is much larger than oldString. Re-read the file and provide the full exact oldString for the intended replacement.")]
|
||||
DisproportionateMatch,
|
||||
}
|
||||
|
||||
const SINGLE_CANDIDATE_SIMILARITY_THRESHOLD: f64 = 0.65;
|
||||
const MULTIPLE_CANDIDATES_SIMILARITY_THRESHOLD: f64 = 0.65;
|
||||
|
||||
fn levenshtein(a: &str, b: &str) -> usize {
|
||||
let a: Vec<char> = a.chars().collect();
|
||||
let b: Vec<char> = b.chars().collect();
|
||||
if a.is_empty() || b.is_empty() {
|
||||
return a.len().max(b.len());
|
||||
}
|
||||
let mut matrix = vec![vec![0usize; b.len() + 1]; a.len() + 1];
|
||||
for (i, row) in matrix.iter_mut().enumerate() {
|
||||
row[0] = i;
|
||||
}
|
||||
for (j, cell) in matrix[0].iter_mut().enumerate() {
|
||||
*cell = j;
|
||||
}
|
||||
for i in 1..=a.len() {
|
||||
for j in 1..=b.len() {
|
||||
let cost = if a[i - 1] == b[j - 1] { 0 } else { 1 };
|
||||
matrix[i][j] = (matrix[i - 1][j] + 1)
|
||||
.min(matrix[i][j - 1] + 1)
|
||||
.min(matrix[i - 1][j - 1] + cost);
|
||||
}
|
||||
}
|
||||
matrix[a.len()][b.len()]
|
||||
}
|
||||
|
||||
fn simple(_content: &str, find: &str) -> Vec<String> {
|
||||
vec![find.to_string()]
|
||||
}
|
||||
|
||||
fn line_trimmed(content: &str, find: &str) -> Vec<String> {
|
||||
let original_lines: Vec<&str> = content.split('\n').collect();
|
||||
let mut search_lines: Vec<&str> = find.split('\n').collect();
|
||||
if search_lines.last() == Some(&"") {
|
||||
search_lines.pop();
|
||||
}
|
||||
let mut out = Vec::new();
|
||||
if search_lines.is_empty() || original_lines.len() < search_lines.len() {
|
||||
return out;
|
||||
}
|
||||
|
||||
for i in 0..=(original_lines.len() - search_lines.len()) {
|
||||
let matches =
|
||||
(0..search_lines.len()).all(|j| original_lines[i + j].trim() == search_lines[j].trim());
|
||||
if matches {
|
||||
let mut match_start = 0usize;
|
||||
for line in &original_lines[..i] {
|
||||
match_start += line.len() + 1;
|
||||
}
|
||||
let mut match_end = match_start;
|
||||
for k in 0..search_lines.len() {
|
||||
match_end += original_lines[i + k].len();
|
||||
if k < search_lines.len() - 1 {
|
||||
match_end += 1;
|
||||
}
|
||||
}
|
||||
out.push(content[match_start..match_end].to_string());
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn extract_lines(original_lines: &[&str], content: &str, start: usize, end: usize) -> String {
|
||||
let mut match_start = 0usize;
|
||||
for line in &original_lines[..start] {
|
||||
match_start += line.len() + 1;
|
||||
}
|
||||
let mut match_end = match_start;
|
||||
for (k, line) in original_lines.iter().enumerate().take(end + 1).skip(start) {
|
||||
match_end += line.len();
|
||||
if k < end {
|
||||
match_end += 1;
|
||||
}
|
||||
}
|
||||
content[match_start..match_end].to_string()
|
||||
}
|
||||
|
||||
fn block_anchor(content: &str, find: &str) -> Vec<String> {
|
||||
let original_lines: Vec<&str> = content.split('\n').collect();
|
||||
let mut search_lines: Vec<&str> = find.split('\n').collect();
|
||||
if search_lines.len() < 3 {
|
||||
return vec![];
|
||||
}
|
||||
if search_lines.last() == Some(&"") {
|
||||
search_lines.pop();
|
||||
}
|
||||
|
||||
let first_line_search = search_lines[0].trim();
|
||||
let last_line_search = search_lines[search_lines.len() - 1].trim();
|
||||
let search_block_size = search_lines.len();
|
||||
let max_line_delta = ((search_block_size as f64 * 0.25).floor() as i64).max(1);
|
||||
|
||||
let mut candidates: Vec<(usize, usize)> = Vec::new();
|
||||
for i in 0..original_lines.len() {
|
||||
if original_lines[i].trim() != first_line_search {
|
||||
continue;
|
||||
}
|
||||
let mut j = i + 2;
|
||||
while j < original_lines.len() {
|
||||
if original_lines[j].trim() == last_line_search {
|
||||
let actual_block_size = (j - i + 1) as i64;
|
||||
if (actual_block_size - search_block_size as i64).abs() <= max_line_delta {
|
||||
candidates.push((i, j));
|
||||
}
|
||||
break;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
}
|
||||
if candidates.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
|
||||
let middle_similarity = |start: usize, end: usize, early_exit: bool| -> f64 {
|
||||
let actual_block_size = end - start + 1;
|
||||
let lines_to_check = (search_block_size as i64 - 2).min(actual_block_size as i64 - 2);
|
||||
if lines_to_check <= 0 {
|
||||
return 1.0;
|
||||
}
|
||||
let lines_to_check = lines_to_check as usize;
|
||||
let upper = (search_block_size - 1).min(actual_block_size - 1);
|
||||
let mut similarity = 0.0f64;
|
||||
for j in 1..upper {
|
||||
let original_line = original_lines[start + j].trim();
|
||||
let search_line = search_lines[j].trim();
|
||||
let max_len = original_line.len().max(search_line.len());
|
||||
if max_len == 0 {
|
||||
continue;
|
||||
}
|
||||
let distance = levenshtein(original_line, search_line);
|
||||
similarity += (1.0 - distance as f64 / max_len as f64) / lines_to_check as f64;
|
||||
if early_exit && similarity >= SINGLE_CANDIDATE_SIMILARITY_THRESHOLD {
|
||||
break;
|
||||
}
|
||||
}
|
||||
similarity
|
||||
};
|
||||
|
||||
if candidates.len() == 1 {
|
||||
let (start, end) = candidates[0];
|
||||
let similarity = middle_similarity(start, end, true);
|
||||
if similarity >= SINGLE_CANDIDATE_SIMILARITY_THRESHOLD {
|
||||
return vec![extract_lines(&original_lines, content, start, end)];
|
||||
}
|
||||
return vec![];
|
||||
}
|
||||
|
||||
let mut best: Option<(usize, usize)> = None;
|
||||
let mut max_similarity = -1.0f64;
|
||||
for &(start, end) in &candidates {
|
||||
let actual_block_size = end - start + 1;
|
||||
let lines_to_check = (search_block_size as i64 - 2).min(actual_block_size as i64 - 2);
|
||||
let similarity = if lines_to_check <= 0 {
|
||||
1.0
|
||||
} else {
|
||||
let upper = (search_block_size - 1).min(actual_block_size - 1);
|
||||
let mut sum = 0.0;
|
||||
for j in 1..upper {
|
||||
let original_line = original_lines[start + j].trim();
|
||||
let search_line = search_lines[j].trim();
|
||||
let max_len = original_line.len().max(search_line.len());
|
||||
if max_len == 0 {
|
||||
continue;
|
||||
}
|
||||
sum += 1.0 - levenshtein(original_line, search_line) as f64 / max_len as f64;
|
||||
}
|
||||
sum / lines_to_check as f64
|
||||
};
|
||||
if similarity > max_similarity {
|
||||
max_similarity = similarity;
|
||||
best = Some((start, end));
|
||||
}
|
||||
}
|
||||
if max_similarity >= MULTIPLE_CANDIDATES_SIMILARITY_THRESHOLD {
|
||||
if let Some((start, end)) = best {
|
||||
return vec![extract_lines(&original_lines, content, start, end)];
|
||||
}
|
||||
}
|
||||
vec![]
|
||||
}
|
||||
|
||||
fn normalize_whitespace(text: &str) -> String {
|
||||
let mut out = String::new();
|
||||
let mut in_ws = false;
|
||||
for c in text.chars() {
|
||||
if c.is_whitespace() {
|
||||
if !in_ws {
|
||||
out.push(' ');
|
||||
in_ws = true;
|
||||
}
|
||||
} else {
|
||||
out.push(c);
|
||||
in_ws = false;
|
||||
}
|
||||
}
|
||||
out.trim().to_string()
|
||||
}
|
||||
|
||||
/// Finds `words` in `line` separated by runs of whitespace, returning the exact
|
||||
/// (non-normalized) substring — a hand-rolled stand-in for the TS version's regex built
|
||||
/// from escaped words joined by `\s+`.
|
||||
fn find_flexible_whitespace_match(line: &str, words: &[&str]) -> Option<String> {
|
||||
if words.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let chars: Vec<char> = line.chars().collect();
|
||||
let n = chars.len();
|
||||
'outer: for start in 0..n {
|
||||
let mut pos = start;
|
||||
for (wi, word) in words.iter().enumerate() {
|
||||
let wchars: Vec<char> = word.chars().collect();
|
||||
if pos + wchars.len() > n || chars[pos..pos + wchars.len()] != wchars[..] {
|
||||
continue 'outer;
|
||||
}
|
||||
pos += wchars.len();
|
||||
if wi < words.len() - 1 {
|
||||
let ws_start = pos;
|
||||
while pos < n && chars[pos].is_whitespace() {
|
||||
pos += 1;
|
||||
}
|
||||
if pos == ws_start {
|
||||
continue 'outer;
|
||||
}
|
||||
}
|
||||
}
|
||||
return Some(chars[start..pos].iter().collect());
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn whitespace_normalized(content: &str, find: &str) -> Vec<String> {
|
||||
let mut out = Vec::new();
|
||||
let normalized_find = normalize_whitespace(find);
|
||||
|
||||
let lines: Vec<&str> = content.split('\n').collect();
|
||||
for line in &lines {
|
||||
if normalize_whitespace(line) == normalized_find {
|
||||
out.push(line.to_string());
|
||||
} else if normalize_whitespace(line).contains(&normalized_find) {
|
||||
let words: Vec<&str> = find.split_whitespace().collect();
|
||||
if let Some(m) = find_flexible_whitespace_match(line, &words) {
|
||||
out.push(m);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let find_lines: Vec<&str> = find.split('\n').collect();
|
||||
if find_lines.len() > 1 && lines.len() >= find_lines.len() {
|
||||
for i in 0..=(lines.len() - find_lines.len()) {
|
||||
let block = lines[i..i + find_lines.len()].join("\n");
|
||||
if normalize_whitespace(&block) == normalized_find {
|
||||
out.push(block);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn remove_indentation(text: &str) -> String {
|
||||
let lines: Vec<&str> = text.split('\n').collect();
|
||||
let non_empty: Vec<&&str> = lines.iter().filter(|l| !l.trim().is_empty()).collect();
|
||||
if non_empty.is_empty() {
|
||||
return text.to_string();
|
||||
}
|
||||
let min_indent = non_empty
|
||||
.iter()
|
||||
.map(|l| l.len() - l.trim_start().len())
|
||||
.min()
|
||||
.unwrap_or(0);
|
||||
lines
|
||||
.iter()
|
||||
.map(|l| {
|
||||
if l.trim().is_empty() {
|
||||
l.to_string()
|
||||
} else {
|
||||
l[min_indent.min(l.len())..].to_string()
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
|
||||
fn indentation_flexible(content: &str, find: &str) -> Vec<String> {
|
||||
let normalized_find = remove_indentation(find);
|
||||
let content_lines: Vec<&str> = content.split('\n').collect();
|
||||
let find_lines: Vec<&str> = find.split('\n').collect();
|
||||
let mut out = Vec::new();
|
||||
if content_lines.len() < find_lines.len() {
|
||||
return out;
|
||||
}
|
||||
for i in 0..=(content_lines.len() - find_lines.len()) {
|
||||
let block = content_lines[i..i + find_lines.len()].join("\n");
|
||||
if remove_indentation(&block) == normalized_find {
|
||||
out.push(block);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn unescape_string(s: &str) -> String {
|
||||
let mut out = String::new();
|
||||
let mut chars = s.chars().peekable();
|
||||
while let Some(c) = chars.next() {
|
||||
if c == '\\' {
|
||||
match chars.peek().copied() {
|
||||
Some('n') => {
|
||||
out.push('\n');
|
||||
chars.next();
|
||||
}
|
||||
Some('t') => {
|
||||
out.push('\t');
|
||||
chars.next();
|
||||
}
|
||||
Some('r') => {
|
||||
out.push('\r');
|
||||
chars.next();
|
||||
}
|
||||
Some('\'') => {
|
||||
out.push('\'');
|
||||
chars.next();
|
||||
}
|
||||
Some('"') => {
|
||||
out.push('"');
|
||||
chars.next();
|
||||
}
|
||||
Some('`') => {
|
||||
out.push('`');
|
||||
chars.next();
|
||||
}
|
||||
Some('\\') => {
|
||||
out.push('\\');
|
||||
chars.next();
|
||||
}
|
||||
Some('\n') => {
|
||||
out.push('\n');
|
||||
chars.next();
|
||||
}
|
||||
Some('$') => {
|
||||
out.push('$');
|
||||
chars.next();
|
||||
}
|
||||
_ => out.push(c),
|
||||
}
|
||||
} else {
|
||||
out.push(c);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn escape_normalized(content: &str, find: &str) -> Vec<String> {
|
||||
let mut out = Vec::new();
|
||||
let unescaped_find = unescape_string(find);
|
||||
if content.contains(&unescaped_find) {
|
||||
out.push(unescaped_find.clone());
|
||||
}
|
||||
let lines: Vec<&str> = content.split('\n').collect();
|
||||
let find_lines: Vec<&str> = unescaped_find.split('\n').collect();
|
||||
if lines.len() >= find_lines.len() {
|
||||
for i in 0..=(lines.len() - find_lines.len()) {
|
||||
let block = lines[i..i + find_lines.len()].join("\n");
|
||||
if unescape_string(&block) == unescaped_find {
|
||||
out.push(block);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn multi_occurrence(content: &str, find: &str) -> Vec<String> {
|
||||
let mut out = Vec::new();
|
||||
if find.is_empty() {
|
||||
return out;
|
||||
}
|
||||
let mut start = 0usize;
|
||||
while let Some(idx) = content[start..].find(find) {
|
||||
out.push(find.to_string());
|
||||
start += idx + find.len();
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn trimmed_boundary(content: &str, find: &str) -> Vec<String> {
|
||||
let trimmed_find = find.trim();
|
||||
let mut out = Vec::new();
|
||||
if trimmed_find == find {
|
||||
return out;
|
||||
}
|
||||
if content.contains(trimmed_find) {
|
||||
out.push(trimmed_find.to_string());
|
||||
}
|
||||
let lines: Vec<&str> = content.split('\n').collect();
|
||||
let find_lines: Vec<&str> = find.split('\n').collect();
|
||||
if lines.len() >= find_lines.len() {
|
||||
for i in 0..=(lines.len() - find_lines.len()) {
|
||||
let block = lines[i..i + find_lines.len()].join("\n");
|
||||
if block.trim() == trimmed_find {
|
||||
out.push(block);
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn context_aware(content: &str, find: &str) -> Vec<String> {
|
||||
let mut find_lines: Vec<&str> = find.split('\n').collect();
|
||||
if find_lines.len() < 3 {
|
||||
return vec![];
|
||||
}
|
||||
if find_lines.last() == Some(&"") {
|
||||
find_lines.pop();
|
||||
}
|
||||
let content_lines: Vec<&str> = content.split('\n').collect();
|
||||
let first_line = find_lines[0].trim();
|
||||
let last_line = find_lines[find_lines.len() - 1].trim();
|
||||
let mut out = Vec::new();
|
||||
|
||||
for i in 0..content_lines.len() {
|
||||
if content_lines[i].trim() != first_line {
|
||||
continue;
|
||||
}
|
||||
let mut j = i + 2;
|
||||
while j < content_lines.len() {
|
||||
if content_lines[j].trim() == last_line {
|
||||
let block_lines = &content_lines[i..=j];
|
||||
if block_lines.len() == find_lines.len() {
|
||||
let mut matching = 0usize;
|
||||
let mut total_non_empty = 0usize;
|
||||
for k in 1..block_lines.len() - 1 {
|
||||
let block_line = block_lines[k].trim();
|
||||
let find_line = find_lines[k].trim();
|
||||
if !block_line.is_empty() || !find_line.is_empty() {
|
||||
total_non_empty += 1;
|
||||
if block_line == find_line {
|
||||
matching += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
if total_non_empty == 0 || (matching as f64 / total_non_empty as f64) >= 0.5 {
|
||||
out.push(block_lines.join("\n"));
|
||||
}
|
||||
}
|
||||
break;
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
type ReplacerFn = fn(&str, &str) -> Vec<String>;
|
||||
|
||||
const REPLACER_CHAIN: &[ReplacerFn] = &[
|
||||
simple,
|
||||
line_trimmed,
|
||||
block_anchor,
|
||||
whitespace_normalized,
|
||||
indentation_flexible,
|
||||
escape_normalized,
|
||||
trimmed_boundary,
|
||||
context_aware,
|
||||
multi_occurrence,
|
||||
];
|
||||
|
||||
fn is_disproportionate_match(search: &str, old_string: &str) -> bool {
|
||||
let old_lines = old_string.split('\n').count() as i64;
|
||||
let search_lines = search.split('\n').count() as i64;
|
||||
if search_lines >= (old_lines + 3).max(old_lines * 2) {
|
||||
return true;
|
||||
}
|
||||
if old_lines == 1 {
|
||||
return false;
|
||||
}
|
||||
let old_trimmed_len = old_string.trim().len() as i64;
|
||||
search.trim().len() as i64 > (old_trimmed_len + 500).max(old_trimmed_len * 4)
|
||||
}
|
||||
|
||||
/// Applies the replacer chain in order, returning the new content on the first
|
||||
/// unambiguous match. Mirrors opencode's `replace()` exactly, including the
|
||||
/// disproportionate-match guard short-circuiting immediately (not skipping ahead).
|
||||
pub fn replace(
|
||||
content: &str,
|
||||
old_string: &str,
|
||||
new_string: &str,
|
||||
replace_all: bool,
|
||||
) -> Result<String, EditError> {
|
||||
if old_string == new_string {
|
||||
return Err(EditError::Identical);
|
||||
}
|
||||
if old_string.is_empty() {
|
||||
return Err(EditError::EmptyOldString);
|
||||
}
|
||||
|
||||
let mut not_found = true;
|
||||
for replacer in REPLACER_CHAIN {
|
||||
for search in replacer(content, old_string) {
|
||||
let index = match content.find(&search) {
|
||||
Some(i) => i,
|
||||
None => continue,
|
||||
};
|
||||
not_found = false;
|
||||
if is_disproportionate_match(&search, old_string) {
|
||||
return Err(EditError::DisproportionateMatch);
|
||||
}
|
||||
if replace_all {
|
||||
return Ok(content.replace(&search, new_string));
|
||||
}
|
||||
let last_index = content.rfind(&search).unwrap();
|
||||
if index != last_index {
|
||||
continue;
|
||||
}
|
||||
return Ok(format!(
|
||||
"{}{}{}",
|
||||
&content[..index],
|
||||
new_string,
|
||||
&content[index + search.len()..]
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if not_found {
|
||||
Err(EditError::NotFound)
|
||||
} else {
|
||||
Err(EditError::MultipleMatches)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn replaces_exact_text_in_existing_content() {
|
||||
let result = replace("old content here", "old content", "new content", false).unwrap();
|
||||
assert_eq!(result, "new content here");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_identical_old_and_new() {
|
||||
assert_eq!(
|
||||
replace("content", "same", "same", false),
|
||||
Err(EditError::Identical)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_old_string() {
|
||||
assert_eq!(
|
||||
replace("content", "", "x", false),
|
||||
Err(EditError::EmptyOldString)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_not_found() {
|
||||
let err = replace("actual content", "not in file", "replacement", false).unwrap_err();
|
||||
assert_eq!(err, EditError::NotFound);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_loose_block_anchor_matches() {
|
||||
let original = [
|
||||
"function configure() {",
|
||||
" keepImportantState()",
|
||||
" removeAllUserData()",
|
||||
" archiveBackups()",
|
||||
" auditLog()",
|
||||
"}",
|
||||
]
|
||||
.join("\n");
|
||||
let old = ["function configure() {", " const enabled = true", "}"].join("\n");
|
||||
let new = ["function configure() {", " const enabled = false", "}"].join("\n");
|
||||
assert_eq!(
|
||||
replace(&original, &old, &new, false),
|
||||
Err(EditError::NotFound)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_block_anchor_matches_with_unrelated_middle_content() {
|
||||
let original = ["function configure() {", " removeAllUserData()", "}"].join("\n");
|
||||
let old = ["function configure() {", " const enabled = true", "}"].join("\n");
|
||||
let new = ["function configure() {", " const enabled = false", "}"].join("\n");
|
||||
assert_eq!(
|
||||
replace(&original, &old, &new, false),
|
||||
Err(EditError::NotFound)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replaces_all_occurrences_with_replace_all() {
|
||||
let result = replace("foo bar foo baz foo", "foo", "qux", true).unwrap();
|
||||
assert_eq!(result, "qux bar qux baz qux");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_ambiguous_single_occurrence_request() {
|
||||
let err = replace("foo bar foo", "foo", "qux", false).unwrap_err();
|
||||
assert_eq!(err, EditError::MultipleMatches);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn handles_multiline_replacements() {
|
||||
let result = replace(
|
||||
"line1\nline2\nline3",
|
||||
"line2",
|
||||
"new line 2\nextra line",
|
||||
false,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(result, "line1\nnew line 2\nextra line\nline3");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn line_trimmed_replacer_matches_despite_surrounding_whitespace_changes() {
|
||||
let content = " line one \n line two \n";
|
||||
let matches = line_trimmed(content, "line one\nline two");
|
||||
assert_eq!(matches, vec![" line one \n line two ".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn whitespace_normalized_matches_reformatted_line() {
|
||||
let content = "const x = 1;\n";
|
||||
let matches = whitespace_normalized(content, "const x = 1;");
|
||||
assert!(matches.iter().any(|m| m == "const x = 1;"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn indentation_flexible_matches_reindented_block() {
|
||||
let content = " if true {\n foo()\n }\n";
|
||||
let find = "if true {\n foo()\n}";
|
||||
let matches = indentation_flexible(content, find);
|
||||
assert_eq!(matches.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_normalized_matches_escaped_newline_in_find() {
|
||||
let content = "line one\nline two\n";
|
||||
let matches = escape_normalized(content, "line one\\nline two");
|
||||
assert!(matches.contains(&"line one\nline two".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trimmed_boundary_matches_padded_find() {
|
||||
let content = "exact text here\n";
|
||||
let matches = trimmed_boundary(content, " exact text here ");
|
||||
assert!(matches.contains(&"exact text here".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disproportionate_match_flags_far_larger_line_count() {
|
||||
let old = "one line";
|
||||
let search = "a\nb\nc\nd"; // 4 lines >= max(1+3, 1*2) = 4
|
||||
assert!(is_disproportionate_match(search, old));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disproportionate_match_flags_far_larger_char_span_same_line_count() {
|
||||
let old = "short";
|
||||
let search = "x".repeat(600); // > max(5+500, 5*4) = 505, still 1 line like `old`
|
||||
assert!(!is_disproportionate_match(&search, old)); // old_lines == 1 short-circuits to false
|
||||
|
||||
let old_multiline = "short\nline";
|
||||
let search_multiline = format!("{}\nline", "x".repeat(600));
|
||||
assert!(is_disproportionate_match(&search_multiline, old_multiline));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disproportionate_match_allows_reasonably_sized_matches() {
|
||||
assert!(!is_disproportionate_match(
|
||||
"old content",
|
||||
"old content here"
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
use std::path::PathBuf;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use globset::GlobBuilder;
|
||||
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
|
||||
use ignore::WalkBuilder;
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::paths;
|
||||
|
||||
const MAX_RESULTS: usize = 100;
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
struct GlobParams {
|
||||
pattern: String,
|
||||
path: Option<String>,
|
||||
}
|
||||
|
||||
/// Gitignore-aware file finder — read-only, no permission ask.
|
||||
pub struct GlobTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GlobTool {
|
||||
fn name(&self) -> &str {
|
||||
"glob"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Finds files by glob pattern (gitignore-aware), most recently modified first."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::to_value(schemars::schema_for!(GlobParams)).unwrap()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: GlobParams =
|
||||
serde_json::from_value(input).map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
let root = params
|
||||
.path
|
||||
.as_ref()
|
||||
.map(|p| paths::resolve(&ctx.cwd, p))
|
||||
.unwrap_or_else(|| ctx.cwd.clone());
|
||||
let pattern = params.pattern.clone();
|
||||
|
||||
let matcher = GlobBuilder::new(&pattern)
|
||||
.literal_separator(true)
|
||||
.build()
|
||||
.map_err(|e| ToolError::Invalid(e.to_string()))?
|
||||
.compile_matcher();
|
||||
|
||||
let mut matches: Vec<(PathBuf, SystemTime)> = Vec::new();
|
||||
for entry in WalkBuilder::new(&root)
|
||||
.hidden(false)
|
||||
.require_git(false)
|
||||
.build()
|
||||
{
|
||||
let entry = match entry {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if !entry.file_type().map(|t| t.is_file()).unwrap_or(false) {
|
||||
continue;
|
||||
}
|
||||
let rel = entry.path().strip_prefix(&root).unwrap_or(entry.path());
|
||||
if !matcher.is_match(rel) {
|
||||
continue;
|
||||
}
|
||||
let mtime = entry
|
||||
.metadata()
|
||||
.ok()
|
||||
.and_then(|m| m.modified().ok())
|
||||
.unwrap_or(SystemTime::UNIX_EPOCH);
|
||||
matches.push((entry.path().to_path_buf(), mtime));
|
||||
}
|
||||
|
||||
matches.sort_by_key(|(_, mtime)| std::cmp::Reverse(*mtime));
|
||||
let truncated = matches.len() > MAX_RESULTS;
|
||||
matches.truncate(MAX_RESULTS);
|
||||
|
||||
let mut output = matches
|
||||
.iter()
|
||||
.map(|(p, _)| p.display().to_string())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
if truncated {
|
||||
output.push_str(
|
||||
"\n\n(Results are truncated: showing first 100 results. Consider using a more specific path or pattern.)",
|
||||
);
|
||||
}
|
||||
if output.is_empty() {
|
||||
output = "(no matches)".to_string();
|
||||
}
|
||||
Ok(ToolOutput::new(pattern, output))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use harness_core::event::EventBus;
|
||||
use harness_core::permission::PermissionService;
|
||||
use harness_core::tool::{MetadataSink, PermissionHandle};
|
||||
use harness_core::types::SessionId;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
fn ctx(cwd: std::path::PathBuf) -> ToolCtx {
|
||||
let bus = EventBus::new();
|
||||
let service = Arc::new(PermissionService::new(bus));
|
||||
let (metadata, _rx) = MetadataSink::channel();
|
||||
ToolCtx {
|
||||
session_id: SessionId::new(),
|
||||
message_id: harness_core::types::MessageId::new(),
|
||||
call_id: "call_1".into(),
|
||||
data_dir: cwd.join("tool-output"),
|
||||
cwd,
|
||||
cancel: CancellationToken::new(),
|
||||
ask: PermissionHandle::new(
|
||||
service,
|
||||
SessionId::new(),
|
||||
Vec::new(),
|
||||
Arc::new(Mutex::new(Vec::new())),
|
||||
CancellationToken::new(),
|
||||
),
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finds_files_matching_pattern() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.rs"), "").unwrap();
|
||||
std::fs::write(dir.path().join("b.txt"), "").unwrap();
|
||||
let output = GlobTool
|
||||
.execute(
|
||||
serde_json::json!({"pattern": "*.rs"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.output.ends_with("a.rs"));
|
||||
assert!(!output.output.contains("b.txt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn respects_gitignore() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join(".gitignore"), "ignored.rs\n").unwrap();
|
||||
std::fs::write(dir.path().join("ignored.rs"), "").unwrap();
|
||||
std::fs::write(dir.path().join("kept.rs"), "").unwrap();
|
||||
let output = GlobTool
|
||||
.execute(
|
||||
serde_json::json!({"pattern": "*.rs"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.output.contains("kept.rs"));
|
||||
assert!(!output.output.contains("ignored.rs"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_matches_reports_clearly() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let output = GlobTool
|
||||
.execute(
|
||||
serde_json::json!({"pattern": "*.nonexistent"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(output.output, "(no matches)");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
use std::path::PathBuf;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use globset::GlobBuilder;
|
||||
use grep_regex::RegexMatcher;
|
||||
use grep_searcher::sinks::UTF8;
|
||||
use grep_searcher::SearcherBuilder;
|
||||
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
|
||||
use ignore::WalkBuilder;
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::paths;
|
||||
|
||||
const MAX_RESULTS: usize = 100;
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
struct GrepParams {
|
||||
pattern: String,
|
||||
path: Option<String>,
|
||||
include: Option<String>,
|
||||
}
|
||||
|
||||
struct Hit {
|
||||
file: PathBuf,
|
||||
line_number: u64,
|
||||
line: String,
|
||||
mtime: SystemTime,
|
||||
}
|
||||
|
||||
/// Content search via ripgrep's engine as a library — read-only, no permission ask.
|
||||
pub struct GrepTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GrepTool {
|
||||
fn name(&self) -> &str {
|
||||
"grep"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Searches file contents by regex (gitignore-aware), most recently modified first."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::to_value(schemars::schema_for!(GrepParams)).unwrap()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: GrepParams =
|
||||
serde_json::from_value(input).map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
let root = params
|
||||
.path
|
||||
.as_ref()
|
||||
.map(|p| paths::resolve(&ctx.cwd, p))
|
||||
.unwrap_or_else(|| ctx.cwd.clone());
|
||||
|
||||
let matcher =
|
||||
RegexMatcher::new(¶ms.pattern).map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
let include_matcher = params
|
||||
.include
|
||||
.as_ref()
|
||||
.map(|inc| GlobBuilder::new(inc).literal_separator(true).build())
|
||||
.transpose()
|
||||
.map_err(|e| ToolError::Invalid(e.to_string()))?
|
||||
.map(|g| g.compile_matcher());
|
||||
|
||||
let mut hits: Vec<Hit> = Vec::new();
|
||||
for entry in WalkBuilder::new(&root)
|
||||
.hidden(false)
|
||||
.require_git(false)
|
||||
.build()
|
||||
{
|
||||
let entry = match entry {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if !entry.file_type().map(|t| t.is_file()).unwrap_or(false) {
|
||||
continue;
|
||||
}
|
||||
let rel = entry.path().strip_prefix(&root).unwrap_or(entry.path());
|
||||
if let Some(glob) = &include_matcher {
|
||||
if !glob.is_match(rel) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
let mtime = entry
|
||||
.metadata()
|
||||
.ok()
|
||||
.and_then(|m| m.modified().ok())
|
||||
.unwrap_or(SystemTime::UNIX_EPOCH);
|
||||
|
||||
let path = entry.path().to_path_buf();
|
||||
let mut searcher = SearcherBuilder::new().line_number(true).build();
|
||||
let _ = searcher.search_path(
|
||||
&matcher,
|
||||
&path,
|
||||
UTF8(|line_number, line| {
|
||||
hits.push(Hit {
|
||||
file: path.clone(),
|
||||
line_number,
|
||||
line: line.trim_end().to_string(),
|
||||
mtime,
|
||||
});
|
||||
Ok(true)
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
hits.sort_by_key(|h| std::cmp::Reverse(h.mtime));
|
||||
let truncated = hits.len() > MAX_RESULTS;
|
||||
hits.truncate(MAX_RESULTS);
|
||||
|
||||
let mut output = hits
|
||||
.iter()
|
||||
.map(|h| format!("{}:{}:{}", h.file.display(), h.line_number, h.line))
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
if truncated {
|
||||
output.push_str("\n\n(Results are truncated: showing first 100 matches. Consider a more specific pattern.)");
|
||||
}
|
||||
if output.is_empty() {
|
||||
output = "(no matches)".to_string();
|
||||
}
|
||||
Ok(ToolOutput::new(params.pattern, output))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use harness_core::event::EventBus;
|
||||
use harness_core::permission::PermissionService;
|
||||
use harness_core::tool::{MetadataSink, PermissionHandle};
|
||||
use harness_core::types::SessionId;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
fn ctx(cwd: std::path::PathBuf) -> ToolCtx {
|
||||
let bus = EventBus::new();
|
||||
let service = Arc::new(PermissionService::new(bus));
|
||||
let (metadata, _rx) = MetadataSink::channel();
|
||||
ToolCtx {
|
||||
session_id: SessionId::new(),
|
||||
message_id: harness_core::types::MessageId::new(),
|
||||
call_id: "call_1".into(),
|
||||
data_dir: cwd.join("tool-output"),
|
||||
cwd,
|
||||
cancel: CancellationToken::new(),
|
||||
ask: PermissionHandle::new(
|
||||
service,
|
||||
SessionId::new(),
|
||||
Vec::new(),
|
||||
Arc::new(Mutex::new(Vec::new())),
|
||||
CancellationToken::new(),
|
||||
),
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finds_matching_lines() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.rs"), "fn main() {}\nfn helper() {}\n").unwrap();
|
||||
let output = GrepTool
|
||||
.execute(
|
||||
serde_json::json!({"pattern": "fn main"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.output.contains("a.rs:1:fn main() {}"));
|
||||
assert!(!output.output.contains("helper"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn filters_by_include_glob() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.rs"), "needle\n").unwrap();
|
||||
std::fs::write(dir.path().join("b.txt"), "needle\n").unwrap();
|
||||
let output = GrepTool
|
||||
.execute(
|
||||
serde_json::json!({"pattern": "needle", "include": "*.rs"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.output.contains("a.rs"));
|
||||
assert!(!output.output.contains("b.txt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn no_matches_reports_clearly() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.rs"), "content\n").unwrap();
|
||||
let output = GrepTool
|
||||
.execute(
|
||||
serde_json::json!({"pattern": "nonexistent"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(output.output, "(no matches)");
|
||||
}
|
||||
}
|
||||
@@ -1 +1,28 @@
|
||||
// Built-in tools (bash, read, edit, write, glob, grep, todo, webfetch, task) land here in M1/M4.
|
||||
mod bash;
|
||||
mod edit;
|
||||
mod glob;
|
||||
mod grep;
|
||||
mod paths;
|
||||
mod read;
|
||||
mod write;
|
||||
|
||||
pub use bash::BashTool;
|
||||
pub use edit::EditTool;
|
||||
pub use glob::GlobTool;
|
||||
pub use grep::GrepTool;
|
||||
pub use read::ReadTool;
|
||||
pub use write::WriteTool;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use harness_core::tool::ToolRegistry;
|
||||
|
||||
/// Registers all M1 built-ins: read, write, edit, bash, glob, grep.
|
||||
pub fn register_builtins(registry: &mut ToolRegistry) {
|
||||
registry.register(Arc::new(ReadTool));
|
||||
registry.register(Arc::new(WriteTool));
|
||||
registry.register(Arc::new(EditTool::new()));
|
||||
registry.register(Arc::new(BashTool));
|
||||
registry.register(Arc::new(GlobTool));
|
||||
registry.register(Arc::new(GrepTool));
|
||||
}
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Resolves a (possibly relative) tool-supplied path against `cwd`.
|
||||
pub fn resolve(cwd: &Path, file_path: &str) -> PathBuf {
|
||||
let path = Path::new(file_path);
|
||||
if path.is_absolute() {
|
||||
path.to_path_buf()
|
||||
} else {
|
||||
cwd.join(path)
|
||||
}
|
||||
}
|
||||
|
||||
/// Slash-separated path relative to `cwd`, used as the permission pattern — falls back to
|
||||
/// the absolute path (still slash-separated) when `path` isn't under `cwd`.
|
||||
pub fn relative_pattern(cwd: &Path, path: &Path) -> String {
|
||||
let rel = path.strip_prefix(cwd).unwrap_or(path);
|
||||
rel.to_string_lossy().replace('\\', "/")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn resolve_joins_relative_paths_to_cwd() {
|
||||
let cwd = Path::new("/home/user/project");
|
||||
assert_eq!(
|
||||
resolve(cwd, "src/main.rs"),
|
||||
PathBuf::from("/home/user/project/src/main.rs")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_leaves_absolute_paths_untouched() {
|
||||
let cwd = Path::new("/home/user/project");
|
||||
assert_eq!(resolve(cwd, "/etc/hosts"), PathBuf::from("/etc/hosts"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn relative_pattern_strips_cwd_prefix() {
|
||||
let cwd = Path::new("/home/user/project");
|
||||
let path = Path::new("/home/user/project/src/main.rs");
|
||||
assert_eq!(relative_pattern(cwd, path), "src/main.rs");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn relative_pattern_falls_back_to_absolute_outside_cwd() {
|
||||
let cwd = Path::new("/home/user/project");
|
||||
let path = Path::new("/etc/hosts");
|
||||
assert_eq!(relative_pattern(cwd, path), "/etc/hosts");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
use async_trait::async_trait;
|
||||
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::paths;
|
||||
|
||||
const DEFAULT_LIMIT: usize = 2000;
|
||||
const LINE_CHAR_CLAMP: usize = 2000;
|
||||
const IMAGE_EXTENSIONS: &[&str] = &["png", "jpg", "jpeg", "gif", "webp", "bmp", "ico", "svg"];
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
struct ReadParams {
|
||||
file_path: String,
|
||||
offset: Option<usize>,
|
||||
limit: Option<usize>,
|
||||
}
|
||||
|
||||
pub struct ReadTool;
|
||||
|
||||
fn is_binary(bytes: &[u8]) -> bool {
|
||||
bytes.iter().take(8000).any(|&b| b == 0)
|
||||
}
|
||||
|
||||
fn clamp_line(line: &str) -> String {
|
||||
if line.chars().count() <= LINE_CHAR_CLAMP {
|
||||
line.to_string()
|
||||
} else {
|
||||
let clamped: String = line.chars().take(LINE_CHAR_CLAMP).collect();
|
||||
format!("{clamped}... [line truncated]")
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ReadTool {
|
||||
fn name(&self) -> &str {
|
||||
"read"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Reads a file from the local filesystem, returning `cat -n`-style numbered lines."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::to_value(schemars::schema_for!(ReadParams)).unwrap()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: ReadParams =
|
||||
serde_json::from_value(input).map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
let path = paths::resolve(&ctx.cwd, ¶ms.file_path);
|
||||
let pattern = paths::relative_pattern(&ctx.cwd, &path);
|
||||
ctx.ask
|
||||
.ask(
|
||||
"read",
|
||||
pattern.clone(),
|
||||
pattern,
|
||||
serde_json::json!({"file_path": params.file_path}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
|
||||
if IMAGE_EXTENSIONS.contains(&ext.to_lowercase().as_str()) {
|
||||
return Ok(ToolOutput::new(
|
||||
params.file_path.clone(),
|
||||
format!(
|
||||
"{}: image file — binary content not shown",
|
||||
params.file_path
|
||||
),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let bytes = tokio::fs::read(&path)
|
||||
.await
|
||||
.map_err(|e| ToolError::Other(format!("{}: {e}", path.display())))?;
|
||||
if is_binary(&bytes) {
|
||||
return Ok(ToolOutput::new(
|
||||
params.file_path.clone(),
|
||||
format!("{}: binary file — content not shown", params.file_path),
|
||||
));
|
||||
}
|
||||
|
||||
let text = String::from_utf8_lossy(&bytes);
|
||||
let offset = params.offset.unwrap_or(0);
|
||||
let limit = params.limit.unwrap_or(DEFAULT_LIMIT);
|
||||
let numbered: Vec<String> = text
|
||||
.lines()
|
||||
.enumerate()
|
||||
.skip(offset)
|
||||
.take(limit)
|
||||
.map(|(i, line)| format!("{:>6}\t{}", i + 1, clamp_line(line)))
|
||||
.collect();
|
||||
|
||||
if numbered.is_empty() {
|
||||
return Ok(ToolOutput::new(
|
||||
params.file_path.clone(),
|
||||
"(empty file or offset past end)".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(ToolOutput::new(
|
||||
params.file_path.clone(),
|
||||
numbered.join("\n"),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use harness_core::event::EventBus;
|
||||
use harness_core::permission::{spawn_auto_approve, PermissionService};
|
||||
use harness_core::tool::{MetadataSink, PermissionHandle};
|
||||
use harness_core::types::SessionId;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
fn ctx(cwd: std::path::PathBuf) -> ToolCtx {
|
||||
let bus = EventBus::new();
|
||||
let service = Arc::new(PermissionService::new(bus.clone()));
|
||||
spawn_auto_approve(bus, service.clone());
|
||||
let (metadata, _rx) = MetadataSink::channel();
|
||||
ToolCtx {
|
||||
session_id: SessionId::new(),
|
||||
message_id: harness_core::types::MessageId::new(),
|
||||
call_id: "call_1".into(),
|
||||
data_dir: cwd.join("tool-output"),
|
||||
cwd,
|
||||
cancel: CancellationToken::new(),
|
||||
ask: PermissionHandle::new(
|
||||
service,
|
||||
SessionId::new(),
|
||||
Vec::new(),
|
||||
Arc::new(Mutex::new(Vec::new())),
|
||||
CancellationToken::new(),
|
||||
),
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reads_and_numbers_lines() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.txt"), "line one\nline two\nline three").unwrap();
|
||||
let output = ReadTool
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "a.txt"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
output.output,
|
||||
" 1\tline one\n 2\tline two\n 3\tline three"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn respects_offset_and_limit() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.txt"), "a\nb\nc\nd\ne").unwrap();
|
||||
let output = ReadTool
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "a.txt", "offset": 1, "limit": 2}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(output.output, " 2\tb\n 3\tc");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn detects_binary_content() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("bin.dat"), [0u8, 1, 2, 3]).unwrap();
|
||||
let output = ReadTool
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "bin.dat"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(output.output.contains("binary file"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_file_is_a_tool_error() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let err = ReadTool
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "missing.txt"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, ToolError::Other(_)));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
use async_trait::async_trait;
|
||||
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
use similar::TextDiff;
|
||||
|
||||
use crate::paths;
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
struct WriteParams {
|
||||
file_path: String,
|
||||
content: String,
|
||||
}
|
||||
|
||||
pub struct WriteTool;
|
||||
|
||||
fn unified_diff(before: &str, after: &str, file_path: &str) -> String {
|
||||
TextDiff::from_lines(before, after)
|
||||
.unified_diff()
|
||||
.header(&format!("a/{file_path}"), &format!("b/{file_path}"))
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for WriteTool {
|
||||
fn name(&self) -> &str {
|
||||
"write"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Writes content to a file, creating it (and parent directories) if needed."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::to_value(schemars::schema_for!(WriteParams)).unwrap()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: WriteParams =
|
||||
serde_json::from_value(input).map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
let path = paths::resolve(&ctx.cwd, ¶ms.file_path);
|
||||
let pattern = paths::relative_pattern(&ctx.cwd, &path);
|
||||
|
||||
let before = tokio::fs::read_to_string(&path).await.unwrap_or_default();
|
||||
let diff = unified_diff(&before, ¶ms.content, ¶ms.file_path);
|
||||
|
||||
// Same permission key as opencode's edit tool — writing is an edit.
|
||||
ctx.ask
|
||||
.ask(
|
||||
"edit",
|
||||
pattern,
|
||||
"*",
|
||||
serde_json::json!({"file_path": params.file_path, "diff": diff}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
if let Some(parent) = path.parent() {
|
||||
tokio::fs::create_dir_all(parent)
|
||||
.await
|
||||
.map_err(|e| ToolError::Other(format!("{}: {e}", parent.display())))?;
|
||||
}
|
||||
tokio::fs::write(&path, ¶ms.content)
|
||||
.await
|
||||
.map_err(|e| ToolError::Other(format!("{}: {e}", path.display())))?;
|
||||
|
||||
let added = diff
|
||||
.lines()
|
||||
.filter(|l| l.starts_with('+') && !l.starts_with("+++"))
|
||||
.count();
|
||||
let removed = diff
|
||||
.lines()
|
||||
.filter(|l| l.starts_with('-') && !l.starts_with("---"))
|
||||
.count();
|
||||
let mut output = ToolOutput::new(
|
||||
params.file_path.clone(),
|
||||
format!("wrote {} bytes", params.content.len()),
|
||||
);
|
||||
output.metadata = serde_json::json!({"diff": diff, "added": added, "removed": removed});
|
||||
Ok(output)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use harness_core::event::EventBus;
|
||||
use harness_core::permission::{spawn_auto_approve, PermissionService};
|
||||
use harness_core::tool::{MetadataSink, PermissionHandle};
|
||||
use harness_core::types::SessionId;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
fn ctx(cwd: std::path::PathBuf) -> ToolCtx {
|
||||
let bus = EventBus::new();
|
||||
let service = Arc::new(PermissionService::new(bus.clone()));
|
||||
spawn_auto_approve(bus, service.clone());
|
||||
let (metadata, _rx) = MetadataSink::channel();
|
||||
ToolCtx {
|
||||
session_id: SessionId::new(),
|
||||
message_id: harness_core::types::MessageId::new(),
|
||||
call_id: "call_1".into(),
|
||||
data_dir: cwd.join("tool-output"),
|
||||
cwd,
|
||||
cancel: CancellationToken::new(),
|
||||
ask: PermissionHandle::new(
|
||||
service,
|
||||
SessionId::new(),
|
||||
Vec::new(),
|
||||
Arc::new(Mutex::new(Vec::new())),
|
||||
CancellationToken::new(),
|
||||
),
|
||||
metadata,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn creates_file_and_parent_dirs() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
WriteTool
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "nested/dir/a.txt", "content": "hello"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let content = std::fs::read_to_string(dir.path().join("nested/dir/a.txt")).unwrap();
|
||||
assert_eq!(content, "hello");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn overwrites_existing_file_and_reports_diff() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.txt"), "old\n").unwrap();
|
||||
let output = WriteTool
|
||||
.execute(
|
||||
serde_json::json!({"file_path": "a.txt", "content": "new\n"}),
|
||||
ctx(dir.path().to_path_buf()),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let content = std::fs::read_to_string(dir.path().join("a.txt")).unwrap();
|
||||
assert_eq!(content, "new\n");
|
||||
assert_eq!(output.metadata["added"], 1);
|
||||
assert_eq!(output.metadata["removed"], 1);
|
||||
}
|
||||
}
|
||||
@@ -10,6 +10,8 @@ path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
harness-app = { workspace = true }
|
||||
harness-core = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -1,3 +1,86 @@
|
||||
fn main() {
|
||||
println!("harness {}", env!("CARGO_PKG_VERSION"));
|
||||
use harness_core::event::RunOutcome;
|
||||
|
||||
const DEFAULT_MODEL: &str = "anthropic/claude-sonnet-4-5";
|
||||
|
||||
fn print_usage() {
|
||||
eprintln!("usage: harness run -p \"<prompt>\" [-m provider/model]");
|
||||
}
|
||||
|
||||
fn parse_run_args(args: &[String]) -> Option<(String, Option<String>)> {
|
||||
let mut prompt = None;
|
||||
let mut model = None;
|
||||
let mut i = 0;
|
||||
while i < args.len() {
|
||||
match args[i].as_str() {
|
||||
"-p" | "--prompt" => {
|
||||
prompt = args.get(i + 1).cloned();
|
||||
i += 2;
|
||||
}
|
||||
"-m" | "--model" => {
|
||||
model = args.get(i + 1).cloned();
|
||||
i += 2;
|
||||
}
|
||||
_ => i += 1,
|
||||
}
|
||||
}
|
||||
prompt.map(|p| (p, model))
|
||||
}
|
||||
|
||||
async fn run(args: &[String]) -> i32 {
|
||||
let Some((prompt, model_arg)) = parse_run_args(args) else {
|
||||
print_usage();
|
||||
return 2;
|
||||
};
|
||||
|
||||
let cwd = match std::env::current_dir() {
|
||||
Ok(cwd) => cwd,
|
||||
Err(e) => {
|
||||
eprintln!("error: could not read cwd: {e}");
|
||||
return 1;
|
||||
}
|
||||
};
|
||||
|
||||
let app = match harness_app::App::init(cwd) {
|
||||
Ok(app) => app,
|
||||
Err(e) => {
|
||||
eprintln!("error: {e}");
|
||||
return 1;
|
||||
}
|
||||
};
|
||||
|
||||
let model_ref = model_arg
|
||||
.or_else(|| app.config.model.clone())
|
||||
.unwrap_or_else(|| DEFAULT_MODEL.to_string());
|
||||
|
||||
match app.run_prompt(prompt, &model_ref).await {
|
||||
Ok((RunOutcome::Stopped, session_id)) => {
|
||||
let text = app.final_text(&session_id).await.unwrap_or_default();
|
||||
println!("{text}");
|
||||
0
|
||||
}
|
||||
Ok((RunOutcome::Aborted, _)) => {
|
||||
eprintln!("run aborted");
|
||||
1
|
||||
}
|
||||
Ok((RunOutcome::Errored { message }, _)) => {
|
||||
eprintln!("error: {message}");
|
||||
1
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("error: {e}");
|
||||
1
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main(flavor = "current_thread")]
|
||||
async fn main() {
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
let exit_code = if args.first().map(String::as_str) == Some("run") {
|
||||
run(&args[1..]).await
|
||||
} else {
|
||||
println!("harness {}", env!("CARGO_PKG_VERSION"));
|
||||
0
|
||||
};
|
||||
std::process::exit(exit_code);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user