Compare commits
24 Commits
develop
...
feature/m6
| Author | SHA1 | Date | |
|---|---|---|---|
| 4ba2f7af74 | |||
| 08b3cd8fb7 | |||
| 5fd1698d47 | |||
| dbc676332a | |||
| d6af65f494 | |||
| fe6d00ce6d | |||
| f71a347061 | |||
| e32f61e4f8 | |||
| 9ee5a1940d | |||
| ac516a7e23 | |||
| 8976b5dc09 | |||
| 07476d1318 | |||
| f3b7666d3b | |||
| 8ed17bf091 | |||
| ab0269df1a | |||
| a99edc116a | |||
| a3f6fd6378 | |||
| 607bd88ae3 | |||
| d4a846f827 | |||
| d83f8c84b7 | |||
| 9ed278bcb5 | |||
| 4118279da1 | |||
| a5af873924 | |||
| 5ac646b3c6 |
Generated
+2167
-12
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,19 @@ harness-providers = { workspace = true }
|
||||
harness-tools = { workspace = true }
|
||||
harness-mcp = { workspace = true }
|
||||
harness-lsp = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
dirs = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
+1915
-1
File diff suppressed because it is too large
Load Diff
@@ -11,15 +11,19 @@ futures = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
serde_yaml_ng = { 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,12 @@
|
||||
---
|
||||
description: UI and interaction design specialist for user-facing surfaces
|
||||
mode: subagent
|
||||
temperature: 0.4
|
||||
tools: { task: false }
|
||||
---
|
||||
You are Designer, a UI and interaction design specialist. You handle user-facing surfaces:
|
||||
layout, component structure, styling, and interaction details.
|
||||
|
||||
- Understand the existing design language before proposing changes; stay consistent with it.
|
||||
- When implementing, make focused edits and describe the visual/interaction effect.
|
||||
- Call out accessibility and responsive concerns relevant to the change.
|
||||
@@ -0,0 +1,16 @@
|
||||
---
|
||||
description: Read-only reconnaissance specialist for mapping code and finding relevant files
|
||||
mode: subagent
|
||||
temperature: 0.1
|
||||
tools: { write: false, edit: false, task: false }
|
||||
permission:
|
||||
- { permission: "edit", pattern: "*", action: deny }
|
||||
- { permission: "write", pattern: "*", action: deny }
|
||||
---
|
||||
You are Explorer, a read-only reconnaissance specialist. You locate the code, files, and
|
||||
facts the orchestrator needs and report back concisely.
|
||||
|
||||
- Use read, glob, grep, and bash (read-only commands) to investigate.
|
||||
- Never modify files. Return a focused summary with concrete `path:line` references, not
|
||||
file dumps.
|
||||
- State what you found and, briefly, what you could not find.
|
||||
@@ -0,0 +1,13 @@
|
||||
---
|
||||
description: Implementation specialist that makes focused code changes and verifies them
|
||||
mode: subagent
|
||||
temperature: 0.2
|
||||
tools: { task: false }
|
||||
---
|
||||
You are Fixer, an implementation specialist. You take a well-scoped change, implement it,
|
||||
and verify it compiles/tests.
|
||||
|
||||
- Make the smallest change that satisfies the request; match the surrounding code's style.
|
||||
- Use read/grep to understand context before editing; use bash to build and run tests.
|
||||
- Report exactly what you changed (files and the essence of the diff) and the result of any
|
||||
verification you ran.
|
||||
@@ -0,0 +1,15 @@
|
||||
---
|
||||
description: Documentation and knowledge lookup specialist
|
||||
mode: subagent
|
||||
temperature: 0.1
|
||||
tools: { write: false, edit: false, task: false }
|
||||
permission:
|
||||
- { permission: "edit", pattern: "*", action: deny }
|
||||
- { permission: "write", pattern: "*", action: deny }
|
||||
---
|
||||
You are Librarian. You find and summarize documentation, comments, READMEs, config, and
|
||||
other in-repo knowledge on request.
|
||||
|
||||
- Search docs and source for the relevant material with read, glob, and grep.
|
||||
- Quote the authoritative source with its `path:line`; do not invent details.
|
||||
- Return a concise, well-organized summary with pointers back to the sources.
|
||||
@@ -0,0 +1,15 @@
|
||||
---
|
||||
description: Deep-reasoning analyst for architecture, debugging, and design trade-offs
|
||||
mode: subagent
|
||||
temperature: 0.3
|
||||
tools: { write: false, edit: false, task: false }
|
||||
permission:
|
||||
- { permission: "edit", pattern: "*", action: deny }
|
||||
- { permission: "write", pattern: "*", action: deny }
|
||||
---
|
||||
You are Oracle, a deep-reasoning analyst. You are consulted for hard questions:
|
||||
root-causing bugs, weighing architectural trade-offs, and reviewing designs.
|
||||
|
||||
- Read whatever code and context you need, but do not modify anything.
|
||||
- Reason carefully and explicitly; state assumptions and the evidence behind conclusions.
|
||||
- Return a decisive recommendation with the reasoning that supports it.
|
||||
@@ -0,0 +1,19 @@
|
||||
---
|
||||
description: Primary coordinator that plans work and delegates to specialists
|
||||
mode: primary
|
||||
temperature: 0.2
|
||||
---
|
||||
You are the orchestrator. You plan the work, delegate focused pieces to specialist
|
||||
subagents via the `task` tool, and synthesize their results into a final answer.
|
||||
|
||||
Guidelines:
|
||||
- Break the request into concrete, independently-verifiable pieces.
|
||||
- Prefer delegating reconnaissance and analysis to subagents so your own context stays
|
||||
focused; do the integration and final write-up yourself.
|
||||
- Launch background tasks for long-running independent work, then continue planning.
|
||||
Do not poll running jobs — wait for completion and reconcile terminal jobs before your
|
||||
final response.
|
||||
- Reuse a completed specialist session (by its job alias) when following up on the same
|
||||
thread of work.
|
||||
|
||||
{{SUBAGENTS}}
|
||||
@@ -0,0 +1,496 @@
|
||||
//! Agent definitions and registry.
|
||||
//!
|
||||
//! All agent *behavior* lives in markdown + config — the engine only understands `mode`, tool
|
||||
//! filters, permissions, model, and depth. Definitions are layered (bundled → global → project
|
||||
//! → config patch); later layers win by name. See `docs/04-multiagent.md`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::config::AgentPatch;
|
||||
use crate::permission::{Rule, Ruleset};
|
||||
use crate::types::ModelRef;
|
||||
|
||||
/// Marker in a primary agent's prompt replaced at load time with the routing list of enabled
|
||||
/// subagents (name + description). Keeps routing text in sync with the enabled agent set.
|
||||
const SUBAGENTS_MARKER: &str = "{{SUBAGENTS}}";
|
||||
|
||||
const BUNDLED: &[(&str, &str)] = &[
|
||||
(
|
||||
"orchestrator",
|
||||
include_str!("../../assets/agents/orchestrator.md"),
|
||||
),
|
||||
("explorer", include_str!("../../assets/agents/explorer.md")),
|
||||
("oracle", include_str!("../../assets/agents/oracle.md")),
|
||||
(
|
||||
"librarian",
|
||||
include_str!("../../assets/agents/librarian.md"),
|
||||
),
|
||||
("fixer", include_str!("../../assets/agents/fixer.md")),
|
||||
("designer", include_str!("../../assets/agents/designer.md")),
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
|
||||
pub enum AgentMode {
|
||||
Primary,
|
||||
#[default]
|
||||
Subagent,
|
||||
All,
|
||||
}
|
||||
|
||||
impl AgentMode {
|
||||
fn parse(s: &str) -> Option<Self> {
|
||||
match s.trim().to_ascii_lowercase().as_str() {
|
||||
"primary" => Some(Self::Primary),
|
||||
"subagent" => Some(Self::Subagent),
|
||||
"all" => Some(Self::All),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether this agent can be invoked as a subagent via the `task` tool.
|
||||
pub fn is_subagent(self) -> bool {
|
||||
matches!(self, Self::Subagent | Self::All)
|
||||
}
|
||||
|
||||
/// Whether this agent can drive a top-level (primary) session.
|
||||
pub fn is_primary(self) -> bool {
|
||||
matches!(self, Self::Primary | Self::All)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AgentSource {
|
||||
Bundled,
|
||||
Global,
|
||||
Project,
|
||||
Config,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AgentDef {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub mode: AgentMode,
|
||||
/// `None` = follow the session model.
|
||||
pub model: Option<ModelRef>,
|
||||
pub temperature: Option<f32>,
|
||||
pub prompt: String,
|
||||
pub permissions: Ruleset,
|
||||
/// Tool enable/disable overrides (wildcard keys allowed); absent = inherit default.
|
||||
pub tools: HashMap<String, bool>,
|
||||
pub max_steps: Option<u32>,
|
||||
pub source: AgentSource,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum AgentError {
|
||||
#[error("agent {0}: {1}")]
|
||||
Frontmatter(String, String),
|
||||
}
|
||||
|
||||
/// YAML frontmatter shape (opencode-compatible).
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(default)]
|
||||
struct Frontmatter {
|
||||
description: Option<String>,
|
||||
mode: Option<String>,
|
||||
model: Option<String>,
|
||||
temperature: Option<f32>,
|
||||
tools: Option<HashMap<String, bool>>,
|
||||
permission: Option<Vec<Rule>>,
|
||||
max_steps: Option<u32>,
|
||||
disable: Option<bool>,
|
||||
}
|
||||
|
||||
fn parse_model_ref(s: &str) -> Option<ModelRef> {
|
||||
s.split_once('/')
|
||||
.map(|(p, m)| ModelRef::new(p.trim(), m.trim()))
|
||||
}
|
||||
|
||||
/// Splits a markdown agent file into (frontmatter, body). A file without a leading `---`
|
||||
/// fence is treated as an all-body prompt with empty frontmatter.
|
||||
fn split_frontmatter(content: &str) -> (&str, &str) {
|
||||
let rest = match content
|
||||
.strip_prefix("---\n")
|
||||
.or_else(|| content.strip_prefix("---\r\n"))
|
||||
{
|
||||
Some(r) => r,
|
||||
None => return ("", content),
|
||||
};
|
||||
// Find the closing fence line.
|
||||
for delim in ["\n---\n", "\n---\r\n"] {
|
||||
if let Some(idx) = rest.find(delim) {
|
||||
let body_start = idx + delim.len();
|
||||
return (&rest[..idx], &rest[body_start..]);
|
||||
}
|
||||
}
|
||||
// Trailing fence with no body / no trailing newline.
|
||||
if let Some(fm) = rest.strip_suffix("\n---").or(Some(rest)) {
|
||||
if rest.ends_with("\n---") {
|
||||
return (fm, "");
|
||||
}
|
||||
}
|
||||
("", content)
|
||||
}
|
||||
|
||||
/// Parses one markdown agent definition. Returns `Ok(None)` when the file marks itself
|
||||
/// `disable: true`.
|
||||
fn parse_agent(
|
||||
name: &str,
|
||||
source: AgentSource,
|
||||
content: &str,
|
||||
) -> Result<Option<AgentDef>, AgentError> {
|
||||
let (fm_raw, body) = split_frontmatter(content);
|
||||
let fm: Frontmatter = if fm_raw.trim().is_empty() {
|
||||
Frontmatter::default()
|
||||
} else {
|
||||
serde_yaml_ng::from_str(fm_raw)
|
||||
.map_err(|e| AgentError::Frontmatter(name.to_string(), e.to_string()))?
|
||||
};
|
||||
|
||||
if fm.disable == Some(true) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(AgentDef {
|
||||
name: name.to_string(),
|
||||
description: fm.description.unwrap_or_default(),
|
||||
mode: fm
|
||||
.mode
|
||||
.as_deref()
|
||||
.and_then(AgentMode::parse)
|
||||
.unwrap_or_default(),
|
||||
model: fm.model.as_deref().and_then(parse_model_ref),
|
||||
temperature: fm.temperature,
|
||||
prompt: body.trim_end().to_string(),
|
||||
permissions: fm.permission.unwrap_or_default(),
|
||||
tools: fm.tools.unwrap_or_default(),
|
||||
max_steps: fm.max_steps,
|
||||
source,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Applies a config `AgentPatch` onto an existing definition (only set fields override).
|
||||
fn apply_patch(def: &mut AgentDef, patch: &AgentPatch) {
|
||||
if let Some(mode) = patch.mode.as_deref().and_then(AgentMode::parse) {
|
||||
def.mode = mode;
|
||||
}
|
||||
if let Some(model) = patch.model.as_deref().and_then(parse_model_ref) {
|
||||
def.model = Some(model);
|
||||
}
|
||||
if let Some(temp) = patch.temperature {
|
||||
def.temperature = Some(temp);
|
||||
}
|
||||
if let Some(prompt) = &patch.prompt {
|
||||
def.prompt = prompt.clone();
|
||||
}
|
||||
if let Some(tools) = &patch.tools {
|
||||
def.tools.extend(tools.clone());
|
||||
}
|
||||
if let Some(permission) = &patch.permission {
|
||||
def.permissions = permission.clone();
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct AgentRegistry {
|
||||
agents: HashMap<String, AgentDef>,
|
||||
}
|
||||
|
||||
impl AgentRegistry {
|
||||
pub fn get(&self, name: &str) -> Option<&AgentDef> {
|
||||
self.agents.get(name)
|
||||
}
|
||||
|
||||
pub fn all(&self) -> Vec<&AgentDef> {
|
||||
self.agents.values().collect()
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.agents.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.agents.is_empty()
|
||||
}
|
||||
|
||||
/// Just the bundled agents — the default when no overrides are configured (tests, headless).
|
||||
pub fn bundled() -> Self {
|
||||
let mut reg = Self::default();
|
||||
reg.load_markdown_layer(
|
||||
BUNDLED.iter().map(|(n, c)| (n.to_string(), *c)),
|
||||
AgentSource::Bundled,
|
||||
);
|
||||
reg.generate_routing();
|
||||
reg
|
||||
}
|
||||
|
||||
/// Full layered load: bundled → global dir → project dir → config patches.
|
||||
pub fn load(
|
||||
config_agents: &HashMap<String, AgentPatch>,
|
||||
global_dir: Option<&Path>,
|
||||
project_dir: Option<&Path>,
|
||||
) -> Self {
|
||||
let mut reg = Self::default();
|
||||
reg.load_markdown_layer(
|
||||
BUNDLED.iter().map(|(n, c)| (n.to_string(), *c)),
|
||||
AgentSource::Bundled,
|
||||
);
|
||||
if let Some(dir) = global_dir {
|
||||
reg.load_dir(dir, AgentSource::Global);
|
||||
}
|
||||
if let Some(dir) = project_dir {
|
||||
reg.load_dir(dir, AgentSource::Project);
|
||||
}
|
||||
reg.apply_config(config_agents);
|
||||
reg.generate_routing();
|
||||
reg
|
||||
}
|
||||
|
||||
fn load_markdown_layer<I, S>(&mut self, files: I, source: AgentSource)
|
||||
where
|
||||
I: IntoIterator<Item = (String, S)>,
|
||||
S: AsRef<str>,
|
||||
{
|
||||
for (name, content) in files {
|
||||
match parse_agent(&name, source, content.as_ref()) {
|
||||
Ok(Some(def)) => {
|
||||
self.agents.insert(name, def);
|
||||
}
|
||||
Ok(None) => {
|
||||
self.agents.remove(&name); // disable: true removes an earlier layer
|
||||
}
|
||||
Err(e) => tracing::warn!(error = %e, "skipping malformed agent"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn load_dir(&mut self, dir: &Path, source: AgentSource) {
|
||||
let Ok(entries) = std::fs::read_dir(dir) else {
|
||||
return;
|
||||
};
|
||||
let mut files: Vec<(String, String)> = Vec::new();
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("md") {
|
||||
continue;
|
||||
}
|
||||
let Some(name) = path.file_stem().and_then(|s| s.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
if let Ok(content) = std::fs::read_to_string(&path) {
|
||||
files.push((name.to_string(), content));
|
||||
}
|
||||
}
|
||||
files.sort();
|
||||
self.load_markdown_layer(files, source);
|
||||
}
|
||||
|
||||
fn apply_config(&mut self, config_agents: &HashMap<String, AgentPatch>) {
|
||||
for (name, patch) in config_agents {
|
||||
if patch.disable == Some(true) {
|
||||
self.agents.remove(name);
|
||||
continue;
|
||||
}
|
||||
if let Some(def) = self.agents.get_mut(name) {
|
||||
apply_patch(def, patch);
|
||||
} else if patch.model.is_some() || patch.prompt.is_some() {
|
||||
// Unknown name with enough to stand on its own → custom agent.
|
||||
let mut def = AgentDef {
|
||||
name: name.clone(),
|
||||
description: String::new(),
|
||||
mode: AgentMode::default(),
|
||||
model: None,
|
||||
temperature: None,
|
||||
prompt: String::new(),
|
||||
permissions: Vec::new(),
|
||||
tools: HashMap::new(),
|
||||
max_steps: None,
|
||||
source: AgentSource::Config,
|
||||
};
|
||||
apply_patch(&mut def, patch);
|
||||
self.agents.insert(name.clone(), def);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Replaces `{{SUBAGENTS}}` in every primary agent's prompt with a generated routing list
|
||||
/// of the enabled subagents (so disabling an agent removes it from routing text).
|
||||
fn generate_routing(&mut self) {
|
||||
let mut subagents: Vec<(String, String)> = self
|
||||
.agents
|
||||
.values()
|
||||
.filter(|a| a.mode.is_subagent())
|
||||
.map(|a| (a.name.clone(), a.description.clone()))
|
||||
.collect();
|
||||
subagents.sort();
|
||||
|
||||
let routing = if subagents.is_empty() {
|
||||
"## Agents\n\nNo specialist subagents are available.".to_string()
|
||||
} else {
|
||||
let mut s = String::from("## Agents\n\nDelegate to these specialists via `task`:\n");
|
||||
for (name, desc) in &subagents {
|
||||
s.push_str(&format!("- **{name}** — {desc}\n"));
|
||||
}
|
||||
s
|
||||
};
|
||||
|
||||
for agent in self.agents.values_mut() {
|
||||
if agent.prompt.contains(SUBAGENTS_MARKER) {
|
||||
agent.prompt = agent.prompt.replace(SUBAGENTS_MARKER, routing.trim_end());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::permission::Action;
|
||||
|
||||
#[test]
|
||||
fn bundled_loads_all_six_agents() {
|
||||
let reg = AgentRegistry::bundled();
|
||||
for name in [
|
||||
"orchestrator",
|
||||
"explorer",
|
||||
"oracle",
|
||||
"librarian",
|
||||
"fixer",
|
||||
"designer",
|
||||
] {
|
||||
assert!(reg.get(name).is_some(), "missing {name}");
|
||||
}
|
||||
assert_eq!(reg.len(), 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_mode_model_temperature_tools_and_permissions() {
|
||||
let md = "---\n\
|
||||
description: test agent\n\
|
||||
mode: subagent\n\
|
||||
model: anthropic/claude-haiku-4-5\n\
|
||||
temperature: 0.1\n\
|
||||
tools: { write: false, bash: true }\n\
|
||||
permission:\n\
|
||||
\x20 - { permission: \"edit\", pattern: \"*\", action: deny }\n\
|
||||
---\n\
|
||||
You are a test agent.\n";
|
||||
let def = parse_agent("tester", AgentSource::Bundled, md)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(def.description, "test agent");
|
||||
assert_eq!(def.mode, AgentMode::Subagent);
|
||||
assert_eq!(
|
||||
def.model,
|
||||
Some(ModelRef::new("anthropic", "claude-haiku-4-5"))
|
||||
);
|
||||
assert_eq!(def.temperature, Some(0.1));
|
||||
assert_eq!(def.tools.get("write"), Some(&false));
|
||||
assert_eq!(def.tools.get("bash"), Some(&true));
|
||||
assert_eq!(def.permissions.len(), 1);
|
||||
assert_eq!(def.permissions[0].action, Action::Deny);
|
||||
assert_eq!(def.prompt, "You are a test agent.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_without_frontmatter_is_all_prompt() {
|
||||
let def = parse_agent("x", AgentSource::Global, "just a prompt")
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(def.prompt, "just a prompt");
|
||||
assert_eq!(def.mode, AgentMode::Subagent); // default
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disable_true_removes_agent() {
|
||||
assert!(
|
||||
parse_agent("x", AgentSource::Config, "---\ndisable: true\n---\nbody")
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explorer_denies_writes_and_disables_edit_tool() {
|
||||
let reg = AgentRegistry::bundled();
|
||||
let explorer = reg.get("explorer").unwrap();
|
||||
assert_eq!(explorer.mode, AgentMode::Subagent);
|
||||
assert_eq!(explorer.tools.get("edit"), Some(&false));
|
||||
assert!(explorer
|
||||
.permissions
|
||||
.iter()
|
||||
.any(|r| r.permission == "write" && r.action == Action::Deny));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn orchestrator_routing_lists_subagents_and_drops_marker() {
|
||||
let reg = AgentRegistry::bundled();
|
||||
let prompt = ®.get("orchestrator").unwrap().prompt;
|
||||
assert!(!prompt.contains("{{SUBAGENTS}}"));
|
||||
assert!(prompt.contains("explorer"));
|
||||
assert!(prompt.contains("fixer"));
|
||||
// The orchestrator itself is primary and must not list itself.
|
||||
assert!(!prompt.contains("- **orchestrator**"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_patch_overrides_only_set_fields() {
|
||||
let mut patches = HashMap::new();
|
||||
patches.insert(
|
||||
"explorer".to_string(),
|
||||
AgentPatch {
|
||||
temperature: Some(0.9),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let reg = AgentRegistry::load(&patches, None, None);
|
||||
let explorer = reg.get("explorer").unwrap();
|
||||
assert_eq!(explorer.temperature, Some(0.9)); // overridden
|
||||
assert_eq!(explorer.mode, AgentMode::Subagent); // untouched
|
||||
assert_eq!(explorer.tools.get("edit"), Some(&false)); // untouched
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_disable_removes_and_unknown_with_model_creates() {
|
||||
let mut patches = HashMap::new();
|
||||
patches.insert(
|
||||
"designer".to_string(),
|
||||
AgentPatch {
|
||||
disable: Some(true),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
patches.insert(
|
||||
"custom".to_string(),
|
||||
AgentPatch {
|
||||
model: Some("openai/gpt-5".into()),
|
||||
prompt: Some("custom prompt".into()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let reg = AgentRegistry::load(&patches, None, None);
|
||||
assert!(reg.get("designer").is_none());
|
||||
let custom = reg.get("custom").unwrap();
|
||||
assert_eq!(custom.source, AgentSource::Config);
|
||||
assert_eq!(custom.model, Some(ModelRef::new("openai", "gpt-5")));
|
||||
assert_eq!(custom.prompt, "custom prompt");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_config_agent_without_model_or_prompt_is_ignored() {
|
||||
let mut patches = HashMap::new();
|
||||
patches.insert(
|
||||
"ghost".to_string(),
|
||||
AgentPatch {
|
||||
temperature: Some(0.5),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
let reg = AgentRegistry::load(&patches, None, None);
|
||||
assert!(reg.get("ghost").is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,270 @@
|
||||
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"),
|
||||
("opencode", "OPENCODE_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,296 @@
|
||||
//! Markdown-defined commands and skills (M5). Both are YAML-frontmatter + body files loaded
|
||||
//! from a global dir (`~/.config/ai-harness/<kind>/`) and a project dir (`<project>/.harness/
|
||||
//! <kind>/`), project winning by name. See `docs/06-config.md` and `docs/09-integrations.md`.
|
||||
//!
|
||||
//! - Commands (`command/*.md`) are a pure input-layer concern: `/name args` expands the body
|
||||
//! template (`$ARGUMENTS`, `$1..$9`) into the user message, optionally switching agent/model.
|
||||
//! - Skills (`skill/<name>/SKILL.md`) advertise `name + description` in the system prompt; the
|
||||
//! model pulls a skill's body on demand via the built-in `skill` tool.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
/// A slash command: a named prompt template with an optional agent/model override.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CommandDef {
|
||||
/// Invocation name (the file stem); used as `/name`.
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
/// Run the expanded prompt under this agent instead of the session's default.
|
||||
pub agent: Option<String>,
|
||||
/// Run under this `provider/model` instead of the session's default.
|
||||
pub model: Option<String>,
|
||||
/// The body, with `$ARGUMENTS` / `$1..$9` placeholders.
|
||||
pub template: String,
|
||||
}
|
||||
|
||||
impl CommandDef {
|
||||
/// Substitutes `$ARGUMENTS` (the whole argument string) and `$1..$9` (whitespace-split
|
||||
/// positionals; missing ones become empty) into the template.
|
||||
pub fn expand(&self, arguments: &str) -> String {
|
||||
let positionals: Vec<&str> = arguments.split_whitespace().collect();
|
||||
let mut out = self.template.replace("$ARGUMENTS", arguments);
|
||||
for i in 1..=9 {
|
||||
let value = positionals.get(i - 1).copied().unwrap_or("");
|
||||
out = out.replace(&format!("${i}"), value);
|
||||
}
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
/// A skill: advertised by `name + description`, body loaded on demand by the `skill` tool.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct SkillDef {
|
||||
/// Skill name (the containing directory name); used as the `skill` tool argument.
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub body: String,
|
||||
}
|
||||
|
||||
/// Frontmatter fields shared by commands (all optional).
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(default)]
|
||||
struct CommandFrontmatter {
|
||||
description: Option<String>,
|
||||
agent: Option<String>,
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(default)]
|
||||
struct SkillFrontmatter {
|
||||
description: Option<String>,
|
||||
}
|
||||
|
||||
/// Splits `---\n…\n---\n` frontmatter from the body. A file with no leading fence is all body.
|
||||
fn split_frontmatter(content: &str) -> (&str, &str) {
|
||||
let rest = match content
|
||||
.strip_prefix("---\n")
|
||||
.or_else(|| content.strip_prefix("---\r\n"))
|
||||
{
|
||||
Some(r) => r,
|
||||
None => return ("", content),
|
||||
};
|
||||
for delim in ["\n---\n", "\n---\r\n"] {
|
||||
if let Some(idx) = rest.find(delim) {
|
||||
return (&rest[..idx], &rest[idx + delim.len()..]);
|
||||
}
|
||||
}
|
||||
if rest.ends_with("\n---") {
|
||||
return (rest.trim_end_matches("\n---"), "");
|
||||
}
|
||||
("", content)
|
||||
}
|
||||
|
||||
fn parse_command(name: &str, content: &str) -> CommandDef {
|
||||
let (fm_raw, body) = split_frontmatter(content);
|
||||
let fm: CommandFrontmatter = if fm_raw.trim().is_empty() {
|
||||
CommandFrontmatter::default()
|
||||
} else {
|
||||
serde_yaml_ng::from_str(fm_raw).unwrap_or_default()
|
||||
};
|
||||
CommandDef {
|
||||
name: name.to_string(),
|
||||
description: fm.description.unwrap_or_default(),
|
||||
agent: fm.agent,
|
||||
model: fm.model,
|
||||
template: body.trim().to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_skill(name: &str, content: &str) -> SkillDef {
|
||||
let (fm_raw, body) = split_frontmatter(content);
|
||||
let fm: SkillFrontmatter = if fm_raw.trim().is_empty() {
|
||||
SkillFrontmatter::default()
|
||||
} else {
|
||||
serde_yaml_ng::from_str(fm_raw).unwrap_or_default()
|
||||
};
|
||||
SkillDef {
|
||||
name: name.to_string(),
|
||||
description: fm.description.unwrap_or_default(),
|
||||
body: body.trim().to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Loads `command/*.md` from global then project dirs (project wins by name).
|
||||
pub fn load_commands(
|
||||
global_dir: Option<&Path>,
|
||||
project_dir: Option<&Path>,
|
||||
) -> HashMap<String, CommandDef> {
|
||||
let mut commands = HashMap::new();
|
||||
for dir in [global_dir, project_dir].into_iter().flatten() {
|
||||
let Ok(entries) = std::fs::read_dir(dir) else {
|
||||
continue;
|
||||
};
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|e| e.to_str()) != Some("md") {
|
||||
continue;
|
||||
}
|
||||
let Some(name) = path.file_stem().and_then(|s| s.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
if let Ok(content) = std::fs::read_to_string(&path) {
|
||||
commands.insert(name.to_string(), parse_command(name, &content));
|
||||
}
|
||||
}
|
||||
}
|
||||
commands
|
||||
}
|
||||
|
||||
/// Loads `skill/<name>/SKILL.md` from global then project dirs (project wins by name),
|
||||
/// returned name-sorted for a stable system-prompt listing.
|
||||
pub fn load_skills(global_dir: Option<&Path>, project_dir: Option<&Path>) -> Vec<SkillDef> {
|
||||
let mut by_name: HashMap<String, SkillDef> = HashMap::new();
|
||||
for dir in [global_dir, project_dir].into_iter().flatten() {
|
||||
let Ok(entries) = std::fs::read_dir(dir) else {
|
||||
continue;
|
||||
};
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if !path.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let Some(name) = path.file_name().and_then(|s| s.to_str()) else {
|
||||
continue;
|
||||
};
|
||||
let skill_file = path.join("SKILL.md");
|
||||
if let Ok(content) = std::fs::read_to_string(&skill_file) {
|
||||
by_name.insert(name.to_string(), parse_skill(name, &content));
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut skills: Vec<SkillDef> = by_name.into_values().collect();
|
||||
skills.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
skills
|
||||
}
|
||||
|
||||
/// The system-prompt section advertising available skills (name + description). `None` when
|
||||
/// there are no skills, so no empty section is injected.
|
||||
pub fn skills_prompt(skills: &[SkillDef]) -> Option<String> {
|
||||
if skills.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut section = String::from(
|
||||
"## Skills\n\nThese skills are available. Load a skill's full instructions on demand \
|
||||
by calling the `skill` tool with its name before doing the related work:\n",
|
||||
);
|
||||
for skill in skills {
|
||||
section.push_str(&format!("- **{}** — {}\n", skill.name, skill.description));
|
||||
}
|
||||
Some(section.trim_end().to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn expand_substitutes_arguments_and_positionals() {
|
||||
let cmd = CommandDef {
|
||||
name: "greet".into(),
|
||||
description: String::new(),
|
||||
agent: None,
|
||||
model: None,
|
||||
template: "Say $1 to $2. All: $ARGUMENTS".into(),
|
||||
};
|
||||
assert_eq!(cmd.expand("hi there"), "Say hi to there. All: hi there");
|
||||
// Missing positionals collapse to empty.
|
||||
assert_eq!(cmd.expand("solo"), "Say solo to . All: solo");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_command_reads_frontmatter_and_body() {
|
||||
let md = "---\ndescription: review a PR\nagent: oracle\nmodel: openai/gpt-5\n---\nReview $ARGUMENTS please.\n";
|
||||
let cmd = parse_command("review", md);
|
||||
assert_eq!(cmd.description, "review a PR");
|
||||
assert_eq!(cmd.agent.as_deref(), Some("oracle"));
|
||||
assert_eq!(cmd.model.as_deref(), Some("openai/gpt-5"));
|
||||
assert_eq!(cmd.template, "Review $ARGUMENTS please.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_command_without_frontmatter_is_all_template() {
|
||||
let cmd = parse_command("x", "just do $1");
|
||||
assert_eq!(cmd.template, "just do $1");
|
||||
assert!(cmd.agent.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_skill_reads_description_and_body() {
|
||||
let md = "---\ndescription: format code\n---\nRun the formatter.\n";
|
||||
let skill = parse_skill("formatter", md);
|
||||
assert_eq!(skill.name, "formatter");
|
||||
assert_eq!(skill.description, "format code");
|
||||
assert_eq!(skill.body, "Run the formatter.");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skills_prompt_lists_each_and_is_none_when_empty() {
|
||||
assert!(skills_prompt(&[]).is_none());
|
||||
let skills = vec![
|
||||
SkillDef {
|
||||
name: "a".into(),
|
||||
description: "does a".into(),
|
||||
body: "…".into(),
|
||||
},
|
||||
SkillDef {
|
||||
name: "b".into(),
|
||||
description: "does b".into(),
|
||||
body: "…".into(),
|
||||
},
|
||||
];
|
||||
let prompt = skills_prompt(&skills).unwrap();
|
||||
assert!(prompt.contains("- **a** — does a"));
|
||||
assert!(prompt.contains("- **b** — does b"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_skills_reads_skill_dirs_and_project_wins() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let global = dir.path().join("global");
|
||||
let project = dir.path().join("project");
|
||||
std::fs::create_dir_all(global.join("fmt")).unwrap();
|
||||
std::fs::create_dir_all(project.join("fmt")).unwrap();
|
||||
std::fs::create_dir_all(global.join("lint")).unwrap();
|
||||
std::fs::write(
|
||||
global.join("fmt/SKILL.md"),
|
||||
"---\ndescription: global fmt\n---\nglobal body",
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(
|
||||
project.join("fmt/SKILL.md"),
|
||||
"---\ndescription: project fmt\n---\nproject body",
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::write(
|
||||
global.join("lint/SKILL.md"),
|
||||
"---\ndescription: lint\n---\nlint body",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let skills = load_skills(Some(&global), Some(&project));
|
||||
assert_eq!(skills.len(), 2);
|
||||
// Sorted by name: fmt, lint.
|
||||
assert_eq!(skills[0].name, "fmt");
|
||||
assert_eq!(skills[0].description, "project fmt"); // project overrode global
|
||||
assert_eq!(skills[0].body, "project body");
|
||||
assert_eq!(skills[1].name, "lint");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_commands_project_overrides_global() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let global = dir.path().join("g");
|
||||
let project = dir.path().join("p");
|
||||
std::fs::create_dir_all(&global).unwrap();
|
||||
std::fs::create_dir_all(&project).unwrap();
|
||||
std::fs::write(global.join("deploy.md"), "global deploy").unwrap();
|
||||
std::fs::write(project.join("deploy.md"), "project deploy").unwrap();
|
||||
let commands = load_commands(Some(&global), Some(&project));
|
||||
assert_eq!(commands.len(), 1);
|
||||
assert_eq!(commands["deploy"].template, "project deploy");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
pub mod load;
|
||||
pub mod markdown;
|
||||
pub mod schema;
|
||||
|
||||
pub use load::{load, ConfigError};
|
||||
pub use markdown::{load_commands, load_skills, skills_prompt, CommandDef, SkillDef};
|
||||
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,146 @@
|
||||
//! Auto-compaction (M6). When a session's context approaches the model's window the loop
|
||||
//! summarizes the conversation so far via a small model, writes a `Compaction` marker, and
|
||||
//! continues — subsequent requests replace the summarized history with the summary. See
|
||||
//! `docs/02-engine.md`.
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::event::AppEvent;
|
||||
use crate::llm::{ProviderError, WireMessage};
|
||||
use crate::types::{Message, MessageId, Part, PartBody, PartId, SessionId};
|
||||
|
||||
use super::processor::StepContext;
|
||||
|
||||
/// Produces a compact recap of a conversation. Implemented by the composition root over the
|
||||
/// configured `small_model`; absent in headless/test contexts (compaction then disabled).
|
||||
#[async_trait]
|
||||
pub trait Compactor: Send + Sync {
|
||||
async fn summarize(&self, messages: &[WireMessage]) -> Result<String, ProviderError>;
|
||||
}
|
||||
|
||||
/// System prompt handed to the small model to summarize the conversation for continuation.
|
||||
pub const SUMMARY_SYSTEM_PROMPT: &str = "You are compacting a long coding-assistant \
|
||||
conversation so it can continue within a smaller context window. Write a dense, factual \
|
||||
summary that preserves: the user's goal and constraints, decisions made and why, files \
|
||||
and symbols touched, commands run and their results, and the exact next step in progress. \
|
||||
Omit pleasantries. Output only the summary.";
|
||||
|
||||
/// Frames a raw summary as the synthetic user turn the model sees after compaction.
|
||||
pub fn frame_summary(summary: &str) -> String {
|
||||
format!(
|
||||
"The earlier part of this conversation was summarized to save context. \
|
||||
Continue from this summary:\n\n{summary}"
|
||||
)
|
||||
}
|
||||
|
||||
/// Persists a compaction: a new user message carrying the `Compaction` marker (used to cut
|
||||
/// history on the next load) plus a synthetic text part holding the framed summary (what the
|
||||
/// model actually reads). Returns the new message id.
|
||||
pub async fn write_compaction(
|
||||
ctx: &StepContext,
|
||||
session_id: &SessionId,
|
||||
replaces_up_to: MessageId,
|
||||
summary: String,
|
||||
now: i64,
|
||||
) -> Result<MessageId, ProviderError> {
|
||||
let message = Message::new_user(session_id.clone(), now);
|
||||
let message_id = message.id.clone();
|
||||
let store_err = |e: crate::store::StoreError| ProviderError::Decode(e.to_string());
|
||||
|
||||
ctx.store
|
||||
.upsert_message(message.clone())
|
||||
.await
|
||||
.map_err(store_err)?;
|
||||
ctx.bus.publish(AppEvent::MessageCreated {
|
||||
message: message.clone(),
|
||||
});
|
||||
|
||||
let marker = Part {
|
||||
id: PartId::new(),
|
||||
message_id: message_id.clone(),
|
||||
session_id: session_id.clone(),
|
||||
idx: 0,
|
||||
body: PartBody::Compaction {
|
||||
replaces_up_to,
|
||||
summary: summary.clone(),
|
||||
},
|
||||
};
|
||||
ctx.store.upsert_part(marker).await.map_err(store_err)?;
|
||||
|
||||
let text = Part {
|
||||
id: PartId::new(),
|
||||
message_id: message_id.clone(),
|
||||
session_id: session_id.clone(),
|
||||
idx: 1,
|
||||
body: PartBody::Text {
|
||||
text: frame_summary(&summary),
|
||||
synthetic: true,
|
||||
},
|
||||
};
|
||||
ctx.store
|
||||
.upsert_part(text.clone())
|
||||
.await
|
||||
.map_err(store_err)?;
|
||||
ctx.bus.publish(AppEvent::PartUpdated { part: text });
|
||||
|
||||
Ok(message_id)
|
||||
}
|
||||
|
||||
/// Index of the last message carrying a `Compaction` part, given each message's parts. History
|
||||
/// before it is dropped when building the request; `None` means no compaction yet.
|
||||
pub fn last_compaction_index(parts_per_message: &[Vec<Part>]) -> Option<usize> {
|
||||
parts_per_message.iter().rposition(|parts| {
|
||||
parts
|
||||
.iter()
|
||||
.any(|p| matches!(p.body, PartBody::Compaction { .. }))
|
||||
})
|
||||
}
|
||||
|
||||
/// Whether accumulated `used` tokens have crossed the compaction threshold (90% of the
|
||||
/// model's context window). `context_limit == 0` (unknown) disables the trigger.
|
||||
pub fn over_threshold(used: u64, context_limit: u64) -> bool {
|
||||
context_limit > 0 && used.saturating_mul(10) > context_limit.saturating_mul(9)
|
||||
}
|
||||
|
||||
/// Effective context occupancy for a step's usage: prompt (incl. cache) plus generated output.
|
||||
pub fn tokens_used(usage: &crate::types::TokenUsage) -> u64 {
|
||||
usage.input + usage.cache_read + usage.cache_write + usage.output
|
||||
}
|
||||
|
||||
/// True if a message list has real history to compact before `cut_start` (avoids a useless
|
||||
/// compaction that would summarize nothing and could loop).
|
||||
pub fn has_history_to_compact(messages: &[Message], cut_start: usize) -> bool {
|
||||
messages.len().saturating_sub(cut_start) > 1
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn over_threshold_respects_ninety_percent_and_unknown_limit() {
|
||||
assert!(over_threshold(91, 100));
|
||||
assert!(!over_threshold(90, 100)); // exactly 90% is not yet over
|
||||
assert!(!over_threshold(50, 100));
|
||||
assert!(!over_threshold(1_000_000, 0)); // unknown limit disables the trigger
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tokens_used_sums_prompt_cache_and_output() {
|
||||
let usage = crate::types::TokenUsage {
|
||||
input: 10,
|
||||
output: 5,
|
||||
reasoning: 0,
|
||||
cache_read: 3,
|
||||
cache_write: 2,
|
||||
};
|
||||
assert_eq!(tokens_used(&usage), 20);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn frame_summary_wraps_text() {
|
||||
let framed = frame_summary("did X");
|
||||
assert!(framed.contains("did X"));
|
||||
assert!(framed.contains("summarized"));
|
||||
}
|
||||
}
|
||||
@@ -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,579 @@
|
||||
//! Background job board — tracks subagent tasks spawned via the `task` tool so the
|
||||
//! orchestrator can see running work, reconcile terminal results, and reuse completed
|
||||
//! child sessions by alias. Simplified native port of oh-my-opencode-slim's
|
||||
//! `background-job-board.ts`. See `docs/04-multiagent.md`.
|
||||
//!
|
||||
//! The board is an in-memory `RwLock<HashMap>` mirrored to the `job` table so it survives a
|
||||
//! resume. All mutations persist through the passed-in `Store`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::event::{AppEvent, EventBus, JobRecordEvent};
|
||||
use crate::store::{Store, StoreError};
|
||||
use crate::types::SessionId;
|
||||
|
||||
/// A file a child session read, surfaced on the board so the orchestrator knows what a
|
||||
/// completed specialist already looked at.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ContextFile {
|
||||
pub path: String,
|
||||
pub lines: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum JobState {
|
||||
Running,
|
||||
Completed,
|
||||
Error,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
impl JobState {
|
||||
/// Terminal jobs are candidates for reconciliation; a completed one is reusable.
|
||||
pub fn is_terminal(self) -> bool {
|
||||
!matches!(self, JobState::Running)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct JobRecord {
|
||||
pub task_id: String,
|
||||
/// Human-friendly handle: first 3 chars of the agent name + a per-agent counter (`exp-1`).
|
||||
pub alias: String,
|
||||
pub parent_session: SessionId,
|
||||
pub child_session: SessionId,
|
||||
pub agent: String,
|
||||
pub description: String,
|
||||
pub objective: Option<String>,
|
||||
pub state: JobState,
|
||||
/// Whether the orchestrator has already seen this job's terminal result.
|
||||
pub reconciled: bool,
|
||||
pub result_summary: Option<String>,
|
||||
pub context_files: Vec<ContextFile>,
|
||||
pub launched_at: i64,
|
||||
pub updated_at: i64,
|
||||
pub last_used_at: i64,
|
||||
}
|
||||
|
||||
/// Parameters for registering a newly launched task on the board.
|
||||
pub struct LaunchSpec {
|
||||
pub task_id: String,
|
||||
pub parent_session: SessionId,
|
||||
pub child_session: SessionId,
|
||||
pub agent: String,
|
||||
pub description: String,
|
||||
pub objective: Option<String>,
|
||||
}
|
||||
|
||||
/// Result summaries are truncated to this many chars on the board.
|
||||
const SUMMARY_MAX: usize = 2000;
|
||||
/// Context files shown per job in the prompt injection.
|
||||
const CONTEXT_FILES_SHOWN: usize = 8;
|
||||
/// A read must cover at least this many lines to be worth reporting to the board.
|
||||
pub const MIN_REPORTED_LINES: u32 = 10;
|
||||
|
||||
pub struct JobBoard {
|
||||
store: Store,
|
||||
bus: EventBus,
|
||||
jobs: RwLock<HashMap<String, JobRecord>>,
|
||||
max_reusable_per_agent: u32,
|
||||
}
|
||||
|
||||
impl JobBoard {
|
||||
/// Builds a board and loads any persisted jobs for `parent_session`'s tree from the store.
|
||||
pub async fn load(
|
||||
store: Store,
|
||||
bus: EventBus,
|
||||
parent_session: &SessionId,
|
||||
max_reusable_per_agent: u32,
|
||||
) -> Result<Self, StoreError> {
|
||||
let existing = store.jobs_for_parent(parent_session.clone()).await?;
|
||||
let mut jobs = HashMap::new();
|
||||
for job in existing {
|
||||
jobs.insert(job.task_id.clone(), job);
|
||||
}
|
||||
Ok(Self {
|
||||
store,
|
||||
bus,
|
||||
jobs: RwLock::new(jobs),
|
||||
max_reusable_per_agent,
|
||||
})
|
||||
}
|
||||
|
||||
/// Assigns the next alias for `agent` under this board: `<3-char-prefix>-<n>`.
|
||||
fn next_alias(&self, agent: &str) -> String {
|
||||
let prefix: String = agent.chars().take(3).collect();
|
||||
let prefix = if prefix.is_empty() {
|
||||
"job".to_string()
|
||||
} else {
|
||||
prefix.to_ascii_lowercase()
|
||||
};
|
||||
let n = self
|
||||
.jobs
|
||||
.read()
|
||||
.unwrap()
|
||||
.values()
|
||||
.filter(|j| j.agent == agent)
|
||||
.count()
|
||||
+ 1;
|
||||
format!("{prefix}-{n}")
|
||||
}
|
||||
|
||||
/// Registers a freshly launched task and returns its assigned alias.
|
||||
pub async fn register_launch(&self, spec: LaunchSpec, now: i64) -> Result<String, StoreError> {
|
||||
let alias = self.next_alias(&spec.agent);
|
||||
let record = JobRecord {
|
||||
task_id: spec.task_id,
|
||||
alias: alias.clone(),
|
||||
parent_session: spec.parent_session,
|
||||
child_session: spec.child_session,
|
||||
agent: spec.agent,
|
||||
description: spec.description,
|
||||
objective: spec.objective,
|
||||
state: JobState::Running,
|
||||
reconciled: false,
|
||||
result_summary: None,
|
||||
context_files: Vec::new(),
|
||||
launched_at: now,
|
||||
updated_at: now,
|
||||
last_used_at: now,
|
||||
};
|
||||
self.upsert(record).await?;
|
||||
Ok(alias)
|
||||
}
|
||||
|
||||
/// Marks a job terminal with an optional result summary (truncated).
|
||||
pub async fn finish(
|
||||
&self,
|
||||
task_id: &str,
|
||||
state: JobState,
|
||||
result_summary: Option<String>,
|
||||
now: i64,
|
||||
) -> Result<(), StoreError> {
|
||||
let Some(mut record) = self.get(task_id) else {
|
||||
return Ok(());
|
||||
};
|
||||
record.state = state;
|
||||
record.result_summary = result_summary.map(|s| truncate_summary(&s));
|
||||
record.updated_at = now;
|
||||
record.last_used_at = now;
|
||||
self.upsert(record).await?;
|
||||
if state == JobState::Completed {
|
||||
self.trim_reusable(now).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Records a file a child session read (deduping by path, keeping the largest read).
|
||||
pub async fn report_context_file(
|
||||
&self,
|
||||
task_id: &str,
|
||||
path: String,
|
||||
lines: u32,
|
||||
now: i64,
|
||||
) -> Result<(), StoreError> {
|
||||
if lines < MIN_REPORTED_LINES {
|
||||
return Ok(());
|
||||
}
|
||||
let Some(mut record) = self.get(task_id) else {
|
||||
return Ok(());
|
||||
};
|
||||
match record.context_files.iter_mut().find(|f| f.path == path) {
|
||||
Some(existing) => existing.lines = existing.lines.max(lines),
|
||||
None => record.context_files.push(ContextFile { path, lines }),
|
||||
}
|
||||
record.updated_at = now;
|
||||
self.upsert(record).await
|
||||
}
|
||||
|
||||
/// Resolves an alias or task id to a job for the given parent (reuse lookup). Only
|
||||
/// completed (reusable) jobs match.
|
||||
pub fn resolve_reusable(&self, parent: &SessionId, alias_or_id: &str) -> Option<JobRecord> {
|
||||
self.jobs
|
||||
.read()
|
||||
.unwrap()
|
||||
.values()
|
||||
.find(|j| {
|
||||
&j.parent_session == parent
|
||||
&& j.state == JobState::Completed
|
||||
&& (j.alias == alias_or_id || j.task_id == alias_or_id)
|
||||
})
|
||||
.cloned()
|
||||
}
|
||||
|
||||
/// Bumps `last_used_at` when a completed session is reused, keeping it fresh in the LRU.
|
||||
pub async fn touch(&self, task_id: &str, now: i64) -> Result<(), StoreError> {
|
||||
let Some(mut record) = self.get(task_id) else {
|
||||
return Ok(());
|
||||
};
|
||||
record.last_used_at = now;
|
||||
self.upsert(record).await
|
||||
}
|
||||
|
||||
/// Marks all terminal jobs `reconciled` — the orchestrator has now seen them.
|
||||
pub async fn reconcile_terminal(&self, now: i64) -> Result<(), StoreError> {
|
||||
let to_update: Vec<JobRecord> = {
|
||||
let jobs = self.jobs.read().unwrap();
|
||||
jobs.values()
|
||||
.filter(|j| j.state.is_terminal() && !j.reconciled)
|
||||
.cloned()
|
||||
.collect()
|
||||
};
|
||||
for mut record in to_update {
|
||||
record.reconciled = true;
|
||||
record.updated_at = now;
|
||||
self.upsert(record).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.jobs.read().unwrap().is_empty()
|
||||
}
|
||||
|
||||
pub fn snapshot(&self) -> Vec<JobRecord> {
|
||||
self.jobs.read().unwrap().values().cloned().collect()
|
||||
}
|
||||
|
||||
/// Renders the board as a synthetic prompt block, or `None` when there is nothing to show.
|
||||
/// Mirrors slim's `formatForPrompt`.
|
||||
pub fn format_for_prompt(&self) -> Option<String> {
|
||||
let jobs = self.jobs.read().unwrap();
|
||||
if jobs.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut active: Vec<&JobRecord> = jobs
|
||||
.values()
|
||||
.filter(|j| !j.state.is_terminal() || !j.reconciled)
|
||||
.collect();
|
||||
let mut reusable: Vec<&JobRecord> = jobs
|
||||
.values()
|
||||
.filter(|j| j.state == JobState::Completed && j.reconciled)
|
||||
.collect();
|
||||
active.sort_by(|a, b| a.alias.cmp(&b.alias));
|
||||
reusable.sort_by(|a, b| a.alias.cmp(&b.alias));
|
||||
|
||||
if active.is_empty() && reusable.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut out = String::from(
|
||||
"### Background Job Board\n\
|
||||
Do not poll running jobs; wait for completion. Reconcile terminal jobs before your \
|
||||
final response.\nCompleted sessions are reusable by alias for the same specialist.\n",
|
||||
);
|
||||
|
||||
if !active.is_empty() {
|
||||
out.push_str("\n#### Active / Unreconciled\n");
|
||||
for job in &active {
|
||||
let state = state_label(job.state);
|
||||
out.push_str(&format!(
|
||||
"- {} / {} / {} / {state}",
|
||||
job.alias, job.child_session, job.agent
|
||||
));
|
||||
if let Some(obj) = &job.objective {
|
||||
out.push_str(&format!(" — Objective: {obj}"));
|
||||
}
|
||||
out.push('\n');
|
||||
if let Some(summary) = &job.result_summary {
|
||||
out.push_str(&format!(" Result: {summary}\n"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !reusable.is_empty() {
|
||||
out.push_str("\n#### Reusable Sessions\n");
|
||||
for job in &reusable {
|
||||
out.push_str(&format!(
|
||||
"- {} / {} / {} / completed\n",
|
||||
job.alias, job.child_session, job.agent
|
||||
));
|
||||
if let Some(obj) = &job.objective {
|
||||
out.push_str(&format!(" Objective: {obj}\n"));
|
||||
}
|
||||
if !job.context_files.is_empty() {
|
||||
let files: Vec<&str> = job
|
||||
.context_files
|
||||
.iter()
|
||||
.take(CONTEXT_FILES_SHOWN)
|
||||
.map(|f| f.path.as_str())
|
||||
.collect();
|
||||
out.push_str(&format!(" Context read: {}\n", files.join(", ")));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some(out)
|
||||
}
|
||||
|
||||
fn get(&self, task_id: &str) -> Option<JobRecord> {
|
||||
self.jobs.read().unwrap().get(task_id).cloned()
|
||||
}
|
||||
|
||||
async fn upsert(&self, record: JobRecord) -> Result<(), StoreError> {
|
||||
self.store.upsert_job(record.clone()).await?;
|
||||
self.jobs
|
||||
.write()
|
||||
.unwrap()
|
||||
.insert(record.task_id.clone(), record.clone());
|
||||
self.bus.publish(AppEvent::JobUpdated {
|
||||
job: JobRecordEvent(serde_json::to_value(&record).unwrap_or(serde_json::Value::Null)),
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Keeps at most `max_reusable_per_agent` completed jobs per agent (LRU by `last_used_at`);
|
||||
/// older completed jobs are dropped from the board and the store.
|
||||
async fn trim_reusable(&self, _now: i64) -> Result<(), StoreError> {
|
||||
let to_remove: Vec<String> = {
|
||||
let jobs = self.jobs.read().unwrap();
|
||||
let mut by_agent: HashMap<&str, Vec<&JobRecord>> = HashMap::new();
|
||||
for job in jobs.values().filter(|j| j.state == JobState::Completed) {
|
||||
by_agent.entry(job.agent.as_str()).or_default().push(job);
|
||||
}
|
||||
let mut remove = Vec::new();
|
||||
for group in by_agent.values_mut() {
|
||||
if group.len() as u32 <= self.max_reusable_per_agent {
|
||||
continue;
|
||||
}
|
||||
// Oldest last_used_at first; drop the excess from the front.
|
||||
group.sort_by_key(|j| j.last_used_at);
|
||||
let excess = group.len() - self.max_reusable_per_agent as usize;
|
||||
for job in group.iter().take(excess) {
|
||||
remove.push(job.task_id.clone());
|
||||
}
|
||||
}
|
||||
remove
|
||||
};
|
||||
for task_id in to_remove {
|
||||
self.store.delete_job(task_id.clone()).await?;
|
||||
self.jobs.write().unwrap().remove(&task_id);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn state_label(state: JobState) -> &'static str {
|
||||
match state {
|
||||
JobState::Running => "running",
|
||||
JobState::Completed => "completed",
|
||||
JobState::Error => "error",
|
||||
JobState::Cancelled => "cancelled",
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_summary(s: &str) -> String {
|
||||
if s.len() <= SUMMARY_MAX {
|
||||
return s.to_string();
|
||||
}
|
||||
let mut end = SUMMARY_MAX;
|
||||
while !s.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
format!("{}…", &s[..end])
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
async fn board(max_reusable: u32) -> (Store, JobBoard, SessionId) {
|
||||
let store = Store::open_in_memory().unwrap();
|
||||
let bus = EventBus::new();
|
||||
let parent = SessionId::new();
|
||||
let board = JobBoard::load(store.clone(), bus, &parent, max_reusable)
|
||||
.await
|
||||
.unwrap();
|
||||
(store, board, parent)
|
||||
}
|
||||
|
||||
fn spec(
|
||||
task_id: &str,
|
||||
parent: SessionId,
|
||||
child: SessionId,
|
||||
agent: &str,
|
||||
objective: Option<&str>,
|
||||
) -> LaunchSpec {
|
||||
LaunchSpec {
|
||||
task_id: task_id.into(),
|
||||
parent_session: parent,
|
||||
child_session: child,
|
||||
agent: agent.into(),
|
||||
description: "d".into(),
|
||||
objective: objective.map(Into::into),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn alias_increments_per_agent() {
|
||||
let (_store, board, parent) = board(2).await;
|
||||
let a1 = board
|
||||
.register_launch(
|
||||
spec("t1", parent.clone(), SessionId::new(), "explorer", None),
|
||||
1,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let a2 = board
|
||||
.register_launch(
|
||||
spec("t2", parent.clone(), SessionId::new(), "explorer", None),
|
||||
2,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let f1 = board
|
||||
.register_launch(
|
||||
spec("t3", parent.clone(), SessionId::new(), "fixer", None),
|
||||
3,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(a1, "exp-1");
|
||||
assert_eq!(a2, "exp-2");
|
||||
assert_eq!(f1, "fix-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn finish_makes_job_reusable_and_resolvable_by_alias() {
|
||||
let (_store, board, parent) = board(2).await;
|
||||
let child = SessionId::new();
|
||||
let alias = board
|
||||
.register_launch(
|
||||
spec(
|
||||
"t1",
|
||||
parent.clone(),
|
||||
child.clone(),
|
||||
"explorer",
|
||||
Some("map the auth flow"),
|
||||
),
|
||||
1,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
// Running jobs are not reusable.
|
||||
assert!(board.resolve_reusable(&parent, &alias).is_none());
|
||||
board
|
||||
.finish("t1", JobState::Completed, Some("done".into()), 2)
|
||||
.await
|
||||
.unwrap();
|
||||
let resolved = board.resolve_reusable(&parent, &alias).unwrap();
|
||||
assert_eq!(resolved.child_session, child);
|
||||
assert_eq!(resolved.result_summary.as_deref(), Some("done"));
|
||||
// Also resolvable by task id.
|
||||
assert!(board.resolve_reusable(&parent, "t1").is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn context_files_dedupe_and_respect_min_lines() {
|
||||
let (_store, board, parent) = board(2).await;
|
||||
board
|
||||
.register_launch(spec("t1", parent, SessionId::new(), "explorer", None), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
// Below threshold — ignored.
|
||||
board
|
||||
.report_context_file("t1", "small.rs".into(), 3, 2)
|
||||
.await
|
||||
.unwrap();
|
||||
board
|
||||
.report_context_file("t1", "a.rs".into(), 20, 2)
|
||||
.await
|
||||
.unwrap();
|
||||
// Same file again with a larger read keeps the max.
|
||||
board
|
||||
.report_context_file("t1", "a.rs".into(), 50, 3)
|
||||
.await
|
||||
.unwrap();
|
||||
let job = board.snapshot().into_iter().next().unwrap();
|
||||
assert_eq!(job.context_files.len(), 1);
|
||||
assert_eq!(job.context_files[0].path, "a.rs");
|
||||
assert_eq!(job.context_files[0].lines, 50);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn trim_reusable_keeps_lru_within_limit() {
|
||||
let (_store, board, parent) = board(2).await;
|
||||
for (i, ts) in [(1, 10), (2, 20), (3, 30)] {
|
||||
board
|
||||
.register_launch(
|
||||
spec(
|
||||
&format!("t{i}"),
|
||||
parent.clone(),
|
||||
SessionId::new(),
|
||||
"explorer",
|
||||
None,
|
||||
),
|
||||
ts,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
board
|
||||
.finish(&format!("t{i}"), JobState::Completed, None, ts)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
// max_reusable = 2, so the oldest (t1, last_used 10) is dropped.
|
||||
let ids: Vec<String> = board.snapshot().into_iter().map(|j| j.task_id).collect();
|
||||
assert_eq!(ids.len(), 2);
|
||||
assert!(!ids.contains(&"t1".to_string()));
|
||||
assert!(ids.contains(&"t2".to_string()));
|
||||
assert!(ids.contains(&"t3".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reconcile_flips_terminal_jobs_and_moves_them_to_reusable_section() {
|
||||
let (_store, board, parent) = board(2).await;
|
||||
board
|
||||
.register_launch(
|
||||
spec("t1", parent, SessionId::new(), "explorer", Some("obj")),
|
||||
1,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
board
|
||||
.finish("t1", JobState::Completed, Some("res".into()), 2)
|
||||
.await
|
||||
.unwrap();
|
||||
// Before reconcile: appears under Active/Unreconciled.
|
||||
let prompt = board.format_for_prompt().unwrap();
|
||||
assert!(prompt.contains("Active / Unreconciled"));
|
||||
board.reconcile_terminal(3).await.unwrap();
|
||||
let prompt = board.format_for_prompt().unwrap();
|
||||
assert!(prompt.contains("Reusable Sessions"));
|
||||
assert!(prompt.contains("exp-1"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn board_reloads_persisted_jobs() {
|
||||
let store = Store::open_in_memory().unwrap();
|
||||
let bus = EventBus::new();
|
||||
let parent = SessionId::new();
|
||||
{
|
||||
let board = JobBoard::load(store.clone(), bus.clone(), &parent, 2)
|
||||
.await
|
||||
.unwrap();
|
||||
board
|
||||
.register_launch(
|
||||
spec("t1", parent.clone(), SessionId::new(), "explorer", None),
|
||||
1,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
// Fresh board over the same store sees the persisted job.
|
||||
let board = JobBoard::load(store, bus, &parent, 2).await.unwrap();
|
||||
assert!(!board.is_empty());
|
||||
assert_eq!(board.snapshot().len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_board_formats_to_none() {
|
||||
let (_store, board, _parent) = board(2).await;
|
||||
assert!(board.format_for_prompt().is_none());
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
pub mod compact;
|
||||
pub mod doomloop;
|
||||
pub mod jobs;
|
||||
pub mod processor;
|
||||
pub mod retry;
|
||||
#[path = "loop.rs"]
|
||||
pub mod session_loop;
|
||||
pub mod system;
|
||||
|
||||
pub use compact::Compactor;
|
||||
pub use doomloop::DoomLoopGuard;
|
||||
pub use jobs::{ContextFile, JobBoard, JobRecord, JobState};
|
||||
pub use processor::{process_step, StepContext, StepError, StepOutcome, StepResult};
|
||||
pub use session_loop::{run_session, RunConfig};
|
||||
@@ -0,0 +1,702 @@
|
||||
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::{
|
||||
ContextReporter, MetadataSink, PermissionHandle, SubagentSpawner, Tool, ToolCtx, ToolError,
|
||||
ToolRegistry,
|
||||
};
|
||||
use crate::types::{Message, MessageId, Part, PartBody, PartId, SessionId, TokenUsage, ToolState};
|
||||
|
||||
use super::doomloop::DoomLoopGuard;
|
||||
use super::jobs::JobBoard;
|
||||
|
||||
#[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,
|
||||
/// Dollar cost of this step's usage (0.0 when no pricing is available).
|
||||
pub cost: f64,
|
||||
pub aborted: bool,
|
||||
/// Whether a file-mutating tool (`edit`/`write`) ran this step — drives the optional
|
||||
/// `after_file_tool` reminder injection on the next turn.
|
||||
pub used_file_tool: 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>>,
|
||||
/// Parent-effective ruleset for a subagent session; empty for a root session. Enables
|
||||
/// permission intersection on this session's tool calls.
|
||||
pub parent_rules: 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,
|
||||
/// Lets the `task` tool spawn subagents. `None` disables delegation (headless/tests).
|
||||
pub spawner: Option<Arc<dyn SubagentSpawner>>,
|
||||
/// This session's background job board (as a parent). Injected into requests when the
|
||||
/// running agent can delegate; `None` disables the board.
|
||||
pub job_board: Option<Arc<JobBoard>>,
|
||||
/// Present in subagent sessions: reports read files to this session's job on the parent
|
||||
/// board. `None` for root sessions (nothing to report to).
|
||||
pub context_reporter: Option<Arc<dyn ContextReporter>>,
|
||||
/// Language-server diagnostics source shared by the session's edit/write tool calls.
|
||||
pub diagnostics: Option<Arc<dyn crate::lsp::DiagnosticsSource>>,
|
||||
/// Summarizes history when the context nears the model's window. `None` disables
|
||||
/// auto-compaction (headless/tests, or when no `small_model` is configured).
|
||||
pub compactor: Option<Arc<dyn super::compact::Compactor>>,
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Tool names that mutate files — after one runs, the optional `after_file_tool` reminder fires.
|
||||
const FILE_TOOLS: &[&str] = &["edit", "write"];
|
||||
|
||||
struct Run<'a> {
|
||||
ctx: &'a StepContext,
|
||||
assistant: Option<Message>,
|
||||
used_file_tool: bool,
|
||||
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,
|
||||
used_file_tool: false,
|
||||
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> {
|
||||
if FILE_TOOLS.contains(&name.as_str()) {
|
||||
self.used_file_tool = true;
|
||||
}
|
||||
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(),
|
||||
)
|
||||
.with_parent_rules(self.ctx.parent_rules.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,
|
||||
spawner: self.ctx.spawner.clone(),
|
||||
context_reporter: self.ctx.context_reporter.clone(),
|
||||
diagnostics: self.ctx.diagnostics.clone(),
|
||||
};
|
||||
|
||||
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,
|
||||
cost: f64,
|
||||
) -> 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,
|
||||
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,
|
||||
cost: Option<crate::types::ModelCost>,
|
||||
doomloop: &mut DoomLoopGuard,
|
||||
) -> Result<StepOutcome, StepError> {
|
||||
let mut run = Run::new(ctx);
|
||||
let mut usage = TokenUsage::default();
|
||||
let mut step_cost = 0.0;
|
||||
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,
|
||||
cost: step_cost,
|
||||
aborted: true,
|
||||
used_file_tool: run.used_file_tool,
|
||||
});
|
||||
}
|
||||
};
|
||||
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;
|
||||
step_cost = cost.map(|c| c.cost_of(&finish_usage)).unwrap_or(0.0);
|
||||
match run.on_finish(reason, finish_usage, step_cost).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,
|
||||
cost: step_cost,
|
||||
aborted: false,
|
||||
used_file_tool: run.used_file_tool,
|
||||
})
|
||||
}
|
||||
@@ -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,76 @@
|
||||
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, the optional skills
|
||||
/// listing, then project instructions (e.g. AGENTS.md contents). Order matches `02-engine.md`.
|
||||
pub fn assemble(
|
||||
env_header: String,
|
||||
agent_prompt: &str,
|
||||
skills: Option<&str>,
|
||||
instructions: &[String],
|
||||
) -> Vec<String> {
|
||||
let mut blocks = vec![env_header, agent_prompt.to_string()];
|
||||
if let Some(skills) = skills {
|
||||
blocks.push(skills.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",
|
||||
None,
|
||||
&["AGENTS.md contents".to_string()],
|
||||
);
|
||||
assert_eq!(blocks, vec!["ENV", "AGENT", "AGENTS.md contents"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assemble_inserts_skills_after_agent_before_instructions() {
|
||||
let blocks = assemble(
|
||||
"ENV".to_string(),
|
||||
"AGENT",
|
||||
Some("SKILLS"),
|
||||
&["INSTR".to_string()],
|
||||
);
|
||||
assert_eq!(blocks, vec!["ENV", "AGENT", "SKILLS", "INSTR"]);
|
||||
}
|
||||
}
|
||||
@@ -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,12 @@
|
||||
pub mod agent;
|
||||
pub mod config;
|
||||
pub mod engine;
|
||||
pub mod event;
|
||||
pub mod llm;
|
||||
pub mod lsp;
|
||||
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>;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
//! Seam trait for language-server diagnostics, implemented by `harness-lsp` and consumed by
|
||||
//! the edit/write tools. Kept in core so `harness-tools` never links the LSP crate directly.
|
||||
//! See `docs/09-integrations.md`.
|
||||
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Severity {
|
||||
Error,
|
||||
Warning,
|
||||
Info,
|
||||
Hint,
|
||||
}
|
||||
|
||||
/// One diagnostic reported by a language server. Line/character are 1-based for display.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct Diagnostic {
|
||||
pub line: u32,
|
||||
pub character: u32,
|
||||
pub severity: Severity,
|
||||
pub message: String,
|
||||
pub source: Option<String>,
|
||||
}
|
||||
|
||||
impl Diagnostic {
|
||||
/// `{file}: L{line}: {message}` — the one-line form appended to tool output.
|
||||
pub fn display_line(&self, file: &str) -> String {
|
||||
format!("{file}: L{}: {}", self.line, self.message)
|
||||
}
|
||||
}
|
||||
|
||||
/// A source of file diagnostics (an LSP pool). All methods are best-effort: failures are
|
||||
/// swallowed (logged) so a broken language server never fails a tool call.
|
||||
#[async_trait]
|
||||
pub trait DiagnosticsSource: Send + Sync {
|
||||
/// Ensure a server for `path`'s language is running and told about the file's current
|
||||
/// contents (spawn-if-needed + didOpen/didChange).
|
||||
async fn touch(&self, path: &Path);
|
||||
|
||||
/// Diagnostics for `path`, waiting up to `wait` for the server to (re)publish after a
|
||||
/// change. Returns whatever is known on timeout.
|
||||
async fn diagnostics(&self, path: &Path, wait: Duration) -> Vec<Diagnostic>;
|
||||
}
|
||||
@@ -1,3 +1,7 @@
|
||||
pub mod rule;
|
||||
pub mod service;
|
||||
|
||||
pub use rule::{evaluate, Action, Rule, Ruleset};
|
||||
pub use rule::{evaluate, evaluate_intersected, Action, Rule, Ruleset};
|
||||
pub use service::{
|
||||
spawn_auto_approve, AskDecision, AskError, AskInput, PermissionReply, PermissionService,
|
||||
};
|
||||
|
||||
@@ -44,6 +44,37 @@ pub fn evaluate(stack: &[&Ruleset], permission: &str, pattern: &str) -> Action {
|
||||
result
|
||||
}
|
||||
|
||||
impl Action {
|
||||
/// How restrictive this verdict is: `Deny` > `Ask` > `Allow`. Used to intersect a
|
||||
/// parent and child verdict when a subagent runs (docs/04-multiagent.md).
|
||||
fn restrictiveness(self) -> u8 {
|
||||
match self {
|
||||
Action::Allow => 0,
|
||||
Action::Ask => 1,
|
||||
Action::Deny => 2,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Evaluate `permission`/`pattern` against a parent-effective stack and a child stack
|
||||
/// independently, returning the **more restrictive** of the two verdicts
|
||||
/// (`deny > ask > allow`). A subagent's tool call must satisfy both the rules it inherits
|
||||
/// from its spawning chain and its own agent ruleset.
|
||||
pub fn evaluate_intersected(
|
||||
parent_stack: &[&Ruleset],
|
||||
child_stack: &[&Ruleset],
|
||||
permission: &str,
|
||||
pattern: &str,
|
||||
) -> Action {
|
||||
let parent = evaluate(parent_stack, permission, pattern);
|
||||
let child = evaluate(child_stack, permission, pattern);
|
||||
if child.restrictiveness() >= parent.restrictiveness() {
|
||||
child
|
||||
} else {
|
||||
parent
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -112,4 +143,49 @@ mod tests {
|
||||
fn empty_stack_defaults_to_ask() {
|
||||
assert_eq!(evaluate(&[], "bash", "ls"), Action::Ask);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn intersected_takes_the_more_restrictive_verdict() {
|
||||
let parent_allow: Ruleset = vec![rule("edit", "*", Action::Allow)];
|
||||
let child_deny: Ruleset = vec![rule("edit", "*", Action::Deny)];
|
||||
// Child denies what the parent would allow → deny wins.
|
||||
assert_eq!(
|
||||
evaluate_intersected(&[&parent_allow], &[&child_deny], "edit", "main.rs"),
|
||||
Action::Deny
|
||||
);
|
||||
// Symmetric: parent denies what the child would allow → deny still wins.
|
||||
assert_eq!(
|
||||
evaluate_intersected(&[&child_deny], &[&parent_allow], "edit", "main.rs"),
|
||||
Action::Deny
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn intersected_ask_beats_allow_but_loses_to_deny() {
|
||||
let allow: Ruleset = vec![rule("bash", "*", Action::Allow)];
|
||||
let ask: Ruleset = vec![rule("bash", "*", Action::Ask)];
|
||||
let deny: Ruleset = vec![rule("bash", "*", Action::Deny)];
|
||||
assert_eq!(
|
||||
evaluate_intersected(&[&allow], &[&ask], "bash", "ls"),
|
||||
Action::Ask
|
||||
);
|
||||
assert_eq!(
|
||||
evaluate_intersected(&[&ask], &[&deny], "bash", "ls"),
|
||||
Action::Deny
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn intersected_allows_only_when_both_allow() {
|
||||
let allow: Ruleset = vec![rule("read", "*", Action::Allow)];
|
||||
assert_eq!(
|
||||
evaluate_intersected(&[&allow], &[&allow], "read", "src/a.rs"),
|
||||
Action::Allow
|
||||
);
|
||||
// Empty child stack defaults to Ask, which is more restrictive than parent Allow.
|
||||
assert_eq!(
|
||||
evaluate_intersected(&[&allow], &[], "read", "src/a.rs"),
|
||||
Action::Ask
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,377 @@
|
||||
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, evaluate_intersected, 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,
|
||||
}
|
||||
}
|
||||
|
||||
/// Like [`ask`](Self::ask), but for a subagent: the verdict is the more restrictive of
|
||||
/// the `parent_stack` (rules inherited from the spawning chain) and `child_stack` (the
|
||||
/// subagent's own rules). Used so a child can never widen what its parent forbids.
|
||||
pub async fn ask_intersected(
|
||||
&self,
|
||||
session_id: &SessionId,
|
||||
parent_stack: &[&Ruleset],
|
||||
child_stack: &[&Ruleset],
|
||||
input: AskInput,
|
||||
cancel: &CancellationToken,
|
||||
) -> Result<AskDecision, AskError> {
|
||||
match evaluate_intersected(parent_stack, child_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));
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,7 @@
|
||||
use rusqlite::{params, Connection};
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
use crate::engine::jobs::JobRecord;
|
||||
use crate::types::{Message, MessageId, Part, Session, SessionId};
|
||||
|
||||
use super::api::StoreError;
|
||||
@@ -13,9 +14,13 @@ pub enum StoreCmd {
|
||||
UpsertSession(Session, Reply<()>),
|
||||
UpsertMessage(Message, Reply<()>),
|
||||
UpsertPart(Part, Reply<()>),
|
||||
Session(SessionId, Reply<Option<Session>>),
|
||||
Sessions(Reply<Vec<Session>>),
|
||||
Messages(SessionId, Reply<Vec<Message>>),
|
||||
Parts(MessageId, Reply<Vec<Part>>),
|
||||
UpsertJob(JobRecord, Reply<()>),
|
||||
DeleteJob(String, Reply<()>),
|
||||
JobsForParent(SessionId, Reply<Vec<JobRecord>>),
|
||||
}
|
||||
|
||||
fn init_schema(conn: &Connection) -> rusqlite::Result<()> {
|
||||
@@ -69,6 +74,15 @@ fn upsert_part(conn: &Connection, part: &Part) -> Result<(), StoreError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn get_session(conn: &Connection, id: &SessionId) -> Result<Option<Session>, StoreError> {
|
||||
let mut stmt = conn.prepare("SELECT data FROM session WHERE id = ?1")?;
|
||||
let mut rows = stmt.query_map(params![id.as_ref()], |row| row.get::<_, String>(0))?;
|
||||
match rows.next() {
|
||||
Some(data) => Ok(Some(serde_json::from_str(&data?)?)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn list_sessions(conn: &Connection) -> Result<Vec<Session>, StoreError> {
|
||||
let mut stmt = conn.prepare("SELECT data FROM session ORDER BY id")?;
|
||||
let rows = stmt
|
||||
@@ -99,6 +113,34 @@ fn list_parts(conn: &Connection, message_id: &MessageId) -> Result<Vec<Part>, St
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn upsert_job(conn: &Connection, job: &JobRecord) -> Result<(), StoreError> {
|
||||
let data = serde_json::to_string(job)?;
|
||||
conn.execute(
|
||||
"INSERT INTO job (task_id, parent_session_id, data) VALUES (?1, ?2, ?3)
|
||||
ON CONFLICT(task_id) DO UPDATE SET data = ?3",
|
||||
params![job.task_id, job.parent_session.as_ref(), data],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn delete_job(conn: &Connection, task_id: &str) -> Result<(), StoreError> {
|
||||
conn.execute("DELETE FROM job WHERE task_id = ?1", params![task_id])?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn list_jobs_for_parent(
|
||||
conn: &Connection,
|
||||
parent: &SessionId,
|
||||
) -> Result<Vec<JobRecord>, StoreError> {
|
||||
let mut stmt = conn.prepare("SELECT data FROM job WHERE parent_session_id = ?1")?;
|
||||
let rows = stmt
|
||||
.query_map(params![parent.as_ref()], |row| row.get::<_, String>(0))?
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
rows.iter()
|
||||
.map(|data| serde_json::from_str(data).map_err(StoreError::from))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Runs on a dedicated OS thread; the async facade in `api.rs` talks to it over `mpsc`.
|
||||
pub fn run(conn: Connection, mut rx: tokio::sync::mpsc::Receiver<StoreCmd>) {
|
||||
if let Err(e) = init_schema(&conn) {
|
||||
@@ -116,6 +158,9 @@ pub fn run(conn: Connection, mut rx: tokio::sync::mpsc::Receiver<StoreCmd>) {
|
||||
StoreCmd::UpsertPart(part, reply) => {
|
||||
let _ = reply.send(upsert_part(&conn, &part));
|
||||
}
|
||||
StoreCmd::Session(id, reply) => {
|
||||
let _ = reply.send(get_session(&conn, &id));
|
||||
}
|
||||
StoreCmd::Sessions(reply) => {
|
||||
let _ = reply.send(list_sessions(&conn));
|
||||
}
|
||||
@@ -125,6 +170,15 @@ pub fn run(conn: Connection, mut rx: tokio::sync::mpsc::Receiver<StoreCmd>) {
|
||||
StoreCmd::Parts(message_id, reply) => {
|
||||
let _ = reply.send(list_parts(&conn, &message_id));
|
||||
}
|
||||
StoreCmd::UpsertJob(job, reply) => {
|
||||
let _ = reply.send(upsert_job(&conn, &job));
|
||||
}
|
||||
StoreCmd::DeleteJob(task_id, reply) => {
|
||||
let _ = reply.send(delete_job(&conn, &task_id));
|
||||
}
|
||||
StoreCmd::JobsForParent(parent, reply) => {
|
||||
let _ = reply.send(list_jobs_for_parent(&conn, &parent));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ use std::path::Path;
|
||||
use rusqlite::Connection;
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
use crate::engine::jobs::JobRecord;
|
||||
use crate::types::{Message, MessageId, Part, Session, SessionId};
|
||||
|
||||
use super::actor::{self, StoreCmd};
|
||||
@@ -77,6 +78,11 @@ impl Store {
|
||||
self.call(|reply| StoreCmd::UpsertPart(part, reply)).await
|
||||
}
|
||||
|
||||
pub async fn session(&self, session_id: SessionId) -> Result<Option<Session>, StoreError> {
|
||||
self.call(|reply| StoreCmd::Session(session_id, reply))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn sessions(&self) -> Result<Vec<Session>, StoreError> {
|
||||
self.call(StoreCmd::Sessions).await
|
||||
}
|
||||
@@ -89,6 +95,19 @@ impl Store {
|
||||
pub async fn parts(&self, message_id: MessageId) -> Result<Vec<Part>, StoreError> {
|
||||
self.call(|reply| StoreCmd::Parts(message_id, reply)).await
|
||||
}
|
||||
|
||||
pub async fn upsert_job(&self, job: JobRecord) -> Result<(), StoreError> {
|
||||
self.call(|reply| StoreCmd::UpsertJob(job, reply)).await
|
||||
}
|
||||
|
||||
pub async fn delete_job(&self, task_id: String) -> Result<(), StoreError> {
|
||||
self.call(|reply| StoreCmd::DeleteJob(task_id, reply)).await
|
||||
}
|
||||
|
||||
pub async fn jobs_for_parent(&self, parent: SessionId) -> Result<Vec<JobRecord>, StoreError> {
|
||||
self.call(|reply| StoreCmd::JobsForParent(parent, reply))
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -0,0 +1,374 @@
|
||||
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};
|
||||
|
||||
/// A request from the `task` tool to run a subagent. The spawner (owned by the composition
|
||||
/// root) resolves the agent, enforces the depth limit, applies permission intersection, and
|
||||
/// runs the child session foreground or background. See `docs/04-multiagent.md`.
|
||||
pub struct SpawnRequest {
|
||||
pub parent_session_id: SessionId,
|
||||
pub parent_message_id: MessageId,
|
||||
pub agent: String,
|
||||
pub description: String,
|
||||
pub prompt: String,
|
||||
/// Alias or task id of a completed job to reuse (continue its child session).
|
||||
pub reuse_task_id: Option<String>,
|
||||
pub background: bool,
|
||||
/// The tool call's cancellation token — used for foreground child runs. Background runs
|
||||
/// are childed from the parent session's run token by the spawner instead.
|
||||
pub cancel: CancellationToken,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct SpawnOutcome {
|
||||
pub child_session_id: SessionId,
|
||||
pub background: bool,
|
||||
/// Board alias assigned to a background launch.
|
||||
pub alias: Option<String>,
|
||||
/// Final assistant text of a foreground run.
|
||||
pub final_text: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum SpawnError {
|
||||
#[error("unknown subagent {0:?}, or it is not usable as a subagent")]
|
||||
InvalidAgent(String),
|
||||
#[error("subagent depth limit reached — do this work yourself instead of delegating further")]
|
||||
DepthExceeded,
|
||||
#[error("cannot reuse {0:?}: no completed job with that alias for this session")]
|
||||
ReuseNotFound(String),
|
||||
#[error("{0}")]
|
||||
Other(String),
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait SubagentSpawner: Send + Sync {
|
||||
async fn spawn(&self, req: SpawnRequest) -> Result<SpawnOutcome, SpawnError>;
|
||||
}
|
||||
|
||||
/// Lets a child session report the files it read to its job board entry, so a completed
|
||||
/// specialist advertises what it already looked at (docs/04-multiagent.md). Present only in
|
||||
/// subagent sessions; the spawner wires it to the right board + job.
|
||||
#[async_trait]
|
||||
pub trait ContextReporter: Send + Sync {
|
||||
async fn report_file(&self, path: String, lines: u32);
|
||||
}
|
||||
|
||||
#[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>>,
|
||||
/// Parent-effective ruleset for a subagent session; empty for a root session. When
|
||||
/// non-empty, verdicts are intersected so a child can only ever be *more* restricted.
|
||||
parent_rules: 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,
|
||||
parent_rules: Vec::new(),
|
||||
cancel,
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets the parent-effective ruleset so this handle intersects verdicts (subagent runs).
|
||||
pub fn with_parent_rules(mut self, parent_rules: Ruleset) -> Self {
|
||||
self.parent_rules = parent_rules;
|
||||
self
|
||||
}
|
||||
|
||||
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 child_stack: [&Ruleset; 2] = [&self.static_rules, &extra_snapshot];
|
||||
let input = AskInput {
|
||||
permission: permission.into(),
|
||||
pattern: pattern.into(),
|
||||
always_pattern: always_pattern.into(),
|
||||
metadata,
|
||||
};
|
||||
let decision = if self.parent_rules.is_empty() {
|
||||
self.service
|
||||
.ask(&self.session_id, &child_stack, input, &self.cancel)
|
||||
.await?
|
||||
} else {
|
||||
let parent_stack: [&Ruleset; 1] = [&self.parent_rules];
|
||||
self.service
|
||||
.ask_intersected(
|
||||
&self.session_id,
|
||||
&parent_stack,
|
||||
&child_stack,
|
||||
input,
|
||||
&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,
|
||||
/// Present when the engine can spawn subagents (the `task` tool's capability). `None`
|
||||
/// in headless/test contexts with no orchestration wired in.
|
||||
pub spawner: Option<Arc<dyn SubagentSpawner>>,
|
||||
/// Present in subagent sessions: lets the read tool report files to the job board.
|
||||
pub context_reporter: Option<Arc<dyn ContextReporter>>,
|
||||
/// Language-server diagnostics source (edit/write surface errors after a change). `None`
|
||||
/// disables LSP integration.
|
||||
pub diagnostics: Option<Arc<dyn crate::lsp::DiagnosticsSource>>,
|
||||
}
|
||||
|
||||
#[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,6 @@ pub mod session;
|
||||
|
||||
pub use ids::{MessageId, PartId, SessionId};
|
||||
pub use message::{Message, MessageError, Role};
|
||||
pub use model::{ModelInfo, ModelRef, TokenUsage};
|
||||
pub use model::{ModelCost, ModelInfo, ModelRef, TokenUsage};
|
||||
pub use part::{Part, PartBody, ToolState};
|
||||
pub use session::Session;
|
||||
|
||||
@@ -34,11 +34,71 @@ impl TokenUsage {
|
||||
}
|
||||
}
|
||||
|
||||
// Full cost/context-limit metadata is populated by harness-providers (models.dev) in M1/M3;
|
||||
// this placeholder only carries what harness-core needs to key on.
|
||||
/// Per-model pricing in USD per **one million** tokens. Populated from models.dev metadata
|
||||
/// by `harness-providers`.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ModelCost {
|
||||
pub input: f64,
|
||||
pub output: f64,
|
||||
pub cache_read: f64,
|
||||
pub cache_write: f64,
|
||||
}
|
||||
|
||||
impl ModelCost {
|
||||
/// Dollar cost of a usage sample. Reasoning tokens are billed within `output` by the
|
||||
/// providers we support (OpenAI reports them as a subset of `output_tokens`; Anthropic
|
||||
/// counts thinking in output), so they are intentionally not charged separately here.
|
||||
pub fn cost_of(&self, usage: &TokenUsage) -> f64 {
|
||||
let per_million = |tokens: u64, rate: f64| (tokens as f64) * rate / 1_000_000.0;
|
||||
per_million(usage.input, self.input)
|
||||
+ per_million(usage.output, self.output)
|
||||
+ per_million(usage.cache_read, self.cache_read)
|
||||
+ per_million(usage.cache_write, self.cache_write)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn cost_of_sums_components_per_million_excluding_reasoning() {
|
||||
let cost = ModelCost {
|
||||
input: 3.0,
|
||||
output: 15.0,
|
||||
cache_read: 0.30,
|
||||
cache_write: 3.75,
|
||||
};
|
||||
let usage = TokenUsage {
|
||||
input: 1_000_000,
|
||||
output: 1_000_000,
|
||||
reasoning: 500_000, // billed within output — must not add extra cost
|
||||
cache_read: 1_000_000,
|
||||
cache_write: 1_000_000,
|
||||
};
|
||||
// 3 + 15 + 0.30 + 3.75, with reasoning contributing nothing.
|
||||
assert!((cost.cost_of(&usage) - 22.05).abs() < 1e-9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_cost_is_zero() {
|
||||
assert_eq!(ModelCost::default().cost_of(&TokenUsage::default()), 0.0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Model metadata, keyed on `model`. Cost/limits are populated from models.dev by
|
||||
/// `harness-providers`; `reasoning`/`tool_call`/`attachment` are capability flags.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ModelInfo {
|
||||
pub model: ModelRef,
|
||||
pub context_limit: u64,
|
||||
pub output_limit: u64,
|
||||
#[serde(default)]
|
||||
pub cost: ModelCost,
|
||||
#[serde(default)]
|
||||
pub reasoning: bool,
|
||||
#[serde(default)]
|
||||
pub tool_call: bool,
|
||||
#[serde(default)]
|
||||
pub attachment: bool,
|
||||
}
|
||||
|
||||
@@ -6,6 +6,13 @@ license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
harness-core = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tempfile = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -0,0 +1,325 @@
|
||||
//! Minimal JSON-RPC client over an LSP server's child stdio: `Content-Length` framing, a
|
||||
//! request-id → oneshot map, and a diagnostics store keyed by document URI. Only the handful
|
||||
//! of methods the edit/write flow needs are implemented (docs/09-integrations.md).
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::sync::atomic::{AtomicI64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use harness_core::lsp::{Diagnostic, Severity};
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::process::{Child, Command};
|
||||
use tokio::sync::{oneshot, Notify};
|
||||
|
||||
type Pending = Arc<Mutex<HashMap<i64, oneshot::Sender<Value>>>>;
|
||||
type DiagStore = Arc<Mutex<HashMap<String, Vec<Diagnostic>>>>;
|
||||
|
||||
/// A running language server plus the state needed to talk to it.
|
||||
pub struct LspClient {
|
||||
outgoing: tokio::sync::mpsc::UnboundedSender<String>,
|
||||
next_id: AtomicI64,
|
||||
pending: Pending,
|
||||
diagnostics: DiagStore,
|
||||
diag_notify: Arc<Notify>,
|
||||
// Kept alive so the child is killed on drop (`kill_on_drop`).
|
||||
_child: Child,
|
||||
}
|
||||
|
||||
impl LspClient {
|
||||
/// Spawns `command args`, performs the `initialize`/`initialized` handshake rooted at
|
||||
/// `root`, and returns a ready client.
|
||||
pub async fn spawn(command: &str, args: &[String], root: &Path) -> std::io::Result<Self> {
|
||||
let mut child = Command::new(command)
|
||||
.args(args)
|
||||
.current_dir(root)
|
||||
.stdin(std::process::Stdio::piped())
|
||||
.stdout(std::process::Stdio::piped())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.kill_on_drop(true)
|
||||
.spawn()?;
|
||||
|
||||
let stdin = child.stdin.take().expect("piped stdin");
|
||||
let stdout = child.stdout.take().expect("piped stdout");
|
||||
|
||||
let (outgoing, mut out_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
|
||||
let pending: Pending = Arc::default();
|
||||
let diagnostics: DiagStore = Arc::default();
|
||||
let diag_notify = Arc::new(Notify::new());
|
||||
|
||||
// Writer task: frame and forward outgoing payloads.
|
||||
tokio::spawn(async move {
|
||||
let mut stdin = stdin;
|
||||
while let Some(payload) = out_rx.recv().await {
|
||||
let frame = format!("Content-Length: {}\r\n\r\n{}", payload.len(), payload);
|
||||
if stdin.write_all(frame.as_bytes()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
let _ = stdin.flush().await;
|
||||
}
|
||||
});
|
||||
|
||||
// Reader task: parse frames, route responses, collect diagnostics, ack server requests.
|
||||
{
|
||||
let pending = pending.clone();
|
||||
let diagnostics = diagnostics.clone();
|
||||
let diag_notify = diag_notify.clone();
|
||||
let outgoing_ack = outgoing.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut reader = BufReader::new(stdout);
|
||||
while let Some(msg) = read_message(&mut reader).await {
|
||||
dispatch(msg, &pending, &diagnostics, &diag_notify, &outgoing_ack);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let client = Self {
|
||||
outgoing,
|
||||
next_id: AtomicI64::new(1),
|
||||
pending,
|
||||
diagnostics,
|
||||
diag_notify,
|
||||
_child: child,
|
||||
};
|
||||
|
||||
client.initialize(root).await?;
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
async fn initialize(&self, root: &Path) -> std::io::Result<()> {
|
||||
let root_uri = path_to_uri(root);
|
||||
let params = json!({
|
||||
"processId": std::process::id(),
|
||||
"rootUri": root_uri,
|
||||
"workspaceFolders": [{ "uri": root_uri, "name": "root" }],
|
||||
"capabilities": {
|
||||
"textDocument": {
|
||||
"publishDiagnostics": { "relatedInformation": false }
|
||||
}
|
||||
},
|
||||
"clientInfo": { "name": "ai-harness" }
|
||||
});
|
||||
// A slow server (rust-analyzer indexing) can take a while to answer initialize.
|
||||
let _ = self
|
||||
.request("initialize", params, Duration::from_secs(30))
|
||||
.await;
|
||||
self.notify("initialized", json!({}));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn send(&self, payload: Value) {
|
||||
if let Ok(text) = serde_json::to_string(&payload) {
|
||||
let _ = self.outgoing.send(text);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn notify(&self, method: &str, params: Value) {
|
||||
self.send(json!({ "jsonrpc": "2.0", "method": method, "params": params }));
|
||||
}
|
||||
|
||||
async fn request(&self, method: &str, params: Value, timeout: Duration) -> Option<Value> {
|
||||
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.pending.lock().unwrap().insert(id, tx);
|
||||
self.send(json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }));
|
||||
match tokio::time::timeout(timeout, rx).await {
|
||||
Ok(Ok(value)) => Some(value),
|
||||
_ => {
|
||||
self.pending.lock().unwrap().remove(&id);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn did_open(&self, path: &Path, language_id: &str, version: i32, text: &str) {
|
||||
self.notify(
|
||||
"textDocument/didOpen",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": path_to_uri(path),
|
||||
"languageId": language_id,
|
||||
"version": version,
|
||||
"text": text,
|
||||
}
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
pub fn did_change(&self, path: &Path, version: i32, text: &str) {
|
||||
self.notify(
|
||||
"textDocument/didChange",
|
||||
json!({
|
||||
"textDocument": { "uri": path_to_uri(path), "version": version },
|
||||
"contentChanges": [{ "text": text }],
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
/// Clears the stored diagnostics for `path` so the next `wait_diagnostics` observes a fresh
|
||||
/// publish rather than a stale one.
|
||||
pub fn clear(&self, path: &Path) {
|
||||
self.diagnostics.lock().unwrap().remove(&path_to_uri(path));
|
||||
}
|
||||
|
||||
/// Waits up to `wait` for the server to publish diagnostics for `path`, returning whatever
|
||||
/// is stored on timeout (possibly empty).
|
||||
pub async fn wait_diagnostics(&self, path: &Path, wait: Duration) -> Vec<Diagnostic> {
|
||||
let uri = path_to_uri(path);
|
||||
let deadline = tokio::time::Instant::now() + wait;
|
||||
loop {
|
||||
if let Some(diags) = self.diagnostics.lock().unwrap().get(&uri) {
|
||||
return diags.clone();
|
||||
}
|
||||
let notified = self.diag_notify.notified();
|
||||
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
|
||||
if remaining.is_zero() {
|
||||
break;
|
||||
}
|
||||
if tokio::time::timeout(remaining, notified).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
self.diagnostics
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(&uri)
|
||||
.cloned()
|
||||
.unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch(
|
||||
msg: Value,
|
||||
pending: &Pending,
|
||||
diagnostics: &DiagStore,
|
||||
diag_notify: &Arc<Notify>,
|
||||
outgoing: &tokio::sync::mpsc::UnboundedSender<String>,
|
||||
) {
|
||||
let method = msg.get("method").and_then(|m| m.as_str());
|
||||
let id = msg.get("id");
|
||||
|
||||
match (method, id) {
|
||||
// Server → client request: ack with a null result so the server can proceed
|
||||
// (e.g. client/registerCapability, window/workDoneProgress/create).
|
||||
(Some(_), Some(id)) => {
|
||||
let reply = json!({ "jsonrpc": "2.0", "id": id, "result": Value::Null });
|
||||
if let Ok(text) = serde_json::to_string(&reply) {
|
||||
let _ = outgoing.send(text);
|
||||
}
|
||||
}
|
||||
// Notification from the server.
|
||||
(Some("textDocument/publishDiagnostics"), None) => {
|
||||
if let Some(params) = msg.get("params") {
|
||||
if let Some((uri, diags)) = parse_diagnostics(params) {
|
||||
diagnostics.lock().unwrap().insert(uri, diags);
|
||||
diag_notify.notify_waiters();
|
||||
}
|
||||
}
|
||||
}
|
||||
(Some(_), None) => {}
|
||||
// Response to one of our requests.
|
||||
(None, Some(id)) => {
|
||||
if let Some(id) = id.as_i64() {
|
||||
if let Some(tx) = pending.lock().unwrap().remove(&id) {
|
||||
let result = msg.get("result").cloned().unwrap_or(Value::Null);
|
||||
let _ = tx.send(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
(None, None) => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_diagnostics(params: &Value) -> Option<(String, Vec<Diagnostic>)> {
|
||||
let uri = params.get("uri")?.as_str()?.to_string();
|
||||
let items = params.get("diagnostics")?.as_array()?;
|
||||
let diags = items
|
||||
.iter()
|
||||
.filter_map(|d| {
|
||||
let start = d.get("range")?.get("start")?;
|
||||
Some(Diagnostic {
|
||||
line: start.get("line")?.as_u64().unwrap_or(0) as u32 + 1,
|
||||
character: start.get("character")?.as_u64().unwrap_or(0) as u32 + 1,
|
||||
severity: match d.get("severity").and_then(|s| s.as_u64()) {
|
||||
Some(1) => Severity::Error,
|
||||
Some(2) => Severity::Warning,
|
||||
Some(3) => Severity::Info,
|
||||
_ => Severity::Hint,
|
||||
},
|
||||
message: d.get("message")?.as_str().unwrap_or("").to_string(),
|
||||
source: d
|
||||
.get("source")
|
||||
.and_then(|s| s.as_str())
|
||||
.map(|s| s.to_string()),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
Some((uri, diags))
|
||||
}
|
||||
|
||||
/// Reads one `Content-Length`-framed JSON-RPC message, or `None` at EOF.
|
||||
async fn read_message<R: tokio::io::AsyncBufRead + Unpin>(reader: &mut R) -> Option<Value> {
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
let mut content_length: Option<usize> = None;
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
let n = reader.read_line(&mut line).await.ok()?;
|
||||
if n == 0 {
|
||||
return None; // EOF
|
||||
}
|
||||
let trimmed = line.trim_end();
|
||||
if trimmed.is_empty() {
|
||||
break; // end of headers
|
||||
}
|
||||
if let Some(value) = trimmed.strip_prefix("Content-Length:") {
|
||||
content_length = value.trim().parse().ok();
|
||||
}
|
||||
}
|
||||
let len = content_length?;
|
||||
let mut buf = vec![0u8; len];
|
||||
reader.read_exact(&mut buf).await.ok()?;
|
||||
serde_json::from_slice(&buf).ok()
|
||||
}
|
||||
|
||||
/// `file://` URI for an absolute path (best-effort; assumes UTF-8, no percent-encoding needed
|
||||
/// for the local paths we handle).
|
||||
fn path_to_uri(path: &Path) -> String {
|
||||
let s = path.to_string_lossy();
|
||||
if s.starts_with('/') {
|
||||
format!("file://{s}")
|
||||
} else {
|
||||
format!("file:///{s}")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parses_publish_diagnostics_into_one_based_positions() {
|
||||
let params = json!({
|
||||
"uri": "file:///tmp/a.rs",
|
||||
"diagnostics": [{
|
||||
"range": {"start": {"line": 4, "character": 8}, "end": {"line": 4, "character": 12}},
|
||||
"severity": 1,
|
||||
"message": "cannot find value `x`",
|
||||
"source": "rustc"
|
||||
}]
|
||||
});
|
||||
let (uri, diags) = parse_diagnostics(¶ms).unwrap();
|
||||
assert_eq!(uri, "file:///tmp/a.rs");
|
||||
assert_eq!(diags.len(), 1);
|
||||
assert_eq!(diags[0].line, 5); // 0-based 4 → 1-based 5
|
||||
assert_eq!(diags[0].character, 9);
|
||||
assert_eq!(diags[0].severity, Severity::Error);
|
||||
assert_eq!(diags[0].source.as_deref(), Some("rustc"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn path_to_uri_prefixes_file_scheme() {
|
||||
assert_eq!(path_to_uri(Path::new("/tmp/a.rs")), "file:///tmp/a.rs");
|
||||
}
|
||||
}
|
||||
@@ -1 +1,241 @@
|
||||
// LSP client pool + diagnostics service land here in M5.
|
||||
//! LSP diagnostics pool: maps a file extension to a language server, lazily spawns one server
|
||||
//! per language, and answers the core `DiagnosticsSource` seam. Built-in servers are only
|
||||
//! offered when their binary is on `PATH`; config can add or override them.
|
||||
//! See `docs/09-integrations.md`.
|
||||
|
||||
mod client;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use harness_core::lsp::{Diagnostic, DiagnosticsSource};
|
||||
|
||||
use client::LspClient;
|
||||
|
||||
/// One language server: which binary to run and which extensions it handles.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ServerConfig {
|
||||
pub name: String,
|
||||
pub command: String,
|
||||
pub args: Vec<String>,
|
||||
pub extensions: Vec<String>,
|
||||
}
|
||||
|
||||
/// Built-in servers, tried only when their binary exists on `PATH`.
|
||||
fn builtin_servers() -> Vec<ServerConfig> {
|
||||
vec![
|
||||
ServerConfig {
|
||||
name: "rust".into(),
|
||||
command: "rust-analyzer".into(),
|
||||
args: vec![],
|
||||
extensions: vec!["rs".into()],
|
||||
},
|
||||
ServerConfig {
|
||||
name: "typescript".into(),
|
||||
command: "typescript-language-server".into(),
|
||||
args: vec!["--stdio".into()],
|
||||
extensions: vec!["ts".into(), "tsx".into(), "js".into(), "jsx".into()],
|
||||
},
|
||||
ServerConfig {
|
||||
name: "go".into(),
|
||||
command: "gopls".into(),
|
||||
args: vec![],
|
||||
extensions: vec!["go".into()],
|
||||
},
|
||||
ServerConfig {
|
||||
name: "python".into(),
|
||||
command: "pyright-langserver".into(),
|
||||
args: vec!["--stdio".into()],
|
||||
extensions: vec!["py".into()],
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
/// LSP `languageId` for a file extension (falls back to the extension itself).
|
||||
fn language_id_for_ext(ext: &str) -> &str {
|
||||
match ext {
|
||||
"rs" => "rust",
|
||||
"ts" => "typescript",
|
||||
"tsx" => "typescriptreact",
|
||||
"js" => "javascript",
|
||||
"jsx" => "javascriptreact",
|
||||
"go" => "go",
|
||||
"py" => "python",
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns true if `command` is an absolute existing file or resolvable on `PATH`.
|
||||
fn binary_exists(command: &str) -> bool {
|
||||
let p = Path::new(command);
|
||||
if p.is_absolute() {
|
||||
return p.is_file();
|
||||
}
|
||||
let Some(paths) = std::env::var_os("PATH") else {
|
||||
return false;
|
||||
};
|
||||
std::env::split_paths(&paths).any(|dir| dir.join(command).is_file())
|
||||
}
|
||||
|
||||
enum Slot {
|
||||
Ready(Arc<LspClient>),
|
||||
Failed,
|
||||
}
|
||||
|
||||
pub struct LspPool {
|
||||
root: PathBuf,
|
||||
servers: Vec<ServerConfig>,
|
||||
/// One slot per server name; absent until first spawn attempt.
|
||||
clients: tokio::sync::Mutex<HashMap<String, Slot>>,
|
||||
/// Open documents → last-sent version, so `didChange` bumps monotonically.
|
||||
open: Mutex<HashMap<PathBuf, i32>>,
|
||||
}
|
||||
|
||||
impl LspPool {
|
||||
/// Builds a pool: available built-ins plus any `config` servers (which override built-ins
|
||||
/// by name; an empty command disables a built-in).
|
||||
pub fn new(root: PathBuf, config: Vec<ServerConfig>) -> Self {
|
||||
let mut by_name: HashMap<String, ServerConfig> = HashMap::new();
|
||||
for server in builtin_servers() {
|
||||
if binary_exists(&server.command) {
|
||||
by_name.insert(server.name.clone(), server);
|
||||
}
|
||||
}
|
||||
for server in config {
|
||||
if server.command.is_empty() {
|
||||
by_name.remove(&server.name);
|
||||
} else {
|
||||
by_name.insert(server.name.clone(), server);
|
||||
}
|
||||
}
|
||||
Self {
|
||||
root,
|
||||
servers: by_name.into_values().collect(),
|
||||
clients: tokio::sync::Mutex::new(HashMap::new()),
|
||||
open: Mutex::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn server_for(&self, path: &Path) -> Option<&ServerConfig> {
|
||||
let ext = path.extension().and_then(|e| e.to_str())?;
|
||||
self.servers
|
||||
.iter()
|
||||
.find(|s| s.extensions.iter().any(|e| e == ext))
|
||||
}
|
||||
|
||||
/// Gets or lazily spawns the client for `server`. Caches a failure so we don't respawn a
|
||||
/// broken server on every edit.
|
||||
async fn client_for(&self, server: &ServerConfig) -> Option<Arc<LspClient>> {
|
||||
let mut clients = self.clients.lock().await;
|
||||
match clients.get(&server.name) {
|
||||
Some(Slot::Ready(c)) => return Some(c.clone()),
|
||||
Some(Slot::Failed) => return None,
|
||||
None => {}
|
||||
}
|
||||
match LspClient::spawn(&server.command, &server.args, &self.root).await {
|
||||
Ok(client) => {
|
||||
let client = Arc::new(client);
|
||||
clients.insert(server.name.clone(), Slot::Ready(client.clone()));
|
||||
Some(client)
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(server = %server.name, error = %e, "LSP server failed to start");
|
||||
clients.insert(server.name.clone(), Slot::Failed);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl DiagnosticsSource for LspPool {
|
||||
async fn touch(&self, path: &Path) {
|
||||
let Some(server) = self.server_for(path).cloned() else {
|
||||
return;
|
||||
};
|
||||
let Ok(text) = tokio::fs::read_to_string(path).await else {
|
||||
return;
|
||||
};
|
||||
let Some(client) = self.client_for(&server).await else {
|
||||
return;
|
||||
};
|
||||
let ext = path
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.unwrap_or_default();
|
||||
let language_id = language_id_for_ext(ext);
|
||||
|
||||
let existing_version = {
|
||||
let mut open = self.open.lock().unwrap();
|
||||
match open.get_mut(path) {
|
||||
Some(version) => {
|
||||
*version += 1;
|
||||
Some(*version)
|
||||
}
|
||||
None => {
|
||||
open.insert(path.to_path_buf(), 1);
|
||||
None
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Fresh diagnostics only reflect the current contents — drop any stale set first.
|
||||
client.clear(path);
|
||||
match existing_version {
|
||||
None => client.did_open(path, language_id, 1, &text),
|
||||
Some(version) => client.did_change(path, version, &text),
|
||||
}
|
||||
}
|
||||
|
||||
async fn diagnostics(&self, path: &Path, wait: Duration) -> Vec<Diagnostic> {
|
||||
let Some(server) = self.server_for(path).cloned() else {
|
||||
return Vec::new();
|
||||
};
|
||||
let Some(client) = self.client_for(&server).await else {
|
||||
return Vec::new();
|
||||
};
|
||||
client.wait_diagnostics(path, wait).await
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn language_ids_map_common_extensions() {
|
||||
assert_eq!(language_id_for_ext("rs"), "rust");
|
||||
assert_eq!(language_id_for_ext("tsx"), "typescriptreact");
|
||||
assert_eq!(language_id_for_ext("unknown"), "unknown");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn config_server_overrides_builtin_by_name() {
|
||||
let pool = LspPool::new(
|
||||
std::env::temp_dir(),
|
||||
vec![ServerConfig {
|
||||
name: "rust".into(),
|
||||
command: "my-custom-ra".into(),
|
||||
args: vec!["--flag".into()],
|
||||
extensions: vec!["rs".into()],
|
||||
}],
|
||||
);
|
||||
let server = pool.server_for(Path::new("/x/a.rs")).unwrap();
|
||||
assert_eq!(server.command, "my-custom-ra");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_server_for_unknown_extension() {
|
||||
let pool = LspPool::new(std::env::temp_dir(), vec![]);
|
||||
assert!(pool.server_for(Path::new("/x/a.zzz")).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_exists_finds_a_known_tool() {
|
||||
assert!(binary_exists("sh"));
|
||||
assert!(!binary_exists("definitely-not-a-real-binary-xyz"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
//! End-to-end LSP test against a real `rust-analyzer`. Ignored by default (requires the binary
|
||||
//! on PATH and is slow — RA indexes the project). Run with:
|
||||
//! cargo test -p harness-lsp --test rust_analyzer -- --ignored --nocapture
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use harness_core::lsp::{DiagnosticsSource, Severity};
|
||||
use harness_lsp::LspPool;
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires rust-analyzer on PATH; slow"]
|
||||
async fn rust_analyzer_reports_a_type_error() {
|
||||
// Minimal cargo project with a deliberate type error.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(
|
||||
dir.path().join("Cargo.toml"),
|
||||
"[package]\nname = \"probe\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
|
||||
)
|
||||
.unwrap();
|
||||
std::fs::create_dir_all(dir.path().join("src")).unwrap();
|
||||
let main_rs = dir.path().join("src/main.rs");
|
||||
std::fs::write(
|
||||
&main_rs,
|
||||
"fn main() {\n let x: i32 = \"not an integer\";\n let _ = x;\n}\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let pool = LspPool::new(dir.path().to_path_buf(), vec![]);
|
||||
pool.touch(&main_rs).await;
|
||||
// Generous wait: rust-analyzer must index the workspace before it reports anything.
|
||||
let diags = pool.diagnostics(&main_rs, Duration::from_secs(60)).await;
|
||||
|
||||
println!("diagnostics: {diags:#?}");
|
||||
assert!(
|
||||
diags.iter().any(|d| d.severity == Severity::Error),
|
||||
"expected at least one error diagnostic, got: {diags:?}"
|
||||
);
|
||||
}
|
||||
@@ -6,6 +6,17 @@ license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
harness-core = { workspace = true }
|
||||
rmcp = { workspace = true, features = ["client", "transport-child-process"] }
|
||||
tokio = { workspace = true }
|
||||
async-trait = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio-util = { workspace = true }
|
||||
tempfile = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -1 +1,264 @@
|
||||
// MCP stdio client → Tool adapters land here in M5.
|
||||
//! MCP stdio client → `Tool` adapters (M5). For each configured server we spawn the child
|
||||
//! over rmcp's `TokioChildProcess` transport, `initialize`, `list_tools`, and wrap every
|
||||
//! remote tool as an [`McpTool`] named `{server}_{tool}`. Servers are gated behind the `mcp`
|
||||
//! permission key. See `docs/09-integrations.md`.
|
||||
//!
|
||||
//! Out of scope for v1 (matching the doc): resources, prompts, sampling, and non-stdio
|
||||
//! transports.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
|
||||
use rmcp::model::{CallToolRequestParam, RawContent};
|
||||
use rmcp::service::RunningService;
|
||||
use rmcp::transport::TokioChildProcess;
|
||||
use rmcp::{RoleClient, ServiceExt};
|
||||
use tokio::process::Command;
|
||||
|
||||
/// Sanitized-name cap so a `{server}_{tool}` name stays a legal tool identifier.
|
||||
const MAX_NAME_LEN: usize = 64;
|
||||
|
||||
/// One configured MCP server: the child command plus its environment.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ServerConfig {
|
||||
pub command: String,
|
||||
pub args: Vec<String>,
|
||||
pub env: HashMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum ConnectError {
|
||||
#[error("spawn/transport failed: {0}")]
|
||||
Transport(#[from] std::io::Error),
|
||||
#[error("MCP service error: {0}")]
|
||||
Service(#[from] rmcp::service::ServiceError),
|
||||
}
|
||||
|
||||
/// Connects every configured server and returns ready tool adapters. A server that fails to
|
||||
/// start (or list its tools) is logged and skipped — the rest are unaffected, and the session
|
||||
/// still runs with whatever connected. `named_servers` is `{server_name: config}`.
|
||||
pub async fn connect_all(named_servers: HashMap<String, ServerConfig>) -> Vec<Arc<dyn Tool>> {
|
||||
let mut tools: Vec<Arc<dyn Tool>> = Vec::new();
|
||||
for (name, config) in named_servers {
|
||||
match connect(&name, &config).await {
|
||||
Ok(mut server_tools) => {
|
||||
tracing::info!(server = %name, count = server_tools.len(), "MCP server connected");
|
||||
tools.append(&mut server_tools);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(server = %name, error = %e, "MCP server failed to start; skipping");
|
||||
}
|
||||
}
|
||||
}
|
||||
tools
|
||||
}
|
||||
|
||||
async fn connect(name: &str, config: &ServerConfig) -> Result<Vec<Arc<dyn Tool>>, ConnectError> {
|
||||
let mut command = Command::new(&config.command);
|
||||
command.args(&config.args);
|
||||
for (key, value) in &config.env {
|
||||
command.env(key, value);
|
||||
}
|
||||
// rmcp sets stdin/stdout to piped and kill-on-drop; the child dies with the `RunningService`.
|
||||
let transport = TokioChildProcess::new(&mut command)?;
|
||||
let service = Arc::new(().serve(transport).await?);
|
||||
|
||||
let remote_tools = service.peer().list_all_tools().await?;
|
||||
let adapters = remote_tools
|
||||
.into_iter()
|
||||
.map(|tool| {
|
||||
let full_name = qualified_name(name, &tool.name);
|
||||
let parameters = serde_json::Value::Object((*tool.input_schema).clone());
|
||||
Arc::new(McpTool {
|
||||
full_name,
|
||||
remote_name: tool.name.to_string(),
|
||||
description: tool.description.to_string(),
|
||||
parameters,
|
||||
service: service.clone(),
|
||||
}) as Arc<dyn Tool>
|
||||
})
|
||||
.collect();
|
||||
Ok(adapters)
|
||||
}
|
||||
|
||||
/// `{server}_{tool}` sanitized to `[a-zA-Z0-9_-]` and capped at [`MAX_NAME_LEN`] chars.
|
||||
fn qualified_name(server: &str, tool: &str) -> String {
|
||||
let mut name: String = format!("{server}_{tool}")
|
||||
.chars()
|
||||
.map(|c| {
|
||||
if c.is_ascii_alphanumeric() || c == '_' || c == '-' {
|
||||
c
|
||||
} else {
|
||||
'_'
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
name.truncate(MAX_NAME_LEN);
|
||||
name
|
||||
}
|
||||
|
||||
/// A single remote MCP tool exposed to the engine as a `Tool`. Holds a shared handle to the
|
||||
/// server's `RunningService` (kept alive for the whole session so the child stays up).
|
||||
struct McpTool {
|
||||
/// Engine-facing name: sanitized `{server}_{tool}`; also the permission pattern.
|
||||
full_name: String,
|
||||
/// The server's own tool name, sent back verbatim in `call_tool`.
|
||||
remote_name: String,
|
||||
description: String,
|
||||
parameters: serde_json::Value,
|
||||
service: Arc<RunningService<RoleClient, ()>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for McpTool {
|
||||
fn name(&self) -> &str {
|
||||
&self.full_name
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
&self.description
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
self.parameters.clone()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
ctx.ask
|
||||
.ask(
|
||||
"mcp",
|
||||
self.full_name.clone(),
|
||||
self.full_name.clone(),
|
||||
input.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let arguments = match input {
|
||||
serde_json::Value::Object(map) => Some(map),
|
||||
serde_json::Value::Null => None,
|
||||
other => {
|
||||
return Err(ToolError::Invalid(format!(
|
||||
"MCP tool arguments must be a JSON object, got {other}"
|
||||
)))
|
||||
}
|
||||
};
|
||||
|
||||
let result = self
|
||||
.service
|
||||
.peer()
|
||||
.call_tool(CallToolRequestParam {
|
||||
name: self.remote_name.clone().into(),
|
||||
arguments,
|
||||
})
|
||||
.await
|
||||
.map_err(|e| ToolError::Other(e.to_string()))?;
|
||||
|
||||
let mut text = String::new();
|
||||
let mut image_index = 0;
|
||||
for content in &result.content {
|
||||
match &content.raw {
|
||||
RawContent::Text(t) => {
|
||||
if !text.is_empty() {
|
||||
text.push('\n');
|
||||
}
|
||||
text.push_str(&t.text);
|
||||
}
|
||||
RawContent::Image(image) => {
|
||||
let note = save_image(&ctx.data_dir, &self.full_name, image_index, image).await;
|
||||
if !text.is_empty() {
|
||||
text.push('\n');
|
||||
}
|
||||
text.push_str(¬e);
|
||||
image_index += 1;
|
||||
}
|
||||
RawContent::Resource(resource) => {
|
||||
let embedded = resource_text(resource);
|
||||
if !embedded.is_empty() {
|
||||
if !text.is_empty() {
|
||||
text.push('\n');
|
||||
}
|
||||
text.push_str(&embedded);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MCP surfaces tool-level failures as `is_error` with the message in `content`; map
|
||||
// that to a tool error so the model sees it as a failed call rather than a result.
|
||||
if result.is_error.unwrap_or(false) {
|
||||
return Err(ToolError::Other(if text.is_empty() {
|
||||
"MCP tool reported an error".to_string()
|
||||
} else {
|
||||
text
|
||||
}));
|
||||
}
|
||||
|
||||
Ok(ToolOutput::new(self.full_name.clone(), text))
|
||||
}
|
||||
}
|
||||
|
||||
/// Writes an image payload to the session data dir and returns a one-line note for the tool
|
||||
/// output. Best-effort: a write failure still yields a note (without a path).
|
||||
async fn save_image(
|
||||
data_dir: &PathBuf,
|
||||
tool_name: &str,
|
||||
index: usize,
|
||||
image: &rmcp::model::RawImageContent,
|
||||
) -> String {
|
||||
let ext = image.mime_type.rsplit('/').next().unwrap_or("bin");
|
||||
let file_name = format!("{tool_name}-image-{index}.{ext}.b64");
|
||||
let path = data_dir.join(&file_name);
|
||||
let saved = tokio::fs::create_dir_all(data_dir).await.is_ok()
|
||||
&& tokio::fs::write(&path, &image.data).await.is_ok();
|
||||
if saved {
|
||||
format!(
|
||||
"[image: {} ({} base64 bytes) saved to {}]",
|
||||
image.mime_type,
|
||||
image.data.len(),
|
||||
path.display()
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"[image: {} ({} base64 bytes, not saved)]",
|
||||
image.mime_type,
|
||||
image.data.len()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Best-effort text extraction from an embedded resource (text resources only in v1).
|
||||
fn resource_text(resource: &rmcp::model::RawEmbeddedResource) -> String {
|
||||
match &resource.resource {
|
||||
rmcp::model::ResourceContents::TextResourceContents { text, .. } => text.clone(),
|
||||
_ => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn qualified_name_prefixes_and_sanitizes() {
|
||||
assert_eq!(qualified_name("fs", "read_file"), "fs_read_file");
|
||||
assert_eq!(
|
||||
qualified_name("my.server", "do/thing"),
|
||||
"my_server_do_thing"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn qualified_name_caps_length() {
|
||||
let long_tool = "t".repeat(100);
|
||||
let name = qualified_name("srv", &long_tool);
|
||||
assert_eq!(name.len(), MAX_NAME_LEN);
|
||||
assert!(name.starts_with("srv_t"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
//! End-to-end MCP integration test: spawn a real stdio MCP server (a small Python fixture),
|
||||
//! connect through the real rmcp client, and verify a discovered tool is callable and gated
|
||||
//! behind the `mcp` permission key. This is the M5 milestone's ✅ for MCP.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use harness_core::event::{AppEvent, EventBus};
|
||||
use harness_core::permission::{PermissionReply, PermissionService};
|
||||
use harness_core::tool::{MetadataSink, PermissionHandle, ToolCtx, ToolError};
|
||||
use harness_core::types::{MessageId, SessionId};
|
||||
use harness_mcp::{connect_all, ServerConfig};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
/// A permission frontend that replies `Once` to every ask and records the `permission`/`pattern`
|
||||
/// of each, so a test can assert the call was actually gated.
|
||||
fn recording_auto_approve(
|
||||
bus: EventBus,
|
||||
service: Arc<PermissionService>,
|
||||
) -> Arc<Mutex<Vec<(String, String)>>> {
|
||||
let asks = Arc::new(Mutex::new(Vec::new()));
|
||||
let asks_task = asks.clone();
|
||||
// Subscribe before spawning: a subscription created inside the task could miss the ask
|
||||
// (tokio broadcast only delivers to receivers that exist at publish time).
|
||||
let mut rx = bus.subscribe();
|
||||
tokio::spawn(async move {
|
||||
while let Ok(event) = rx.recv().await {
|
||||
if let AppEvent::PermissionAsked { request } = event {
|
||||
asks_task
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((request.permission.clone(), request.pattern.clone()));
|
||||
service.reply(&request.id, PermissionReply::Once);
|
||||
}
|
||||
}
|
||||
});
|
||||
asks
|
||||
}
|
||||
|
||||
struct Harness {
|
||||
ctx_data_dir: std::path::PathBuf,
|
||||
service: Arc<PermissionService>,
|
||||
}
|
||||
|
||||
impl Harness {
|
||||
fn ctx(&self) -> ToolCtx {
|
||||
let (metadata, _rx) = MetadataSink::channel();
|
||||
ToolCtx {
|
||||
session_id: SessionId::new(),
|
||||
message_id: MessageId::new(),
|
||||
call_id: "call_1".into(),
|
||||
data_dir: self.ctx_data_dir.clone(),
|
||||
cwd: std::env::temp_dir(),
|
||||
cancel: CancellationToken::new(),
|
||||
ask: PermissionHandle::new(
|
||||
self.service.clone(),
|
||||
SessionId::new(),
|
||||
Vec::new(),
|
||||
Arc::new(Mutex::new(Vec::new())),
|
||||
CancellationToken::new(),
|
||||
),
|
||||
metadata,
|
||||
spawner: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn fixture_server() -> HashMap<String, ServerConfig> {
|
||||
let script = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/echo_server.py");
|
||||
HashMap::from([(
|
||||
"fix".to_string(),
|
||||
ServerConfig {
|
||||
command: "python3".to_string(),
|
||||
args: vec![script.to_string()],
|
||||
env: HashMap::new(),
|
||||
},
|
||||
)])
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discovers_and_calls_a_real_mcp_tool_with_permission() {
|
||||
let tools = connect_all(fixture_server()).await;
|
||||
let names: Vec<_> = tools.iter().map(|t| t.name().to_string()).collect();
|
||||
assert!(
|
||||
names.contains(&"fix_echo".to_string()),
|
||||
"expected fix_echo among {names:?}"
|
||||
);
|
||||
assert!(names.contains(&"fix_boom".to_string()));
|
||||
|
||||
let echo = tools.iter().find(|t| t.name() == "fix_echo").unwrap();
|
||||
// Schema passes through untouched from the server.
|
||||
assert_eq!(echo.parameters()["properties"]["text"]["type"], "string");
|
||||
|
||||
let bus = EventBus::new();
|
||||
let service = Arc::new(PermissionService::new(bus.clone()));
|
||||
let asks = recording_auto_approve(bus, service.clone());
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let harness = Harness {
|
||||
ctx_data_dir: dir.path().to_path_buf(),
|
||||
service,
|
||||
};
|
||||
|
||||
let out = echo
|
||||
.execute(serde_json::json!({"text": "hi there"}), harness.ctx())
|
||||
.await
|
||||
.expect("echo call succeeds");
|
||||
assert_eq!(out.output, "hi there");
|
||||
|
||||
// The call was gated on the `mcp` key with the qualified tool name as the pattern.
|
||||
let recorded = asks.lock().unwrap().clone();
|
||||
assert_eq!(recorded, vec![("mcp".to_string(), "fix_echo".to_string())]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tool_error_result_maps_to_tool_error() {
|
||||
let tools = connect_all(fixture_server()).await;
|
||||
let boom = tools.iter().find(|t| t.name() == "fix_boom").unwrap();
|
||||
|
||||
let bus = EventBus::new();
|
||||
let service = Arc::new(PermissionService::new(bus.clone()));
|
||||
let _asks = recording_auto_approve(bus, service.clone());
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let harness = Harness {
|
||||
ctx_data_dir: dir.path().to_path_buf(),
|
||||
service,
|
||||
};
|
||||
|
||||
let err = boom
|
||||
.execute(serde_json::json!({}), harness.ctx())
|
||||
.await
|
||||
.expect_err("boom reports an error result");
|
||||
match err {
|
||||
ToolError::Other(msg) => assert_eq!(msg, "kaboom"),
|
||||
other => panic!("expected ToolError::Other, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_failed_server_is_skipped_not_fatal() {
|
||||
let servers = HashMap::from([(
|
||||
"broken".to_string(),
|
||||
ServerConfig {
|
||||
command: "definitely-not-a-real-binary-xyz".to_string(),
|
||||
args: vec![],
|
||||
env: HashMap::new(),
|
||||
},
|
||||
)]);
|
||||
// No panic, no tools — the missing server is logged and skipped.
|
||||
let tools = connect_all(servers).await;
|
||||
assert!(tools.is_empty());
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Minimal MCP stdio server fixture for harness-mcp integration tests.
|
||||
|
||||
Speaks newline-delimited JSON-RPC (the framing rmcp's child-process transport uses) and
|
||||
implements just enough of the protocol to be discovered and called: `initialize`,
|
||||
`notifications/initialized`, `tools/list`, and `tools/call`. Exposes one tool, `echo`,
|
||||
which returns its `text` argument, plus `boom`, which returns an error result.
|
||||
"""
|
||||
import json
|
||||
import sys
|
||||
|
||||
PROTOCOL_VERSION = "2024-11-05"
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"name": "echo",
|
||||
"description": "Returns the text it is given.",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {"text": {"type": "string"}},
|
||||
"required": ["text"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "boom",
|
||||
"description": "Always fails.",
|
||||
"inputSchema": {"type": "object", "properties": {}},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def reply(msg_id, result):
|
||||
sys.stdout.write(json.dumps({"jsonrpc": "2.0", "id": msg_id, "result": result}) + "\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def main():
|
||||
# readline() rather than `for line in sys.stdin`: the latter's read-ahead buffer blocks
|
||||
# until it fills, which would stall the JSON-RPC handshake line-by-line.
|
||||
while True:
|
||||
line = sys.stdin.readline()
|
||||
if line == "": # EOF: parent closed stdin
|
||||
break
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
msg = json.loads(line)
|
||||
method = msg.get("method")
|
||||
msg_id = msg.get("id")
|
||||
|
||||
if method == "initialize":
|
||||
reply(msg_id, {
|
||||
"protocolVersion": PROTOCOL_VERSION,
|
||||
"capabilities": {"tools": {}},
|
||||
"serverInfo": {"name": "echo-fixture", "version": "0.1.0"},
|
||||
})
|
||||
elif method == "notifications/initialized":
|
||||
pass # notification: no response
|
||||
elif method == "tools/list":
|
||||
reply(msg_id, {"tools": TOOLS})
|
||||
elif method == "tools/call":
|
||||
params = msg.get("params") or {}
|
||||
name = params.get("name")
|
||||
args = params.get("arguments") or {}
|
||||
if name == "echo":
|
||||
reply(msg_id, {
|
||||
"content": [{"type": "text", "text": args.get("text", "")}],
|
||||
"isError": False,
|
||||
})
|
||||
elif name == "boom":
|
||||
reply(msg_id, {
|
||||
"content": [{"type": "text", "text": "kaboom"}],
|
||||
"isError": True,
|
||||
})
|
||||
else:
|
||||
reply(msg_id, {
|
||||
"content": [{"type": "text", "text": f"unknown tool {name}"}],
|
||||
"isError": True,
|
||||
})
|
||||
elif msg_id is not None:
|
||||
# Unknown request: empty result keeps the client happy.
|
||||
reply(msg_id, {})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -6,6 +6,23 @@ 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 }
|
||||
dirs = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["test-util", "macros"] }
|
||||
tempfile = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -0,0 +1,67 @@
|
||||
{
|
||||
"anthropic": {
|
||||
"id": "anthropic",
|
||||
"name": "Anthropic",
|
||||
"models": {
|
||||
"claude-sonnet-4-5": {
|
||||
"id": "claude-sonnet-4-5",
|
||||
"reasoning": true,
|
||||
"tool_call": true,
|
||||
"attachment": true,
|
||||
"cost": { "input": 3, "output": 15, "cache_read": 0.3, "cache_write": 3.75 },
|
||||
"limit": { "context": 200000, "output": 64000 }
|
||||
},
|
||||
"claude-opus-4-1": {
|
||||
"id": "claude-opus-4-1",
|
||||
"reasoning": true,
|
||||
"tool_call": true,
|
||||
"attachment": true,
|
||||
"cost": { "input": 15, "output": 75, "cache_read": 1.5, "cache_write": 18.75 },
|
||||
"limit": { "context": 200000, "output": 32000 }
|
||||
},
|
||||
"claude-haiku-4-5": {
|
||||
"id": "claude-haiku-4-5",
|
||||
"reasoning": true,
|
||||
"tool_call": true,
|
||||
"attachment": true,
|
||||
"cost": { "input": 1, "output": 5, "cache_read": 0.1, "cache_write": 1.25 },
|
||||
"limit": { "context": 200000, "output": 64000 }
|
||||
}
|
||||
}
|
||||
},
|
||||
"openai": {
|
||||
"id": "openai",
|
||||
"name": "OpenAI",
|
||||
"models": {
|
||||
"gpt-4o": {
|
||||
"id": "gpt-4o",
|
||||
"tool_call": true,
|
||||
"attachment": true,
|
||||
"cost": { "input": 2.5, "output": 10, "cache_read": 1.25 },
|
||||
"limit": { "context": 128000, "output": 16384 }
|
||||
},
|
||||
"gpt-4o-mini": {
|
||||
"id": "gpt-4o-mini",
|
||||
"tool_call": true,
|
||||
"attachment": true,
|
||||
"cost": { "input": 0.15, "output": 0.6, "cache_read": 0.075 },
|
||||
"limit": { "context": 128000, "output": 16384 }
|
||||
},
|
||||
"gpt-5": {
|
||||
"id": "gpt-5",
|
||||
"reasoning": true,
|
||||
"tool_call": true,
|
||||
"attachment": true,
|
||||
"cost": { "input": 1.25, "output": 10, "cache_read": 0.125 },
|
||||
"limit": { "context": 400000, "output": 128000 }
|
||||
},
|
||||
"o3": {
|
||||
"id": "o3",
|
||||
"reasoning": true,
|
||||
"tool_call": true,
|
||||
"cost": { "input": 2, "output": 8, "cache_read": 0.5 },
|
||||
"limit": { "context": 200000, "output": 100000 }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,228 @@
|
||||
//! Credential storage for providers that authenticate outside the config file.
|
||||
//!
|
||||
//! A single JSON file (`~/.local/share/ai-harness/auth.json`, mode `0600`) keyed by provider
|
||||
//! id. OAuth providers (Copilot) store the token triple and refresh it in place; API-key
|
||||
//! providers store the bare key. Config-supplied keys take precedence over this file — this
|
||||
//! is only for interactively-obtained credentials.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum AuthRecord {
|
||||
/// OAuth credentials. `expires` is a unix-ms timestamp; `0` means the token never expires
|
||||
/// (opencode's direct-Bearer Copilot mode).
|
||||
OAuth {
|
||||
access: String,
|
||||
refresh: String,
|
||||
expires: i64,
|
||||
},
|
||||
Api {
|
||||
key: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl AuthRecord {
|
||||
/// True when an OAuth token is at or past `expires` (with a safety skew), given `now_ms`.
|
||||
/// API keys and never-expiring OAuth tokens (`expires == 0`) are never considered expired.
|
||||
pub fn is_expired(&self, now_ms: i64, skew_ms: i64) -> bool {
|
||||
match self {
|
||||
AuthRecord::OAuth { expires, .. } if *expires > 0 => now_ms + skew_ms >= *expires,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Read/modify/write access to the auth file. Cheap to construct; every operation reloads so
|
||||
/// concurrent writers (a refresh in one provider, a login in another) don't clobber each other.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AuthStorage {
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
impl AuthStorage {
|
||||
pub fn new(path: PathBuf) -> Self {
|
||||
Self { path }
|
||||
}
|
||||
|
||||
/// `~/.local/share/ai-harness/auth.json` (falling back to the temp dir if there is no
|
||||
/// data directory).
|
||||
pub fn default_path() -> PathBuf {
|
||||
dirs::data_dir()
|
||||
.unwrap_or_else(std::env::temp_dir)
|
||||
.join("ai-harness")
|
||||
.join("auth.json")
|
||||
}
|
||||
|
||||
pub fn with_default_path() -> Self {
|
||||
Self::new(Self::default_path())
|
||||
}
|
||||
|
||||
/// Loads all records. A missing file yields an empty map; a corrupt file is an error.
|
||||
pub fn load(&self) -> io::Result<HashMap<String, AuthRecord>> {
|
||||
match std::fs::read(&self.path) {
|
||||
Ok(bytes) => serde_json::from_slice(&bytes)
|
||||
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)),
|
||||
Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(HashMap::new()),
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get(&self, provider: &str) -> io::Result<Option<AuthRecord>> {
|
||||
Ok(self.load()?.remove(provider))
|
||||
}
|
||||
|
||||
pub fn set(&self, provider: &str, record: AuthRecord) -> io::Result<()> {
|
||||
let mut records = self.load()?;
|
||||
records.insert(provider.to_string(), record);
|
||||
self.write(&records)
|
||||
}
|
||||
|
||||
pub fn remove(&self, provider: &str) -> io::Result<()> {
|
||||
let mut records = self.load()?;
|
||||
if records.remove(provider).is_some() {
|
||||
self.write(&records)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Serializes `records` and writes them with `0600` permissions via a temp-file rename so
|
||||
/// a partial write can never leave a truncated auth file behind.
|
||||
fn write(&self, records: &HashMap<String, AuthRecord>) -> io::Result<()> {
|
||||
if let Some(parent) = self.path.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
let json = serde_json::to_vec_pretty(records)
|
||||
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
|
||||
|
||||
let tmp = self.path.with_extension("json.tmp");
|
||||
std::fs::write(&tmp, &json)?;
|
||||
set_owner_only(&tmp)?;
|
||||
std::fs::rename(&tmp, &self.path)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn set_owner_only(path: &Path) -> io::Result<()> {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn set_owner_only(_path: &Path) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn storage() -> (tempfile::TempDir, AuthStorage) {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("nested").join("auth.json");
|
||||
(dir, AuthStorage::new(path))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_file_loads_empty() {
|
||||
let (_dir, store) = storage();
|
||||
assert!(store.load().unwrap().is_empty());
|
||||
assert_eq!(store.get("anthropic").unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_get_roundtrip_and_overwrite() {
|
||||
let (_dir, store) = storage();
|
||||
store
|
||||
.set("anthropic", AuthRecord::Api { key: "k1".into() })
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store.get("anthropic").unwrap(),
|
||||
Some(AuthRecord::Api { key: "k1".into() })
|
||||
);
|
||||
|
||||
store
|
||||
.set("anthropic", AuthRecord::Api { key: "k2".into() })
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
store.get("anthropic").unwrap(),
|
||||
Some(AuthRecord::Api { key: "k2".into() })
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multiple_providers_coexist() {
|
||||
let (_dir, store) = storage();
|
||||
store
|
||||
.set("openai", AuthRecord::Api { key: "sk".into() })
|
||||
.unwrap();
|
||||
store
|
||||
.set(
|
||||
"github-copilot",
|
||||
AuthRecord::OAuth {
|
||||
access: "a".into(),
|
||||
refresh: "r".into(),
|
||||
expires: 0,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
let all = store.load().unwrap();
|
||||
assert_eq!(all.len(), 2);
|
||||
assert!(all.contains_key("openai"));
|
||||
assert!(all.contains_key("github-copilot"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_deletes_only_named_provider() {
|
||||
let (_dir, store) = storage();
|
||||
store
|
||||
.set("openai", AuthRecord::Api { key: "sk".into() })
|
||||
.unwrap();
|
||||
store
|
||||
.set("anthropic", AuthRecord::Api { key: "an".into() })
|
||||
.unwrap();
|
||||
store.remove("openai").unwrap();
|
||||
assert_eq!(store.get("openai").unwrap(), None);
|
||||
assert!(store.get("anthropic").unwrap().is_some());
|
||||
// Removing an absent provider is a no-op, not an error.
|
||||
store.remove("openai").unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oauth_expiry_respects_skew_and_never_expires() {
|
||||
let never = AuthRecord::OAuth {
|
||||
access: "a".into(),
|
||||
refresh: "r".into(),
|
||||
expires: 0,
|
||||
};
|
||||
assert!(!never.is_expired(i64::MAX, 0));
|
||||
|
||||
let expiring = AuthRecord::OAuth {
|
||||
access: "a".into(),
|
||||
refresh: "r".into(),
|
||||
expires: 1_000,
|
||||
};
|
||||
assert!(!expiring.is_expired(800, 120));
|
||||
assert!(expiring.is_expired(880, 120)); // within skew window
|
||||
assert!(expiring.is_expired(1_000, 0));
|
||||
|
||||
assert!(!AuthRecord::Api { key: "k".into() }.is_expired(i64::MAX, 0));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn file_is_written_owner_only() {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let (_dir, store) = storage();
|
||||
store
|
||||
.set("openai", AuthRecord::Api { key: "sk".into() })
|
||||
.unwrap();
|
||||
let mode = std::fs::metadata(&store.path).unwrap().permissions().mode();
|
||||
assert_eq!(mode & 0o777, 0o600);
|
||||
}
|
||||
}
|
||||
@@ -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,3 @@
|
||||
pub mod anthropic;
|
||||
pub mod openai_chat;
|
||||
pub mod openai_responses;
|
||||
@@ -0,0 +1,502 @@
|
||||
//! Request builder + SSE decoder for OpenAI's `/chat/completions` streaming API.
|
||||
//!
|
||||
//! Also used as the fallback codec for OpenAI-compatible endpoints (including Copilot models
|
||||
//! whose `supported_endpoints` list `/chat/completions`).
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_stream::try_stream;
|
||||
use eventsource_stream::Eventsource;
|
||||
use futures::Stream;
|
||||
use harness_core::llm::{
|
||||
FinishReason, LlmEvent, LlmEventStream, LlmRequest, ProviderError, ReasoningEffort, Role,
|
||||
WireContent,
|
||||
};
|
||||
use harness_core::types::TokenUsage;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn effort_str(effort: ReasoningEffort) -> &'static str {
|
||||
match effort {
|
||||
ReasoningEffort::Low => "low",
|
||||
ReasoningEffort::Medium => "medium",
|
||||
ReasoningEffort::High => "high",
|
||||
}
|
||||
}
|
||||
|
||||
fn role_str(role: Role) -> &'static str {
|
||||
match role {
|
||||
Role::System => "system",
|
||||
Role::User => "user",
|
||||
Role::Assistant => "assistant",
|
||||
Role::Tool => "tool",
|
||||
}
|
||||
}
|
||||
|
||||
/// Flattens our grouped `WireMessage`s into the flat chat-completions message list. Tool
|
||||
/// results become their own `role: "tool"` messages (one per result), which is what the API
|
||||
/// expects regardless of how the engine grouped them.
|
||||
fn build_messages(system: &[String], messages: &[harness_core::llm::WireMessage]) -> Vec<Value> {
|
||||
let mut out: Vec<Value> = Vec::new();
|
||||
|
||||
if !system.is_empty() {
|
||||
out.push(json!({"role": "system", "content": system.join("\n\n")}));
|
||||
}
|
||||
|
||||
for m in messages {
|
||||
// Tool results are always emitted as standalone `tool` messages.
|
||||
for c in &m.content {
|
||||
if let WireContent::ToolResult {
|
||||
call_id, output, ..
|
||||
} = c
|
||||
{
|
||||
out.push(json!({
|
||||
"role": "tool",
|
||||
"tool_call_id": call_id,
|
||||
"content": output,
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
let mut text = String::new();
|
||||
let mut images: Vec<Value> = Vec::new();
|
||||
let mut tool_calls: Vec<Value> = Vec::new();
|
||||
for c in &m.content {
|
||||
match c {
|
||||
WireContent::Text { text: t } => {
|
||||
if !text.is_empty() {
|
||||
text.push('\n');
|
||||
}
|
||||
text.push_str(t);
|
||||
}
|
||||
WireContent::ToolCall {
|
||||
call_id,
|
||||
name,
|
||||
input,
|
||||
} => tool_calls.push(json!({
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
"function": {"name": name, "arguments": input.to_string()},
|
||||
})),
|
||||
WireContent::Image { mime_type, data } => images.push(json!({
|
||||
"type": "image_url",
|
||||
"image_url": {"url": format!("data:{mime_type};base64,{data}")},
|
||||
})),
|
||||
WireContent::ToolResult { .. } => {} // handled above
|
||||
}
|
||||
}
|
||||
|
||||
// A message that carried nothing but tool results contributes no further entry.
|
||||
if text.is_empty() && images.is_empty() && tool_calls.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let mut msg = json!({"role": role_str(m.role)});
|
||||
if images.is_empty() {
|
||||
msg["content"] = json!(text);
|
||||
} else {
|
||||
let mut parts = vec![json!({"type": "text", "text": text})];
|
||||
parts.extend(images);
|
||||
msg["content"] = json!(parts);
|
||||
}
|
||||
if !tool_calls.is_empty() {
|
||||
msg["tool_calls"] = json!(tool_calls);
|
||||
}
|
||||
out.push(msg);
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
pub fn build_request(req: &LlmRequest) -> Value {
|
||||
let mut body = json!({
|
||||
"model": req.model,
|
||||
"messages": build_messages(&req.system, &req.messages),
|
||||
"stream": true,
|
||||
"stream_options": {"include_usage": true},
|
||||
});
|
||||
|
||||
if !req.tools.is_empty() {
|
||||
body["tools"] = Value::Array(
|
||||
req.tools
|
||||
.iter()
|
||||
.map(|t| {
|
||||
json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": t.name,
|
||||
"description": t.description,
|
||||
"parameters": t.parameters,
|
||||
},
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
if let Some(temp) = req.temperature {
|
||||
body["temperature"] = json!(temp);
|
||||
}
|
||||
if let Some(max) = req.max_tokens {
|
||||
body["max_completion_tokens"] = json!(max);
|
||||
}
|
||||
if let Some(effort) = req.reasoning.as_ref().and_then(|r| r.effort) {
|
||||
body["reasoning_effort"] = json!(effort_str(effort));
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
fn map_finish_reason(reason: &str) -> FinishReason {
|
||||
match reason {
|
||||
"stop" => FinishReason::Stop,
|
||||
"tool_calls" | "function_call" => FinishReason::ToolCalls,
|
||||
"length" => FinishReason::Length,
|
||||
"content_filter" => FinishReason::ContentFilter,
|
||||
other => FinishReason::Unknown(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ToolAccum {
|
||||
call_id: String,
|
||||
name: String,
|
||||
args: String,
|
||||
}
|
||||
|
||||
/// Decodes a chat-completions SSE byte stream into normalized `LlmEvent`s. Tool-call deltas
|
||||
/// are accumulated by their `index` and flushed as `ToolCall`s once the stream ends.
|
||||
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 usage = TokenUsage::default();
|
||||
let mut reason = FinishReason::Stop;
|
||||
let mut text_open = false;
|
||||
let mut reasoning_open = false;
|
||||
let mut tools: HashMap<u64, ToolAccum> = HashMap::new();
|
||||
let mut tool_order: Vec<u64> = Vec::new();
|
||||
|
||||
while let Some(item) = futures::StreamExt::next(&mut events).await {
|
||||
let event = item.map_err(|e| ProviderError::Decode(e.to_string()))?;
|
||||
let data = event.data.trim();
|
||||
if data.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if data == "[DONE]" {
|
||||
break;
|
||||
}
|
||||
let value: Value = serde_json::from_str(data)
|
||||
.map_err(|e| ProviderError::Decode(format!("{e}: {data}")))?;
|
||||
|
||||
if let Some(u) = value.get("usage").filter(|u| u.is_object()) {
|
||||
usage.input = u["prompt_tokens"].as_u64().unwrap_or(usage.input);
|
||||
usage.output = u["completion_tokens"].as_u64().unwrap_or(usage.output);
|
||||
usage.reasoning = u["completion_tokens_details"]["reasoning_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(usage.reasoning);
|
||||
usage.cache_read = u["prompt_tokens_details"]["cached_tokens"]
|
||||
.as_u64()
|
||||
.unwrap_or(usage.cache_read);
|
||||
}
|
||||
|
||||
let choice = &value["choices"][0];
|
||||
let delta = &choice["delta"];
|
||||
|
||||
if let Some(rc) = delta["reasoning_content"].as_str().filter(|s| !s.is_empty()) {
|
||||
if !reasoning_open {
|
||||
reasoning_open = true;
|
||||
yield LlmEvent::ReasoningStart { id: "reasoning".into() };
|
||||
}
|
||||
yield LlmEvent::ReasoningDelta { id: "reasoning".into(), text: rc.to_string() };
|
||||
}
|
||||
|
||||
if let Some(text) = delta["content"].as_str().filter(|s| !s.is_empty()) {
|
||||
if reasoning_open {
|
||||
reasoning_open = false;
|
||||
yield LlmEvent::ReasoningEnd { id: "reasoning".into(), signature: None };
|
||||
}
|
||||
if !text_open {
|
||||
text_open = true;
|
||||
yield LlmEvent::TextStart { id: "0".into() };
|
||||
}
|
||||
yield LlmEvent::TextDelta { id: "0".into(), text: text.to_string() };
|
||||
}
|
||||
|
||||
if let Some(calls) = delta["tool_calls"].as_array() {
|
||||
for call in calls {
|
||||
let index = call["index"].as_u64().unwrap_or(0);
|
||||
let entry = tools.entry(index).or_insert_with(|| {
|
||||
tool_order.push(index);
|
||||
ToolAccum::default()
|
||||
});
|
||||
if let Some(id) = call["id"].as_str().filter(|s| !s.is_empty()) {
|
||||
entry.call_id = id.to_string();
|
||||
}
|
||||
if let Some(name) = call["function"]["name"].as_str().filter(|s| !s.is_empty()) {
|
||||
entry.name = name.to_string();
|
||||
yield LlmEvent::ToolInputStart {
|
||||
call_id: entry.call_id.clone(),
|
||||
name: entry.name.clone(),
|
||||
};
|
||||
}
|
||||
if let Some(args) = call["function"]["arguments"].as_str().filter(|s| !s.is_empty()) {
|
||||
entry.args.push_str(args);
|
||||
yield LlmEvent::ToolInputDelta {
|
||||
call_id: entry.call_id.clone(),
|
||||
json: args.to_string(),
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(fr) = choice["finish_reason"].as_str() {
|
||||
reason = map_finish_reason(fr);
|
||||
}
|
||||
}
|
||||
|
||||
if reasoning_open {
|
||||
yield LlmEvent::ReasoningEnd { id: "reasoning".into(), signature: None };
|
||||
}
|
||||
if text_open {
|
||||
yield LlmEvent::TextEnd { id: "0".into() };
|
||||
}
|
||||
for index in tool_order {
|
||||
if let Some(entry) = tools.remove(&index) {
|
||||
let input: Value = if entry.args.trim().is_empty() {
|
||||
json!({})
|
||||
} else {
|
||||
serde_json::from_str(&entry.args).unwrap_or(Value::Null)
|
||||
};
|
||||
yield LlmEvent::ToolCall { call_id: entry.call_id, name: entry.name, input };
|
||||
}
|
||||
}
|
||||
yield LlmEvent::Finish { reason, usage };
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use futures::StreamExt;
|
||||
use harness_core::llm::{Initiator, 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))
|
||||
}
|
||||
|
||||
fn req(model: &str) -> LlmRequest {
|
||||
LlmRequest {
|
||||
model: model.into(),
|
||||
system: vec![],
|
||||
messages: vec![],
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
reasoning: None,
|
||||
initiator: Initiator::User,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_text_only_response() {
|
||||
let raw = concat!(
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"},\"finish_reason\":null}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"lo\"},\"finish_reason\":null}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
|
||||
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5}}\n\n",
|
||||
"data: [DONE]\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: "Hel".into()
|
||||
},
|
||||
LlmEvent::TextDelta {
|
||||
id: "0".into(),
|
||||
text: "lo".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_accumulated_by_index() {
|
||||
let raw = concat!(
|
||||
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"read\",\"arguments\":\"\"}}]}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"file\\\"\"}}]}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\":\\\"a.txt\\\"}\"}}]}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":8}}\n\n",
|
||||
"data: [DONE]\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_reasoning_content_before_text() {
|
||||
let raw = concat!(
|
||||
"data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"hmm\"}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{\"content\":\"answer\"}}]}\n\n",
|
||||
"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
|
||||
"data: [DONE]\n\n",
|
||||
);
|
||||
let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
LlmEvent::ReasoningStart {
|
||||
id: "reasoning".into()
|
||||
},
|
||||
LlmEvent::ReasoningDelta {
|
||||
id: "reasoning".into(),
|
||||
text: "hmm".into()
|
||||
},
|
||||
LlmEvent::ReasoningEnd {
|
||||
id: "reasoning".into(),
|
||||
signature: None
|
||||
},
|
||||
LlmEvent::TextStart { id: "0".into() },
|
||||
LlmEvent::TextDelta {
|
||||
id: "0".into(),
|
||||
text: "answer".into()
|
||||
},
|
||||
LlmEvent::TextEnd { id: "0".into() },
|
||||
LlmEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
usage: TokenUsage::default()
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_wraps_tools_in_function_envelope() {
|
||||
let mut r = req("gpt-4o");
|
||||
r.tools = vec![ToolSchema {
|
||||
name: "read".into(),
|
||||
description: "reads a file".into(),
|
||||
parameters: json!({"type": "object"}),
|
||||
}];
|
||||
let body = build_request(&r);
|
||||
assert_eq!(body["tools"][0]["type"], "function");
|
||||
assert_eq!(body["tools"][0]["function"]["name"], "read");
|
||||
assert_eq!(
|
||||
body["tools"][0]["function"]["parameters"],
|
||||
json!({"type": "object"})
|
||||
);
|
||||
assert_eq!(body["stream_options"]["include_usage"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_prepends_system_message() {
|
||||
let mut r = req("gpt-4o");
|
||||
r.system = vec!["env".into(), "agent".into()];
|
||||
r.messages = vec![WireMessage {
|
||||
role: Role::User,
|
||||
content: vec![WireContent::Text { text: "hi".into() }],
|
||||
}];
|
||||
let body = build_request(&r);
|
||||
let msgs = body["messages"].as_array().unwrap();
|
||||
assert_eq!(msgs[0]["role"], "system");
|
||||
assert_eq!(msgs[0]["content"], "env\n\nagent");
|
||||
assert_eq!(msgs[1]["role"], "user");
|
||||
assert_eq!(msgs[1]["content"], "hi");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_expands_tool_results_to_tool_messages() {
|
||||
let mut r = req("gpt-4o");
|
||||
r.messages = vec![WireMessage {
|
||||
role: Role::Tool,
|
||||
content: vec![WireContent::ToolResult {
|
||||
call_id: "call_1".into(),
|
||||
output: "contents".into(),
|
||||
is_error: false,
|
||||
}],
|
||||
}];
|
||||
let body = build_request(&r);
|
||||
let msgs = body["messages"].as_array().unwrap();
|
||||
assert_eq!(msgs.len(), 1);
|
||||
assert_eq!(msgs[0]["role"], "tool");
|
||||
assert_eq!(msgs[0]["tool_call_id"], "call_1");
|
||||
assert_eq!(msgs[0]["content"], "contents");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_emits_assistant_tool_calls() {
|
||||
let mut r = req("gpt-4o");
|
||||
r.messages = vec![WireMessage {
|
||||
role: Role::Assistant,
|
||||
content: vec![WireContent::ToolCall {
|
||||
call_id: "call_1".into(),
|
||||
name: "read".into(),
|
||||
input: json!({"file": "a.txt"}),
|
||||
}],
|
||||
}];
|
||||
let body = build_request(&r);
|
||||
let msgs = body["messages"].as_array().unwrap();
|
||||
assert_eq!(msgs[0]["tool_calls"][0]["id"], "call_1");
|
||||
assert_eq!(msgs[0]["tool_calls"][0]["function"]["name"], "read");
|
||||
assert_eq!(
|
||||
msgs[0]["tool_calls"][0]["function"]["arguments"],
|
||||
"{\"file\":\"a.txt\"}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_sets_reasoning_effort() {
|
||||
let mut r = req("o3");
|
||||
r.reasoning = Some(ReasoningOpts {
|
||||
effort: Some(ReasoningEffort::High),
|
||||
budget_tokens: None,
|
||||
});
|
||||
let body = build_request(&r);
|
||||
assert_eq!(body["reasoning_effort"], "high");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,402 @@
|
||||
//! Request builder + SSE decoder for OpenAI's `/responses` streaming API.
|
||||
//!
|
||||
//! The responses API is the preferred surface for `gpt-*` / `o-*` models: it carries
|
||||
//! reasoning items natively and reports reasoning-token usage separately.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_stream::try_stream;
|
||||
use eventsource_stream::Eventsource;
|
||||
use futures::Stream;
|
||||
use harness_core::llm::{
|
||||
FinishReason, LlmEvent, LlmEventStream, LlmRequest, ProviderError, ReasoningEffort, Role,
|
||||
WireContent,
|
||||
};
|
||||
use harness_core::types::TokenUsage;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
fn effort_str(effort: ReasoningEffort) -> &'static str {
|
||||
match effort {
|
||||
ReasoningEffort::Low => "low",
|
||||
ReasoningEffort::Medium => "medium",
|
||||
ReasoningEffort::High => "high",
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds the `input` item list. Text/image content become `message` items; tool calls and
|
||||
/// their results become `function_call` / `function_call_output` items (the responses API
|
||||
/// keeps these as top-level items rather than nesting them inside messages).
|
||||
fn build_input(messages: &[harness_core::llm::WireMessage]) -> Vec<Value> {
|
||||
let mut out: Vec<Value> = Vec::new();
|
||||
for m in messages {
|
||||
let (role, text_type) = match m.role {
|
||||
Role::Assistant => ("assistant", "output_text"),
|
||||
_ => ("user", "input_text"),
|
||||
};
|
||||
let mut content_parts: Vec<Value> = Vec::new();
|
||||
for c in &m.content {
|
||||
match c {
|
||||
WireContent::Text { text } => {
|
||||
content_parts.push(json!({"type": text_type, "text": text}));
|
||||
}
|
||||
WireContent::Image { mime_type, data } => {
|
||||
content_parts.push(json!({
|
||||
"type": "input_image",
|
||||
"image_url": format!("data:{mime_type};base64,{data}"),
|
||||
}));
|
||||
}
|
||||
WireContent::ToolCall {
|
||||
call_id,
|
||||
name,
|
||||
input,
|
||||
} => out.push(json!({
|
||||
"type": "function_call",
|
||||
"call_id": call_id,
|
||||
"name": name,
|
||||
"arguments": input.to_string(),
|
||||
})),
|
||||
WireContent::ToolResult {
|
||||
call_id, output, ..
|
||||
} => out.push(json!({
|
||||
"type": "function_call_output",
|
||||
"call_id": call_id,
|
||||
"output": output,
|
||||
})),
|
||||
}
|
||||
}
|
||||
if !content_parts.is_empty() {
|
||||
out.push(json!({"type": "message", "role": role, "content": content_parts}));
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
pub fn build_request(req: &LlmRequest) -> Value {
|
||||
let mut body = json!({
|
||||
"model": req.model,
|
||||
"input": build_input(&req.messages),
|
||||
"stream": true,
|
||||
"store": false,
|
||||
});
|
||||
|
||||
if !req.system.is_empty() {
|
||||
body["instructions"] = json!(req.system.join("\n\n"));
|
||||
}
|
||||
if !req.tools.is_empty() {
|
||||
body["tools"] = Value::Array(
|
||||
req.tools
|
||||
.iter()
|
||||
.map(|t| {
|
||||
json!({
|
||||
"type": "function",
|
||||
"name": t.name,
|
||||
"description": t.description,
|
||||
"parameters": t.parameters,
|
||||
})
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
if let Some(temp) = req.temperature {
|
||||
body["temperature"] = json!(temp);
|
||||
}
|
||||
if let Some(max) = req.max_tokens {
|
||||
body["max_output_tokens"] = json!(max);
|
||||
}
|
||||
if let Some(effort) = req.reasoning.as_ref().and_then(|r| r.effort) {
|
||||
body["reasoning"] = json!({"effort": effort_str(effort), "summary": "auto"});
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
fn map_status(status: Option<&str>, had_tool_call: bool) -> FinishReason {
|
||||
match status {
|
||||
Some("completed") if had_tool_call => FinishReason::ToolCalls,
|
||||
Some("completed") => FinishReason::Stop,
|
||||
Some("incomplete") => FinishReason::Length,
|
||||
Some(other) => FinishReason::Unknown(other.to_string()),
|
||||
None if had_tool_call => FinishReason::ToolCalls,
|
||||
None => FinishReason::Stop,
|
||||
}
|
||||
}
|
||||
|
||||
/// Tracks which normalized stream an `item_id` belongs to so deltas route correctly.
|
||||
enum ItemKind {
|
||||
Text,
|
||||
Reasoning,
|
||||
FunctionCall { call_id: String, name: String },
|
||||
}
|
||||
|
||||
/// Decodes a `/responses` SSE byte stream into normalized `LlmEvent`s.
|
||||
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 items: HashMap<String, ItemKind> = HashMap::new();
|
||||
let mut fn_args: HashMap<String, String> = HashMap::new();
|
||||
let mut usage = TokenUsage::default();
|
||||
let mut had_tool_call = false;
|
||||
|
||||
while let Some(item) = futures::StreamExt::next(&mut events).await {
|
||||
let event = item.map_err(|e| ProviderError::Decode(e.to_string()))?;
|
||||
if event.data.trim().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 {
|
||||
"response.output_item.added" => {
|
||||
let item = &value["item"];
|
||||
let id = item["id"].as_str().unwrap_or_default().to_string();
|
||||
match item["type"].as_str().unwrap_or_default() {
|
||||
"message" => {
|
||||
items.insert(id.clone(), ItemKind::Text);
|
||||
yield LlmEvent::TextStart { id };
|
||||
}
|
||||
"reasoning" => {
|
||||
items.insert(id.clone(), ItemKind::Reasoning);
|
||||
yield LlmEvent::ReasoningStart { id };
|
||||
}
|
||||
"function_call" => {
|
||||
let call_id = item["call_id"].as_str().unwrap_or_default().to_string();
|
||||
let name = item["name"].as_str().unwrap_or_default().to_string();
|
||||
fn_args.insert(id.clone(), String::new());
|
||||
items.insert(id, ItemKind::FunctionCall { call_id: call_id.clone(), name: name.clone() });
|
||||
had_tool_call = true;
|
||||
yield LlmEvent::ToolInputStart { call_id, name };
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
"response.output_text.delta" => {
|
||||
let id = value["item_id"].as_str().unwrap_or_default().to_string();
|
||||
let text = value["delta"].as_str().unwrap_or_default().to_string();
|
||||
yield LlmEvent::TextDelta { id, text };
|
||||
}
|
||||
"response.reasoning_summary_text.delta" => {
|
||||
let id = value["item_id"].as_str().unwrap_or_default().to_string();
|
||||
let text = value["delta"].as_str().unwrap_or_default().to_string();
|
||||
yield LlmEvent::ReasoningDelta { id, text };
|
||||
}
|
||||
"response.function_call_arguments.delta" => {
|
||||
let id = value["item_id"].as_str().unwrap_or_default().to_string();
|
||||
let delta = value["delta"].as_str().unwrap_or_default();
|
||||
if let (Some(buf), Some(ItemKind::FunctionCall { call_id, .. })) =
|
||||
(fn_args.get_mut(&id), items.get(&id))
|
||||
{
|
||||
buf.push_str(delta);
|
||||
yield LlmEvent::ToolInputDelta { call_id: call_id.clone(), json: delta.to_string() };
|
||||
}
|
||||
}
|
||||
"response.output_item.done" => {
|
||||
let id = value["item"]["id"].as_str().unwrap_or_default().to_string();
|
||||
match items.remove(&id) {
|
||||
Some(ItemKind::Text) => yield LlmEvent::TextEnd { id },
|
||||
Some(ItemKind::Reasoning) => yield LlmEvent::ReasoningEnd { id, signature: None },
|
||||
Some(ItemKind::FunctionCall { call_id, name }) => {
|
||||
let raw = fn_args.remove(&id).unwrap_or_default();
|
||||
let input: Value = if raw.trim().is_empty() {
|
||||
json!({})
|
||||
} else {
|
||||
serde_json::from_str(&raw).unwrap_or(Value::Null)
|
||||
};
|
||||
yield LlmEvent::ToolCall { call_id, name, input };
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
"response.completed" | "response.incomplete" => {
|
||||
let response = &value["response"];
|
||||
let u = &response["usage"];
|
||||
usage.input = u["input_tokens"].as_u64().unwrap_or(0);
|
||||
usage.output = u["output_tokens"].as_u64().unwrap_or(0);
|
||||
usage.reasoning = u["output_tokens_details"]["reasoning_tokens"].as_u64().unwrap_or(0);
|
||||
usage.cache_read = u["input_tokens_details"]["cached_tokens"].as_u64().unwrap_or(0);
|
||||
let status = response["status"].as_str();
|
||||
yield LlmEvent::Finish { reason: map_status(status, had_tool_call), usage };
|
||||
}
|
||||
"error" | "response.failed" => {
|
||||
let message = value["response"]["error"]["message"]
|
||||
.as_str()
|
||||
.or_else(|| value["message"].as_str())
|
||||
.unwrap_or("unknown error")
|
||||
.to_string();
|
||||
Err(ProviderError::Http { status: 0, body: message })?;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use futures::StreamExt;
|
||||
use harness_core::llm::{Initiator, 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))
|
||||
}
|
||||
|
||||
fn req(model: &str) -> LlmRequest {
|
||||
LlmRequest {
|
||||
model: model.into(),
|
||||
system: vec![],
|
||||
messages: vec![],
|
||||
tools: vec![],
|
||||
temperature: None,
|
||||
max_tokens: None,
|
||||
reasoning: None,
|
||||
initiator: Initiator::User,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_text_response() {
|
||||
let raw = concat!(
|
||||
"data: {\"type\":\"response.output_item.added\",\"item\":{\"id\":\"msg_1\",\"type\":\"message\"}}\n\n",
|
||||
"data: {\"type\":\"response.output_text.delta\",\"item_id\":\"msg_1\",\"delta\":\"Hello\"}\n\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"msg_1\"}}\n\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":10,\"output_tokens\":5}}}\n\n",
|
||||
);
|
||||
let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![
|
||||
LlmEvent::TextStart { id: "msg_1".into() },
|
||||
LlmEvent::TextDelta {
|
||||
id: "msg_1".into(),
|
||||
text: "Hello".into()
|
||||
},
|
||||
LlmEvent::TextEnd { id: "msg_1".into() },
|
||||
LlmEvent::Finish {
|
||||
reason: FinishReason::Stop,
|
||||
usage: TokenUsage {
|
||||
input: 10,
|
||||
output: 5,
|
||||
..Default::default()
|
||||
},
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decodes_function_call_with_reasoning_tokens() {
|
||||
let raw = concat!(
|
||||
"data: {\"type\":\"response.output_item.added\",\"item\":{\"id\":\"fc_1\",\"type\":\"function_call\",\"call_id\":\"call_1\",\"name\":\"read\"}}\n\n",
|
||||
"data: {\"type\":\"response.function_call_arguments.delta\",\"item_id\":\"fc_1\",\"delta\":\"{\\\"file\\\":\"}\n\n",
|
||||
"data: {\"type\":\"response.function_call_arguments.delta\",\"item_id\":\"fc_1\",\"delta\":\"\\\"a.txt\\\"}\"}\n\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"item\":{\"id\":\"fc_1\"}}\n\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":3,\"output_tokens\":9,\"output_tokens_details\":{\"reasoning_tokens\":4}}}}\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: 3,
|
||||
output: 9,
|
||||
reasoning: 4,
|
||||
..Default::default()
|
||||
},
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failed_response_surfaces_as_err() {
|
||||
let raw = "data: {\"type\":\"response.failed\",\"response\":{\"error\":{\"message\":\"boom\"}}}\n\n";
|
||||
let events: Vec<Result<LlmEvent, ProviderError>> = sse_stream(raw).collect().await;
|
||||
assert!(
|
||||
matches!(events.last(), Some(Err(ProviderError::Http { body, .. })) if body == "boom")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_puts_system_in_instructions() {
|
||||
let mut r = req("gpt-5");
|
||||
r.system = vec!["env".into(), "agent".into()];
|
||||
let body = build_request(&r);
|
||||
assert_eq!(body["instructions"], "env\n\nagent");
|
||||
assert!(body.get("messages").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_flattens_tool_calls_and_results_to_items() {
|
||||
let mut r = req("gpt-5");
|
||||
r.messages = vec![
|
||||
WireMessage {
|
||||
role: Role::Assistant,
|
||||
content: vec![WireContent::ToolCall {
|
||||
call_id: "call_1".into(),
|
||||
name: "read".into(),
|
||||
input: json!({"file": "a.txt"}),
|
||||
}],
|
||||
},
|
||||
WireMessage {
|
||||
role: Role::Tool,
|
||||
content: vec![WireContent::ToolResult {
|
||||
call_id: "call_1".into(),
|
||||
output: "contents".into(),
|
||||
is_error: false,
|
||||
}],
|
||||
},
|
||||
];
|
||||
let body = build_request(&r);
|
||||
let input = body["input"].as_array().unwrap();
|
||||
assert_eq!(input[0]["type"], "function_call");
|
||||
assert_eq!(input[0]["call_id"], "call_1");
|
||||
assert_eq!(input[0]["arguments"], "{\"file\":\"a.txt\"}");
|
||||
assert_eq!(input[1]["type"], "function_call_output");
|
||||
assert_eq!(input[1]["output"], "contents");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_request_uses_flat_function_tool_shape_and_reasoning() {
|
||||
let mut r = req("o3");
|
||||
r.tools = vec![ToolSchema {
|
||||
name: "read".into(),
|
||||
description: "reads a file".into(),
|
||||
parameters: json!({"type": "object"}),
|
||||
}];
|
||||
r.reasoning = Some(ReasoningOpts {
|
||||
effort: Some(ReasoningEffort::Medium),
|
||||
budget_tokens: None,
|
||||
});
|
||||
let body = build_request(&r);
|
||||
assert_eq!(body["tools"][0]["type"], "function");
|
||||
assert_eq!(body["tools"][0]["name"], "read");
|
||||
assert_eq!(body["reasoning"]["effort"], "medium");
|
||||
assert_eq!(body["reasoning"]["summary"], "auto");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
//! GitHub OAuth device flow (RFC 8628) for Copilot login.
|
||||
//!
|
||||
//! ai-harness must register its own GitHub OAuth app and supply its client id (we do not
|
||||
//! hardcode opencode's). The client id is passed in by the caller — see [`request_device_code`].
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
|
||||
const DEVICE_CODE_URL: &str = "https://github.com/login/device/code";
|
||||
const ACCESS_TOKEN_URL: &str = "https://github.com/login/oauth/access_token";
|
||||
const GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:device_code";
|
||||
/// Minimum scope needed to call the Copilot token-exchange endpoint.
|
||||
pub const SCOPE: &str = "read:user";
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct DeviceCode {
|
||||
pub device_code: String,
|
||||
pub user_code: String,
|
||||
pub verification_uri: String,
|
||||
/// Seconds between polls; GitHub requires honoring this and any `slow_down` bumps.
|
||||
#[serde(default = "default_interval")]
|
||||
pub interval: u64,
|
||||
#[serde(default)]
|
||||
pub expires_in: u64,
|
||||
}
|
||||
|
||||
fn default_interval() -> u64 {
|
||||
5
|
||||
}
|
||||
|
||||
/// Result of one poll of the access-token endpoint.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PollOutcome {
|
||||
/// The user hasn't authorized yet; keep polling at the current interval.
|
||||
Pending,
|
||||
/// GitHub asked us to slow down; add 5s to the interval (RFC 8628).
|
||||
SlowDown,
|
||||
/// Authorization complete.
|
||||
Success { access_token: String },
|
||||
/// Terminal failure (expired code, denied, unknown error) with a human-readable reason.
|
||||
Failed(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum DeviceFlowError {
|
||||
#[error("network: {0}")]
|
||||
Network(String),
|
||||
#[error("unexpected response: {0}")]
|
||||
Unexpected(String),
|
||||
}
|
||||
|
||||
/// Pure mapping of an access-token poll response body to a [`PollOutcome`], so the state
|
||||
/// machine is testable without a live GitHub.
|
||||
pub fn parse_poll_response(body: &Value) -> PollOutcome {
|
||||
if let Some(token) = body["access_token"].as_str() {
|
||||
return PollOutcome::Success {
|
||||
access_token: token.to_string(),
|
||||
};
|
||||
}
|
||||
match body["error"].as_str() {
|
||||
Some("authorization_pending") => PollOutcome::Pending,
|
||||
Some("slow_down") => PollOutcome::SlowDown,
|
||||
Some(other) => {
|
||||
let desc = body["error_description"].as_str().unwrap_or(other);
|
||||
PollOutcome::Failed(desc.to_string())
|
||||
}
|
||||
None => PollOutcome::Failed("no access_token and no error in response".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Applies a `slow_down` to a poll interval per RFC 8628 (+5s).
|
||||
pub fn bump_interval(interval: u64) -> u64 {
|
||||
interval + 5
|
||||
}
|
||||
|
||||
/// Step 1: request a device + user code for `client_id`.
|
||||
pub async fn request_device_code(
|
||||
client: &reqwest::Client,
|
||||
client_id: &str,
|
||||
) -> Result<DeviceCode, DeviceFlowError> {
|
||||
let resp = client
|
||||
.post(DEVICE_CODE_URL)
|
||||
.header("accept", "application/json")
|
||||
.json(&serde_json::json!({"client_id": client_id, "scope": SCOPE}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| DeviceFlowError::Network(e.to_string()))?;
|
||||
let value: Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| DeviceFlowError::Network(e.to_string()))?;
|
||||
serde_json::from_value(value.clone())
|
||||
.map_err(|_| DeviceFlowError::Unexpected(value.to_string()))
|
||||
}
|
||||
|
||||
/// Step 2 (single poll): exchange the device code for an access token, once.
|
||||
pub async fn poll_once(
|
||||
client: &reqwest::Client,
|
||||
client_id: &str,
|
||||
device_code: &str,
|
||||
) -> Result<PollOutcome, DeviceFlowError> {
|
||||
let resp = client
|
||||
.post(ACCESS_TOKEN_URL)
|
||||
.header("accept", "application/json")
|
||||
.json(&serde_json::json!({
|
||||
"client_id": client_id,
|
||||
"device_code": device_code,
|
||||
"grant_type": GRANT_TYPE,
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| DeviceFlowError::Network(e.to_string()))?;
|
||||
let value: Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| DeviceFlowError::Network(e.to_string()))?;
|
||||
Ok(parse_poll_response(&value))
|
||||
}
|
||||
|
||||
/// Step 2 (full loop): polls until success or terminal failure, honoring `interval` and
|
||||
/// `slow_down`. `sleep` is injected so tests can drive it without real time.
|
||||
pub async fn poll_for_token<S, Fut>(
|
||||
client: &reqwest::Client,
|
||||
client_id: &str,
|
||||
device: &DeviceCode,
|
||||
sleep: S,
|
||||
) -> Result<String, DeviceFlowError>
|
||||
where
|
||||
S: Fn(Duration) -> Fut,
|
||||
Fut: std::future::Future<Output = ()>,
|
||||
{
|
||||
let mut interval = device.interval;
|
||||
loop {
|
||||
sleep(Duration::from_secs(interval)).await;
|
||||
match poll_once(client, client_id, &device.device_code).await? {
|
||||
PollOutcome::Pending => {}
|
||||
PollOutcome::SlowDown => interval = bump_interval(interval),
|
||||
PollOutcome::Success { access_token } => return Ok(access_token),
|
||||
PollOutcome::Failed(reason) => return Err(DeviceFlowError::Unexpected(reason)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn parses_success() {
|
||||
let outcome =
|
||||
parse_poll_response(&json!({"access_token": "gho_abc", "token_type": "bearer"}));
|
||||
assert_eq!(
|
||||
outcome,
|
||||
PollOutcome::Success {
|
||||
access_token: "gho_abc".into()
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_pending_and_slow_down() {
|
||||
assert_eq!(
|
||||
parse_poll_response(&json!({"error": "authorization_pending"})),
|
||||
PollOutcome::Pending
|
||||
);
|
||||
assert_eq!(
|
||||
parse_poll_response(&json!({"error": "slow_down", "interval": 10})),
|
||||
PollOutcome::SlowDown
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_terminal_errors_with_description() {
|
||||
match parse_poll_response(
|
||||
&json!({"error": "expired_token", "error_description": "code expired"}),
|
||||
) {
|
||||
PollOutcome::Failed(msg) => assert_eq!(msg, "code expired"),
|
||||
other => panic!("expected Failed, got {other:?}"),
|
||||
}
|
||||
assert!(matches!(
|
||||
parse_poll_response(&json!({"error": "access_denied"})),
|
||||
PollOutcome::Failed(_)
|
||||
));
|
||||
assert!(matches!(
|
||||
parse_poll_response(&json!({})),
|
||||
PollOutcome::Failed(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slow_down_adds_five_seconds() {
|
||||
assert_eq!(bump_interval(5), 10);
|
||||
assert_eq!(bump_interval(10), 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn device_code_defaults_interval() {
|
||||
let dc: DeviceCode = serde_json::from_value(json!({
|
||||
"device_code": "d",
|
||||
"user_code": "WXYZ-1234",
|
||||
"verification_uri": "https://github.com/login/device"
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(dc.interval, 5);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
//! GitHub Copilot: OAuth device-flow login, token exchange, and a provider that multiplexes
|
||||
//! the chat/responses/anthropic codecs behind `api.githubcopilot.com`.
|
||||
|
||||
pub mod device_flow;
|
||||
pub mod provider;
|
||||
pub mod token;
|
||||
|
||||
pub use device_flow::{DeviceCode, PollOutcome};
|
||||
pub use provider::{CopilotCodec, CopilotModel, CopilotProvider};
|
||||
pub use token::{CopilotToken, TokenProvider};
|
||||
@@ -0,0 +1,296 @@
|
||||
//! Copilot provider: multiplexes all three wire codecs behind `api.githubcopilot.com`.
|
||||
//!
|
||||
//! Each model's `supported_endpoints` (from `GET /models`) decides which codec/endpoint to
|
||||
//! use. The token comes from [`super::token::TokenProvider`]. Network paths here are the
|
||||
//! flagged live-verification risk — the routing/header logic is unit-tested; end-to-end
|
||||
//! streaming must be checked against the live API.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use harness_core::llm::{
|
||||
Initiator, LlmEventStream, LlmRequest, Provider, ProviderError, WireContent,
|
||||
};
|
||||
use harness_core::types::ModelInfo;
|
||||
use serde_json::Value;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::codec::{anthropic, openai_chat, openai_responses};
|
||||
|
||||
use super::token::TokenProvider;
|
||||
|
||||
pub const DEFAULT_BASE_URL: &str = "https://api.githubcopilot.com";
|
||||
const API_VERSION: &str = "2026-06-01";
|
||||
const ANTHROPIC_BETA: &str = "interleaved-thinking-2025-05-14";
|
||||
const USER_AGENT: &str = concat!("ai-harness/", env!("CARGO_PKG_VERSION"));
|
||||
|
||||
/// Which wire codec a Copilot model speaks.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum CopilotCodec {
|
||||
Chat,
|
||||
Responses,
|
||||
Anthropic,
|
||||
}
|
||||
|
||||
impl CopilotCodec {
|
||||
fn path(self) -> &'static str {
|
||||
match self {
|
||||
CopilotCodec::Chat => "/chat/completions",
|
||||
CopilotCodec::Responses => "/responses",
|
||||
CopilotCodec::Anthropic => "/v1/messages",
|
||||
}
|
||||
}
|
||||
|
||||
fn build_request(self, req: &LlmRequest) -> Value {
|
||||
match self {
|
||||
CopilotCodec::Chat => openai_chat::build_request(req),
|
||||
CopilotCodec::Responses => openai_responses::build_request(req),
|
||||
CopilotCodec::Anthropic => anthropic::build_request(req),
|
||||
}
|
||||
}
|
||||
|
||||
fn decode(
|
||||
self,
|
||||
stream: impl futures::Stream<Item = reqwest::Result<bytes::Bytes>> + Send + 'static,
|
||||
) -> LlmEventStream {
|
||||
match self {
|
||||
CopilotCodec::Chat => openai_chat::decode(stream),
|
||||
CopilotCodec::Responses => openai_responses::decode(stream),
|
||||
CopilotCodec::Anthropic => anthropic::decode(stream),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A single model as advertised by `GET /models`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CopilotModel {
|
||||
pub id: String,
|
||||
pub codec: CopilotCodec,
|
||||
/// Only `model_picker_enabled` models are offered in the UI picker.
|
||||
pub picker_enabled: bool,
|
||||
}
|
||||
|
||||
/// Picks a codec from a model's `supported_endpoints`, preferring the responses API, then the
|
||||
/// Anthropic messages API, then chat completions.
|
||||
pub fn codec_for_endpoints(endpoints: &[String]) -> CopilotCodec {
|
||||
let has = |needle: &str| endpoints.iter().any(|e| e.contains(needle));
|
||||
if has("/responses") {
|
||||
CopilotCodec::Responses
|
||||
} else if has("/v1/messages") || has("/messages") {
|
||||
CopilotCodec::Anthropic
|
||||
} else {
|
||||
CopilotCodec::Chat
|
||||
}
|
||||
}
|
||||
|
||||
/// Parses a Copilot `GET /models` response body.
|
||||
pub fn parse_models(body: &Value) -> Vec<CopilotModel> {
|
||||
body["data"]
|
||||
.as_array()
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|m| {
|
||||
let id = m["id"].as_str()?.to_string();
|
||||
let endpoints: Vec<String> = m["supported_endpoints"]
|
||||
.as_array()
|
||||
.map(|a| {
|
||||
a.iter()
|
||||
.filter_map(|e| e.as_str().map(String::from))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
Some(CopilotModel {
|
||||
id,
|
||||
codec: codec_for_endpoints(&endpoints),
|
||||
picker_enabled: m["model_picker_enabled"].as_bool().unwrap_or(false),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn has_images(req: &LlmRequest) -> bool {
|
||||
req.messages
|
||||
.iter()
|
||||
.flat_map(|m| &m.content)
|
||||
.any(|c| matches!(c, WireContent::Image { .. }))
|
||||
}
|
||||
|
||||
fn initiator_header(initiator: Initiator) -> &'static str {
|
||||
match initiator {
|
||||
Initiator::User => "user",
|
||||
Initiator::Agent => "agent",
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CopilotProvider {
|
||||
tokens: TokenProvider,
|
||||
base_url: String,
|
||||
client: reqwest::Client,
|
||||
/// model id → codec, from `GET /models`. Unknown models default to chat completions.
|
||||
codecs: HashMap<String, CopilotCodec>,
|
||||
}
|
||||
|
||||
impl CopilotProvider {
|
||||
pub fn new(oauth_token: impl Into<String>, models: Vec<CopilotModel>) -> Self {
|
||||
Self::with_base_url(oauth_token, DEFAULT_BASE_URL, models)
|
||||
}
|
||||
|
||||
pub fn with_base_url(
|
||||
oauth_token: impl Into<String>,
|
||||
base_url: impl Into<String>,
|
||||
models: Vec<CopilotModel>,
|
||||
) -> Self {
|
||||
let client = reqwest::Client::new();
|
||||
let codecs = models.into_iter().map(|m| (m.id, m.codec)).collect();
|
||||
Self {
|
||||
tokens: TokenProvider::new(oauth_token, client.clone()),
|
||||
base_url: base_url.into(),
|
||||
client,
|
||||
codecs,
|
||||
}
|
||||
}
|
||||
|
||||
fn codec_for(&self, model: &str) -> CopilotCodec {
|
||||
self.codecs
|
||||
.get(model)
|
||||
.copied()
|
||||
.unwrap_or(CopilotCodec::Chat)
|
||||
}
|
||||
|
||||
fn classify_error(status: reqwest::StatusCode, body: String) -> ProviderError {
|
||||
match status.as_u16() {
|
||||
401 | 403 => ProviderError::Auth(body),
|
||||
429 => ProviderError::RateLimited { retry_after: None },
|
||||
s if (500..600).contains(&s) => ProviderError::Overloaded,
|
||||
s => ProviderError::Http { status: s, body },
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for CopilotProvider {
|
||||
fn id(&self) -> &str {
|
||||
"github-copilot"
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
||||
// Cost/limit metadata is layered on by models.dev; routing metadata lives in `codecs`.
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
req: LlmRequest,
|
||||
cancel: CancellationToken,
|
||||
) -> Result<LlmEventStream, ProviderError> {
|
||||
let codec = self.codec_for(&req.model);
|
||||
let body = codec.build_request(&req);
|
||||
let url = format!("{}{}", self.base_url, codec.path());
|
||||
|
||||
let now_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs() as i64;
|
||||
let token = self.tokens.token(now_secs).await;
|
||||
|
||||
let mut builder = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("authorization", format!("Bearer {}", token.token))
|
||||
.header("user-agent", USER_AGENT)
|
||||
.header("x-github-api-version", API_VERSION)
|
||||
.header("openai-intent", "conversation-edits")
|
||||
.header("x-initiator", initiator_header(req.initiator));
|
||||
if has_images(&req) {
|
||||
builder = builder.header("copilot-vision-request", "true");
|
||||
}
|
||||
if codec == CopilotCodec::Anthropic {
|
||||
builder = builder.header("anthropic-beta", ANTHROPIC_BETA);
|
||||
}
|
||||
|
||||
let send = builder.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 body = response.text().await.unwrap_or_default();
|
||||
return Err(Self::classify_error(status, body));
|
||||
}
|
||||
|
||||
Ok(codec.decode(response.bytes_stream()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn codec_routing_prefers_responses_then_anthropic_then_chat() {
|
||||
assert_eq!(
|
||||
codec_for_endpoints(&["/chat/completions".into(), "/responses".into()]),
|
||||
CopilotCodec::Responses
|
||||
);
|
||||
assert_eq!(
|
||||
codec_for_endpoints(&["/v1/messages".into()]),
|
||||
CopilotCodec::Anthropic
|
||||
);
|
||||
assert_eq!(
|
||||
codec_for_endpoints(&["/chat/completions".into()]),
|
||||
CopilotCodec::Chat
|
||||
);
|
||||
assert_eq!(codec_for_endpoints(&[]), CopilotCodec::Chat);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_models_with_endpoints_and_picker_flag() {
|
||||
let body = json!({
|
||||
"data": [
|
||||
{"id": "gpt-4o", "supported_endpoints": ["/chat/completions"], "model_picker_enabled": true},
|
||||
{"id": "claude-sonnet-4-5", "supported_endpoints": ["/v1/messages"], "model_picker_enabled": true},
|
||||
{"id": "o3", "supported_endpoints": ["/responses"], "model_picker_enabled": false},
|
||||
{"id": "no-endpoints"}
|
||||
]
|
||||
});
|
||||
let models = parse_models(&body);
|
||||
assert_eq!(models.len(), 4);
|
||||
assert_eq!(models[0].codec, CopilotCodec::Chat);
|
||||
assert!(models[0].picker_enabled);
|
||||
assert_eq!(models[1].codec, CopilotCodec::Anthropic);
|
||||
assert_eq!(models[2].codec, CopilotCodec::Responses);
|
||||
assert!(!models[2].picker_enabled);
|
||||
// Missing supported_endpoints → default chat, picker false.
|
||||
assert_eq!(models[3].codec, CopilotCodec::Chat);
|
||||
assert!(!models[3].picker_enabled);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_model_defaults_to_chat_codec() {
|
||||
let provider = CopilotProvider::new("oauth", vec![]);
|
||||
assert_eq!(provider.codec_for("whatever"), CopilotCodec::Chat);
|
||||
assert_eq!(provider.id(), "github-copilot");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn known_model_uses_mapped_codec() {
|
||||
let provider = CopilotProvider::new(
|
||||
"oauth",
|
||||
vec![CopilotModel {
|
||||
id: "claude-sonnet-4-5".into(),
|
||||
codec: CopilotCodec::Anthropic,
|
||||
picker_enabled: true,
|
||||
}],
|
||||
);
|
||||
assert_eq!(
|
||||
provider.codec_for("claude-sonnet-4-5"),
|
||||
CopilotCodec::Anthropic
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
//! Copilot API token strategy (flagged-risk area — verify against the live API).
|
||||
//!
|
||||
//! opencode sends the GitHub OAuth token directly as the Bearer (`expires: 0`). The classic
|
||||
//! Copilot API instead wants a short-lived token from `copilot_internal/v2/token`. We try the
|
||||
//! exchange first and cache it until shortly before expiry; if the endpoint rejects us, we fall
|
||||
//! back to the direct-Bearer behavior. Refresh is single-flight under a `tokio::Mutex`.
|
||||
|
||||
use serde::Deserialize;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
const EXCHANGE_URL: &str = "https://api.github.com/copilot_internal/v2/token";
|
||||
/// Refresh this many seconds before the reported expiry.
|
||||
const EXPIRY_SKEW_SECS: i64 = 120;
|
||||
const USER_AGENT: &str = concat!("ai-harness/", env!("CARGO_PKG_VERSION"));
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CopilotToken {
|
||||
pub token: String,
|
||||
/// Unix seconds; `0` means "never expires" (direct-Bearer fallback).
|
||||
pub expires_at: i64,
|
||||
}
|
||||
|
||||
impl CopilotToken {
|
||||
/// True when the token should be refreshed at `now_secs` (never, for `expires_at == 0`).
|
||||
pub fn needs_refresh(&self, now_secs: i64) -> bool {
|
||||
self.expires_at != 0 && now_secs + EXPIRY_SKEW_SECS >= self.expires_at
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum TokenError {
|
||||
#[error("network: {0}")]
|
||||
Network(String),
|
||||
/// The exchange endpoint is unavailable/unauthorized — the caller should fall back to the
|
||||
/// direct-Bearer strategy.
|
||||
#[error("exchange unsupported (status {0})")]
|
||||
ExchangeUnsupported(u16),
|
||||
#[error("unexpected response: {0}")]
|
||||
Unexpected(String),
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ExchangeResponse {
|
||||
token: String,
|
||||
#[serde(default)]
|
||||
expires_at: i64,
|
||||
}
|
||||
|
||||
/// Performs the token exchange once. `ExchangeUnsupported` signals the caller to fall back.
|
||||
pub async fn exchange(
|
||||
client: &reqwest::Client,
|
||||
oauth_token: &str,
|
||||
) -> Result<CopilotToken, TokenError> {
|
||||
let resp = client
|
||||
.get(EXCHANGE_URL)
|
||||
.header("authorization", format!("token {oauth_token}"))
|
||||
.header("accept", "application/json")
|
||||
.header("user-agent", USER_AGENT)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| TokenError::Network(e.to_string()))?;
|
||||
|
||||
let status = resp.status();
|
||||
if !status.is_success() {
|
||||
// 401/403/404 → this deployment doesn't support the exchange; fall back to direct Bearer.
|
||||
if matches!(status.as_u16(), 401 | 403 | 404) {
|
||||
return Err(TokenError::ExchangeUnsupported(status.as_u16()));
|
||||
}
|
||||
return Err(TokenError::Network(format!("status {}", status.as_u16())));
|
||||
}
|
||||
|
||||
let parsed: ExchangeResponse = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| TokenError::Unexpected(e.to_string()))?;
|
||||
Ok(CopilotToken {
|
||||
token: parsed.token,
|
||||
expires_at: parsed.expires_at,
|
||||
})
|
||||
}
|
||||
|
||||
/// Caches the exchanged Copilot token and refreshes it single-flight. Falls back to using the
|
||||
/// OAuth token directly (never-expiring) when the exchange endpoint rejects the request.
|
||||
pub struct TokenProvider {
|
||||
oauth_token: String,
|
||||
client: reqwest::Client,
|
||||
cached: Mutex<Option<CopilotToken>>,
|
||||
}
|
||||
|
||||
impl TokenProvider {
|
||||
pub fn new(oauth_token: impl Into<String>, client: reqwest::Client) -> Self {
|
||||
Self {
|
||||
oauth_token: oauth_token.into(),
|
||||
client,
|
||||
cached: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a usable Bearer token, refreshing/exchanging as needed. `now_secs` is injected
|
||||
/// for testability.
|
||||
pub async fn token(&self, now_secs: i64) -> CopilotToken {
|
||||
let mut guard = self.cached.lock().await;
|
||||
if let Some(tok) = guard.as_ref() {
|
||||
if !tok.needs_refresh(now_secs) {
|
||||
return tok.clone();
|
||||
}
|
||||
}
|
||||
let fresh = match exchange(&self.client, &self.oauth_token).await {
|
||||
Ok(tok) => tok,
|
||||
Err(_) => CopilotToken {
|
||||
// Direct-Bearer fallback: use the OAuth token itself, never expiring.
|
||||
token: self.oauth_token.clone(),
|
||||
expires_at: 0,
|
||||
},
|
||||
};
|
||||
*guard = Some(fresh.clone());
|
||||
fresh
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn never_expiring_token_does_not_refresh() {
|
||||
let tok = CopilotToken {
|
||||
token: "t".into(),
|
||||
expires_at: 0,
|
||||
};
|
||||
assert!(!tok.needs_refresh(i64::MAX));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refresh_triggers_within_skew_window() {
|
||||
let tok = CopilotToken {
|
||||
token: "t".into(),
|
||||
expires_at: 1_000,
|
||||
};
|
||||
assert!(!tok.needs_refresh(800)); // 800 + 120 < 1000
|
||||
assert!(tok.needs_refresh(881)); // 881 + 120 >= 1000
|
||||
assert!(tok.needs_refresh(1_000));
|
||||
}
|
||||
}
|
||||
@@ -1 +1,14 @@
|
||||
// Provider implementations (Copilot, OpenAI, Anthropic) land here in M1/M3.
|
||||
pub mod anthropic;
|
||||
pub mod auth;
|
||||
pub mod codec;
|
||||
pub mod copilot;
|
||||
pub mod modelsdev;
|
||||
pub mod openai;
|
||||
pub mod registry;
|
||||
|
||||
pub use anthropic::AnthropicProvider;
|
||||
pub use auth::{AuthRecord, AuthStorage};
|
||||
pub use copilot::CopilotProvider;
|
||||
pub use modelsdev::ModelCatalog;
|
||||
pub use openai::OpenAiProvider;
|
||||
pub use registry::ProviderRegistry;
|
||||
|
||||
@@ -0,0 +1,296 @@
|
||||
//! models.dev metadata: per-model context/output limits, pricing, and capability flags.
|
||||
//!
|
||||
//! At runtime we prefer a cached copy of `https://models.dev/api.json` (refreshed every 24h);
|
||||
//! if the cache is missing/stale and the network is unavailable, we fall back to a baked
|
||||
//! snapshot so cost display and limits still work offline.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use harness_core::types::{ModelCost, ModelInfo, ModelRef};
|
||||
use serde::Deserialize;
|
||||
|
||||
const API_URL: &str = "https://models.dev/api.json";
|
||||
const CACHE_TTL: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
/// Baked fallback, refreshed manually. Keeps cost/limits working with no cache and no network.
|
||||
const SNAPSHOT: &str = include_str!("../assets/models-snapshot.json");
|
||||
|
||||
/// models.dev top-level shape: `{ provider_id: { models: { model_id: {...} } } }`. Unknown
|
||||
/// keys (provider metadata, per-model extras) are ignored.
|
||||
#[derive(Deserialize)]
|
||||
struct ApiProvider {
|
||||
#[serde(default)]
|
||||
models: HashMap<String, ApiModel>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ApiModel {
|
||||
#[serde(default)]
|
||||
reasoning: bool,
|
||||
#[serde(default)]
|
||||
tool_call: bool,
|
||||
#[serde(default)]
|
||||
attachment: bool,
|
||||
#[serde(default)]
|
||||
cost: ApiCost,
|
||||
#[serde(default)]
|
||||
limit: ApiLimit,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct ApiCost {
|
||||
#[serde(default)]
|
||||
input: f64,
|
||||
#[serde(default)]
|
||||
output: f64,
|
||||
#[serde(default)]
|
||||
cache_read: f64,
|
||||
#[serde(default)]
|
||||
cache_write: f64,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Default)]
|
||||
struct ApiLimit {
|
||||
#[serde(default)]
|
||||
context: u64,
|
||||
#[serde(default)]
|
||||
output: u64,
|
||||
}
|
||||
|
||||
/// Parsed metadata keyed by `(provider_id, model_id)`.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ModelCatalog {
|
||||
models: HashMap<(String, String), ModelInfo>,
|
||||
}
|
||||
|
||||
impl ModelCatalog {
|
||||
/// Parses a models.dev `api.json` document.
|
||||
pub fn from_api_json(bytes: &[u8]) -> Result<Self, serde_json::Error> {
|
||||
let raw: HashMap<String, ApiProvider> = serde_json::from_slice(bytes)?;
|
||||
let mut models = HashMap::new();
|
||||
for (provider_id, provider) in raw {
|
||||
for (model_id, m) in provider.models {
|
||||
let info = ModelInfo {
|
||||
model: ModelRef::new(provider_id.clone(), model_id.clone()),
|
||||
context_limit: m.limit.context,
|
||||
output_limit: m.limit.output,
|
||||
cost: ModelCost {
|
||||
input: m.cost.input,
|
||||
output: m.cost.output,
|
||||
cache_read: m.cost.cache_read,
|
||||
cache_write: m.cost.cache_write,
|
||||
},
|
||||
reasoning: m.reasoning,
|
||||
tool_call: m.tool_call,
|
||||
attachment: m.attachment,
|
||||
};
|
||||
models.insert((provider_id.clone(), model_id), info);
|
||||
}
|
||||
}
|
||||
Ok(Self { models })
|
||||
}
|
||||
|
||||
/// The baked fallback snapshot — always available, never fails.
|
||||
pub fn baked() -> Self {
|
||||
Self::from_api_json(SNAPSHOT.as_bytes()).expect("baked models snapshot must parse")
|
||||
}
|
||||
|
||||
pub fn get(&self, provider: &str, model: &str) -> Option<&ModelInfo> {
|
||||
self.models.get(&(provider.to_string(), model.to_string()))
|
||||
}
|
||||
|
||||
/// Pricing for a model, or zero-cost if unknown.
|
||||
pub fn cost(&self, provider: &str, model: &str) -> ModelCost {
|
||||
self.get(provider, model)
|
||||
.map(|m| m.cost)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.models.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.models.is_empty()
|
||||
}
|
||||
|
||||
/// Default cache location: `~/.cache/ai-harness/models.json`.
|
||||
pub fn default_cache_path() -> PathBuf {
|
||||
dirs::cache_dir()
|
||||
.unwrap_or_else(std::env::temp_dir)
|
||||
.join("ai-harness")
|
||||
.join("models.json")
|
||||
}
|
||||
|
||||
/// Loads from the cache if present and younger than [`CACHE_TTL`]; otherwise returns the
|
||||
/// baked snapshot. Never touches the network — call [`ModelCatalog::refresh`] for that.
|
||||
pub fn load_cached_or_baked(cache_path: &Path) -> Self {
|
||||
if cache_fresh(cache_path, CACHE_TTL, SystemTime::now()) {
|
||||
if let Ok(bytes) = std::fs::read(cache_path) {
|
||||
if let Ok(catalog) = Self::from_api_json(&bytes) {
|
||||
if !catalog.is_empty() {
|
||||
return catalog;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Self::baked()
|
||||
}
|
||||
|
||||
/// Refreshes the default cache location with a fresh HTTP client. Convenience wrapper for
|
||||
/// callers (the app) that don't want to depend on `reqwest` directly.
|
||||
pub async fn refresh_default_cache() -> Result<Self, RefreshError> {
|
||||
let client = reqwest::Client::new();
|
||||
Self::refresh(&client, &Self::default_cache_path()).await
|
||||
}
|
||||
|
||||
/// Fetches the latest metadata and writes it to `cache_path` (best-effort). Returns the
|
||||
/// freshly-parsed catalog. Intended to run in the background so the *next* launch is current.
|
||||
pub async fn refresh(
|
||||
client: &reqwest::Client,
|
||||
cache_path: &Path,
|
||||
) -> Result<Self, RefreshError> {
|
||||
let bytes = client
|
||||
.get(API_URL)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| RefreshError::Network(e.to_string()))?
|
||||
.error_for_status()
|
||||
.map_err(|e| RefreshError::Network(e.to_string()))?
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| RefreshError::Network(e.to_string()))?;
|
||||
|
||||
let catalog =
|
||||
Self::from_api_json(&bytes).map_err(|e| RefreshError::Parse(e.to_string()))?;
|
||||
|
||||
if let Some(parent) = cache_path.parent() {
|
||||
let _ = std::fs::create_dir_all(parent);
|
||||
}
|
||||
if let Err(e) = std::fs::write(cache_path, &bytes) {
|
||||
tracing::warn!(error = %e, "failed to cache models.dev metadata");
|
||||
}
|
||||
Ok(catalog)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum RefreshError {
|
||||
#[error("network: {0}")]
|
||||
Network(String),
|
||||
#[error("parse: {0}")]
|
||||
Parse(String),
|
||||
}
|
||||
|
||||
/// True when `path` exists and was modified within `ttl` of `now`.
|
||||
fn cache_fresh(path: &Path, ttl: Duration, now: SystemTime) -> bool {
|
||||
let Ok(modified) = std::fs::metadata(path).and_then(|m| m.modified()) else {
|
||||
return false;
|
||||
};
|
||||
match now.duration_since(modified) {
|
||||
Ok(age) => age < ttl,
|
||||
Err(_) => true, // modified in the future (clock skew) — treat as fresh
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const FIXTURE: &str = r#"{
|
||||
"anthropic": {
|
||||
"id": "anthropic",
|
||||
"name": "Anthropic",
|
||||
"models": {
|
||||
"claude-sonnet-4-5": {
|
||||
"id": "claude-sonnet-4-5",
|
||||
"reasoning": true,
|
||||
"tool_call": true,
|
||||
"attachment": true,
|
||||
"cost": {"input": 3, "output": 15, "cache_read": 0.3, "cache_write": 3.75},
|
||||
"limit": {"context": 200000, "output": 64000}
|
||||
}
|
||||
}
|
||||
},
|
||||
"openai": {
|
||||
"id": "openai",
|
||||
"models": {
|
||||
"gpt-4o": {
|
||||
"id": "gpt-4o",
|
||||
"tool_call": true,
|
||||
"cost": {"input": 2.5, "output": 10},
|
||||
"limit": {"context": 128000, "output": 16384}
|
||||
}
|
||||
}
|
||||
}
|
||||
}"#;
|
||||
|
||||
#[test]
|
||||
fn parses_providers_and_models() {
|
||||
let catalog = ModelCatalog::from_api_json(FIXTURE.as_bytes()).unwrap();
|
||||
assert_eq!(catalog.len(), 2);
|
||||
|
||||
let sonnet = catalog.get("anthropic", "claude-sonnet-4-5").unwrap();
|
||||
assert_eq!(sonnet.context_limit, 200_000);
|
||||
assert_eq!(sonnet.output_limit, 64_000);
|
||||
assert_eq!(sonnet.cost.input, 3.0);
|
||||
assert_eq!(sonnet.cost.cache_write, 3.75);
|
||||
assert!(sonnet.reasoning && sonnet.tool_call && sonnet.attachment);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_cost_fields_default_to_zero() {
|
||||
let catalog = ModelCatalog::from_api_json(FIXTURE.as_bytes()).unwrap();
|
||||
let gpt = catalog.get("openai", "gpt-4o").unwrap();
|
||||
assert_eq!(gpt.cost.cache_read, 0.0);
|
||||
assert_eq!(gpt.cost.cache_write, 0.0);
|
||||
assert!(!gpt.reasoning); // absent → false
|
||||
assert!(gpt.tool_call);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cost_helper_is_zero_for_unknown_model() {
|
||||
let catalog = ModelCatalog::from_api_json(FIXTURE.as_bytes()).unwrap();
|
||||
assert_eq!(catalog.cost("openai", "nonexistent"), ModelCost::default());
|
||||
assert_eq!(catalog.cost("anthropic", "claude-sonnet-4-5").input, 3.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn baked_snapshot_parses_and_is_nonempty() {
|
||||
let catalog = ModelCatalog::baked();
|
||||
assert!(!catalog.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_freshness_respects_ttl() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("models.json");
|
||||
assert!(!cache_fresh(&path, CACHE_TTL, SystemTime::now())); // missing
|
||||
|
||||
std::fs::write(&path, "{}").unwrap();
|
||||
let now = SystemTime::now();
|
||||
assert!(cache_fresh(&path, CACHE_TTL, now));
|
||||
// A "now" far in the future makes the file look stale.
|
||||
let future = now + Duration::from_secs(48 * 60 * 60);
|
||||
assert!(!cache_fresh(&path, CACHE_TTL, future));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_cached_or_baked_falls_back_when_cache_absent() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("models.json");
|
||||
let catalog = ModelCatalog::load_cached_or_baked(&path);
|
||||
assert!(!catalog.is_empty()); // baked fallback
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn load_cached_or_baked_reads_fresh_cache() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("models.json");
|
||||
std::fs::write(&path, FIXTURE).unwrap();
|
||||
let catalog = ModelCatalog::load_cached_or_baked(&path);
|
||||
assert_eq!(catalog.len(), 2);
|
||||
assert!(catalog.get("anthropic", "claude-sonnet-4-5").is_some());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
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::{openai_chat, openai_responses};
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
|
||||
/// OpenCode Zen's OpenAI-compatible gateway (chat-completions only).
|
||||
const OPENCODE_BASE_URL: &str = "https://opencode.ai/zen/v1";
|
||||
|
||||
/// Which OpenAI wire format to speak for a given model.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum ApiFlavor {
|
||||
Responses,
|
||||
Chat,
|
||||
}
|
||||
|
||||
/// `gpt-*` and `o-*` models use the responses API; everything else (including
|
||||
/// OpenAI-compatible third-party endpoints) falls back to chat completions.
|
||||
fn flavor_for(model: &str) -> ApiFlavor {
|
||||
let is_native = model.starts_with("gpt-")
|
||||
|| model.starts_with("o1")
|
||||
|| model.starts_with("o3")
|
||||
|| model.starts_with("o4")
|
||||
|| model == "o1"
|
||||
|| model == "o3";
|
||||
if is_native {
|
||||
ApiFlavor::Responses
|
||||
} else {
|
||||
ApiFlavor::Chat
|
||||
}
|
||||
}
|
||||
|
||||
pub struct OpenAiProvider {
|
||||
id: String,
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
/// When set, always speak chat-completions regardless of model name — required for
|
||||
/// OpenAI-compatible gateways (e.g. OpenCode Zen) that don't implement `/responses`.
|
||||
chat_only: bool,
|
||||
client: reqwest::Client,
|
||||
}
|
||||
|
||||
impl OpenAiProvider {
|
||||
pub fn new(api_key: impl Into<String>) -> Self {
|
||||
Self {
|
||||
id: "openai".to_string(),
|
||||
api_key: api_key.into(),
|
||||
base_url: DEFAULT_BASE_URL.to_string(),
|
||||
chat_only: false,
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_base_url(api_key: impl Into<String>, base_url: impl Into<String>) -> Self {
|
||||
Self {
|
||||
id: "openai".to_string(),
|
||||
api_key: api_key.into(),
|
||||
base_url: base_url.into(),
|
||||
chat_only: false,
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// An OpenCode Zen provider: id `opencode`, chat-completions only, defaulting to Zen's
|
||||
/// gateway. Pass `base_url = None` to use the default endpoint.
|
||||
pub fn opencode(api_key: impl Into<String>, base_url: Option<String>) -> Self {
|
||||
Self {
|
||||
id: "opencode".to_string(),
|
||||
api_key: api_key.into(),
|
||||
base_url: base_url.unwrap_or_else(|| OPENCODE_BASE_URL.to_string()),
|
||||
chat_only: true,
|
||||
client: reqwest::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The wire format for `model`, honoring `chat_only`.
|
||||
fn flavor(&self, model: &str) -> ApiFlavor {
|
||||
if self.chat_only {
|
||||
ApiFlavor::Chat
|
||||
} else {
|
||||
flavor_for(model)
|
||||
}
|
||||
}
|
||||
|
||||
fn classify_error(
|
||||
status: reqwest::StatusCode,
|
||||
body: String,
|
||||
retry_after: Option<Duration>,
|
||||
) -> ProviderError {
|
||||
match status.as_u16() {
|
||||
400 if body.contains("context_length_exceeded")
|
||||
|| body.contains("maximum context length") =>
|
||||
{
|
||||
ProviderError::ContextOverflow
|
||||
}
|
||||
401 | 403 => ProviderError::Auth(body),
|
||||
429 => ProviderError::RateLimited { retry_after },
|
||||
s if (500..600).contains(&s) => ProviderError::Overloaded,
|
||||
s => ProviderError::Http { status: s, body },
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Provider for OpenAiProvider {
|
||||
fn id(&self) -> &str {
|
||||
&self.id
|
||||
}
|
||||
|
||||
async fn list_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
|
||||
// models.dev metadata is layered on top by the caller (see `modelsdev.rs`).
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
async fn stream(
|
||||
&self,
|
||||
req: LlmRequest,
|
||||
cancel: CancellationToken,
|
||||
) -> Result<LlmEventStream, ProviderError> {
|
||||
let (path, body) = match self.flavor(&req.model) {
|
||||
ApiFlavor::Responses => ("/responses", openai_responses::build_request(&req)),
|
||||
ApiFlavor::Chat => ("/chat/completions", openai_chat::build_request(&req)),
|
||||
};
|
||||
let url = format!("{}{}", self.base_url, path);
|
||||
|
||||
let send = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("authorization", format!("Bearer {}", self.api_key))
|
||||
.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(match self.flavor(&req.model) {
|
||||
ApiFlavor::Responses => openai_responses::decode(response.bytes_stream()),
|
||||
ApiFlavor::Chat => openai_chat::decode(response.bytes_stream()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn routes_gpt_and_o_series_to_responses() {
|
||||
assert_eq!(flavor_for("gpt-4o"), ApiFlavor::Responses);
|
||||
assert_eq!(flavor_for("gpt-5"), ApiFlavor::Responses);
|
||||
assert_eq!(flavor_for("o1"), ApiFlavor::Responses);
|
||||
assert_eq!(flavor_for("o3-mini"), ApiFlavor::Responses);
|
||||
assert_eq!(flavor_for("o4-mini"), ApiFlavor::Responses);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn routes_other_models_to_chat() {
|
||||
assert_eq!(flavor_for("llama-3.1-70b"), ApiFlavor::Chat);
|
||||
assert_eq!(flavor_for("deepseek-chat"), ApiFlavor::Chat);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn id_is_openai() {
|
||||
assert_eq!(OpenAiProvider::new("k").id(), "openai");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn opencode_provider_is_chat_only_and_ids_as_opencode() {
|
||||
let p = OpenAiProvider::opencode("k", None);
|
||||
assert_eq!(p.id(), "opencode");
|
||||
assert_eq!(p.base_url, OPENCODE_BASE_URL);
|
||||
// Even a gpt-*/o* model must route to chat completions on a chat-only gateway.
|
||||
assert_eq!(p.flavor("gpt-5.5"), ApiFlavor::Chat);
|
||||
assert_eq!(p.flavor("claude-sonnet-5"), ApiFlavor::Chat);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn opencode_honors_custom_base_url() {
|
||||
let p = OpenAiProvider::opencode("k", Some("https://example.test/v1".into()));
|
||||
assert_eq!(p.base_url, "https://example.test/v1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_context_overflow_from_400_body() {
|
||||
assert!(matches!(
|
||||
OpenAiProvider::classify_error(
|
||||
reqwest::StatusCode::BAD_REQUEST,
|
||||
"context_length_exceeded".into(),
|
||||
None,
|
||||
),
|
||||
ProviderError::ContextOverflow
|
||||
));
|
||||
assert!(matches!(
|
||||
OpenAiProvider::classify_error(
|
||||
reqwest::StatusCode::BAD_REQUEST,
|
||||
"some other error".into(),
|
||||
None,
|
||||
),
|
||||
ProviderError::Http { status: 400, .. }
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -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,218 @@
|
||||
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,
|
||||
spawner: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[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,73 @@
|
||||
//! Shared LSP-diagnostics reporting for the edit/write tools. After a file is written we ask
|
||||
//! the (optional) diagnostics source to re-analyze it and append any error-severity items to
|
||||
//! the tool output so the model sees mistakes it just introduced. Best-effort: no source, a
|
||||
//! slow server, or a timeout all just mean "no diagnostics" — never a tool failure.
|
||||
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
use harness_core::lsp::Severity;
|
||||
use harness_core::tool::{ToolCtx, ToolOutput};
|
||||
|
||||
/// docs/09-integrations.md: wait up to 1.5s for the server to (re)publish after the change.
|
||||
const DIAGNOSTICS_WAIT: Duration = Duration::from_millis(1500);
|
||||
|
||||
/// Touches `path` in the language server and appends error-severity diagnostics to `output`
|
||||
/// (both as a human-readable block in the text and the full set in metadata under `diagnostics`).
|
||||
pub async fn append_diagnostics(
|
||||
ctx: &ToolCtx,
|
||||
path: &Path,
|
||||
display_name: &str,
|
||||
output: &mut ToolOutput,
|
||||
) {
|
||||
let Some(source) = &ctx.diagnostics else {
|
||||
return;
|
||||
};
|
||||
source.touch(path).await;
|
||||
let diagnostics = source.diagnostics(path, DIAGNOSTICS_WAIT).await;
|
||||
if diagnostics.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let errors: Vec<_> = diagnostics
|
||||
.iter()
|
||||
.filter(|d| d.severity == Severity::Error)
|
||||
.collect();
|
||||
if !errors.is_empty() {
|
||||
output
|
||||
.output
|
||||
.push_str("\n\nLSP errors detected in this file, please fix:");
|
||||
for diag in &errors {
|
||||
output.output.push('\n');
|
||||
output.output.push_str(&diag.display_line(display_name));
|
||||
}
|
||||
}
|
||||
|
||||
// Full set (all severities) into metadata for the TUI.
|
||||
let items: Vec<serde_json::Value> = diagnostics
|
||||
.iter()
|
||||
.map(|d| {
|
||||
serde_json::json!({
|
||||
"line": d.line,
|
||||
"character": d.character,
|
||||
"severity": severity_str(d.severity),
|
||||
"message": d.message,
|
||||
"source": d.source,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
if let serde_json::Value::Object(map) = &mut output.metadata {
|
||||
map.insert("diagnostics".into(), serde_json::Value::Array(items));
|
||||
} else {
|
||||
output.metadata = serde_json::json!({ "diagnostics": items });
|
||||
}
|
||||
}
|
||||
|
||||
fn severity_str(severity: Severity) -> &'static str {
|
||||
match severity {
|
||||
Severity::Error => "error",
|
||||
Severity::Warning => "warning",
|
||||
Severity::Info => "info",
|
||||
Severity::Hint => "hint",
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,471 @@
|
||||
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.clone(), "Edit applied successfully.".to_string());
|
||||
output.metadata = serde_json::json!({"diff": diff, "added": added, "removed": removed});
|
||||
crate::diagnostics::append_diagnostics(&ctx, &path, &pattern, &mut output).await;
|
||||
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,
|
||||
spawner: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[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,184 @@
|
||||
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,
|
||||
spawner: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[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,212 @@
|
||||
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,
|
||||
spawner: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[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,48 @@
|
||||
// Built-in tools (bash, read, edit, write, glob, grep, todo, webfetch, task) land here in M1/M4.
|
||||
mod bash;
|
||||
mod diagnostics;
|
||||
mod edit;
|
||||
mod glob;
|
||||
mod grep;
|
||||
mod paths;
|
||||
mod read;
|
||||
mod skill;
|
||||
mod task;
|
||||
mod write;
|
||||
|
||||
pub use bash::BashTool;
|
||||
pub use edit::EditTool;
|
||||
pub use glob::GlobTool;
|
||||
pub use grep::GrepTool;
|
||||
pub use read::ReadTool;
|
||||
pub use skill::SkillTool;
|
||||
pub use task::TaskTool;
|
||||
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));
|
||||
}
|
||||
|
||||
/// Registers the multiagent `task` tool (M4). Kept separate from [`register_builtins`] so
|
||||
/// non-orchestrating contexts can omit delegation; the tool no-ops with an error if the
|
||||
/// session has no spawner wired in.
|
||||
pub fn register_task_tool(registry: &mut ToolRegistry) {
|
||||
registry.register(Arc::new(TaskTool));
|
||||
}
|
||||
|
||||
/// Registers the `skill` tool (M5) over a loaded skill set. No-op when there are no skills,
|
||||
/// so the tool is only advertised when something can be loaded.
|
||||
pub fn register_skill_tool(registry: &mut ToolRegistry, skills: &[harness_core::config::SkillDef]) {
|
||||
if !skills.is_empty() {
|
||||
registry.register(Arc::new(SkillTool::new(skills)));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,212 @@
|
||||
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(),
|
||||
));
|
||||
}
|
||||
|
||||
// In a subagent session, advertise this read on the job board.
|
||||
if let Some(reporter) = &ctx.context_reporter {
|
||||
let reported = paths::relative_pattern(&ctx.cwd, &path);
|
||||
reporter.report_file(reported, numbered.len() as u32).await;
|
||||
}
|
||||
|
||||
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,
|
||||
spawner: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[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,154 @@
|
||||
//! The `skill` tool (M5): the system prompt advertises each skill's name + description; when
|
||||
//! the model decides a skill is relevant it calls this tool with the skill name to pull the
|
||||
//! full instructions on demand. See `docs/09-integrations.md`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use harness_core::config::SkillDef;
|
||||
use harness_core::tool::{Tool, ToolCtx, ToolError, ToolOutput};
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
struct SkillParams {
|
||||
/// The name of the skill to load, as advertised in the system prompt.
|
||||
name: String,
|
||||
}
|
||||
|
||||
/// Serves skill bodies by name. Built from the loaded skill set; if there are no skills the
|
||||
/// caller simply doesn't register the tool.
|
||||
pub struct SkillTool {
|
||||
/// name → (description, body).
|
||||
skills: HashMap<String, (String, String)>,
|
||||
/// Sorted names, for a stable "unknown skill" hint.
|
||||
names: Vec<String>,
|
||||
description: String,
|
||||
}
|
||||
|
||||
impl SkillTool {
|
||||
pub fn new(skills: &[SkillDef]) -> Self {
|
||||
let mut names: Vec<String> = skills.iter().map(|s| s.name.clone()).collect();
|
||||
names.sort();
|
||||
let map = skills
|
||||
.iter()
|
||||
.map(|s| (s.name.clone(), (s.description.clone(), s.body.clone())))
|
||||
.collect();
|
||||
let description = format!(
|
||||
"Load the full instructions for a named skill before doing the related work. \
|
||||
Available skills: {}.",
|
||||
names.join(", ")
|
||||
);
|
||||
Self {
|
||||
skills: map,
|
||||
names,
|
||||
description,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SkillTool {
|
||||
fn name(&self) -> &str {
|
||||
"skill"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
&self.description
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
serde_json::to_value(schemars::schema_for!(SkillParams)).unwrap()
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
_ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let params: SkillParams =
|
||||
serde_json::from_value(input).map_err(|e| ToolError::Invalid(e.to_string()))?;
|
||||
match self.skills.get(¶ms.name) {
|
||||
Some((_description, body)) => Ok(ToolOutput::new(
|
||||
format!("skill: {}", params.name),
|
||||
body.clone(),
|
||||
)),
|
||||
None => Err(ToolError::Invalid(format!(
|
||||
"unknown skill {:?}; available: {}",
|
||||
params.name,
|
||||
self.names.join(", ")
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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() -> 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: "c1".into(),
|
||||
data_dir: std::env::temp_dir(),
|
||||
cwd: std::env::temp_dir(),
|
||||
cancel: CancellationToken::new(),
|
||||
ask: PermissionHandle::new(
|
||||
service,
|
||||
SessionId::new(),
|
||||
Vec::new(),
|
||||
Arc::new(Mutex::new(Vec::new())),
|
||||
CancellationToken::new(),
|
||||
),
|
||||
metadata,
|
||||
spawner: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample() -> Vec<SkillDef> {
|
||||
vec![SkillDef {
|
||||
name: "formatter".into(),
|
||||
description: "format code".into(),
|
||||
body: "Run cargo fmt.".into(),
|
||||
}]
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn returns_skill_body_by_name() {
|
||||
let tool = SkillTool::new(&sample());
|
||||
let out = tool
|
||||
.execute(serde_json::json!({"name": "formatter"}), ctx())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(out.output, "Run cargo fmt.");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_skill_is_an_input_error() {
|
||||
let tool = SkillTool::new(&sample());
|
||||
let err = tool
|
||||
.execute(serde_json::json!({"name": "nope"}), ctx())
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, ToolError::Invalid(_)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn description_lists_available_skills() {
|
||||
let tool = SkillTool::new(&sample());
|
||||
assert!(tool.description().contains("formatter"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
//! The `task` tool: delegate work to a specialist subagent, foreground or background.
|
||||
//!
|
||||
//! This tool is deliberately thin — it validates input, gates on a `task/<agent>` permission,
|
||||
//! and hands off to the engine's `SubagentSpawner` (owned by the composition root), which
|
||||
//! resolves the agent, enforces the depth limit, applies permission intersection, and runs
|
||||
//! the child session. See `docs/04-multiagent.md`.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
|
||||
use harness_core::tool::{
|
||||
invalid_input, SpawnError, SpawnRequest, Tool, ToolCtx, ToolError, ToolOutput,
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct TaskInput {
|
||||
/// Short human-facing label for the subtask (shown on the job board).
|
||||
description: String,
|
||||
/// The full instruction handed to the subagent.
|
||||
prompt: String,
|
||||
/// Which specialist to run (must be a subagent-capable agent).
|
||||
subagent_type: String,
|
||||
/// Alias or task id of a completed job to continue instead of starting fresh.
|
||||
#[serde(default)]
|
||||
task_id: Option<String>,
|
||||
/// Run in the background and return immediately (tracked on the job board).
|
||||
#[serde(default)]
|
||||
background: bool,
|
||||
}
|
||||
|
||||
pub struct TaskTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskTool {
|
||||
fn name(&self) -> &str {
|
||||
"task"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Delegate a self-contained unit of work to a specialist subagent. Set `background: \
|
||||
true` to launch it without blocking (track progress on the job board); reuse a \
|
||||
completed subagent by passing its `task_id`/alias to continue the same session."
|
||||
}
|
||||
|
||||
fn parameters(&self) -> serde_json::Value {
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"description": { "type": "string", "description": "Short label for the subtask." },
|
||||
"prompt": { "type": "string", "description": "Full instruction for the subagent." },
|
||||
"subagent_type": { "type": "string", "description": "Specialist to run." },
|
||||
"task_id": { "type": "string", "description": "Alias/id of a completed job to reuse." },
|
||||
"background": { "type": "boolean", "description": "Launch without blocking." }
|
||||
},
|
||||
"required": ["description", "prompt", "subagent_type"]
|
||||
})
|
||||
}
|
||||
|
||||
async fn execute(
|
||||
&self,
|
||||
input: serde_json::Value,
|
||||
ctx: ToolCtx,
|
||||
) -> Result<ToolOutput, ToolError> {
|
||||
let args: TaskInput = serde_json::from_value(input).map_err(|e| invalid_input(self, e))?;
|
||||
|
||||
let Some(spawner) = ctx.spawner.clone() else {
|
||||
return Err(ToolError::Other(
|
||||
"subagent delegation is not available in this session".into(),
|
||||
));
|
||||
};
|
||||
|
||||
// Gate on task/<agent>. `Always` grants blanket delegation to this specialist.
|
||||
ctx.ask
|
||||
.ask(
|
||||
"task",
|
||||
&args.subagent_type,
|
||||
&args.subagent_type,
|
||||
json!({
|
||||
"agent": args.subagent_type,
|
||||
"description": args.description,
|
||||
"background": args.background,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
let req = SpawnRequest {
|
||||
parent_session_id: ctx.session_id.clone(),
|
||||
parent_message_id: ctx.message_id.clone(),
|
||||
agent: args.subagent_type.clone(),
|
||||
description: args.description.clone(),
|
||||
prompt: args.prompt,
|
||||
reuse_task_id: args.task_id,
|
||||
background: args.background,
|
||||
cancel: ctx.cancel.clone(),
|
||||
};
|
||||
|
||||
let outcome = spawner.spawn(req).await.map_err(map_spawn_error)?;
|
||||
|
||||
if outcome.background {
|
||||
let alias = outcome.alias.unwrap_or_default();
|
||||
Ok(ToolOutput {
|
||||
title: format!("launched {} ({alias})", args.subagent_type),
|
||||
output: format!(
|
||||
"Launched background task {alias} ({}). Check the job board; do not poll — \
|
||||
wait for completion.",
|
||||
outcome.child_session_id
|
||||
),
|
||||
metadata: json!({
|
||||
"child_session": outcome.child_session_id,
|
||||
"agent": args.subagent_type,
|
||||
"alias": alias,
|
||||
"background": true,
|
||||
}),
|
||||
})
|
||||
} else {
|
||||
let text = outcome.final_text.unwrap_or_default();
|
||||
Ok(ToolOutput {
|
||||
title: format!("{} — {}", args.subagent_type, args.description),
|
||||
output: text,
|
||||
metadata: json!({
|
||||
"child_session": outcome.child_session_id,
|
||||
"agent": args.subagent_type,
|
||||
"background": false,
|
||||
}),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn map_spawn_error(err: SpawnError) -> ToolError {
|
||||
match err {
|
||||
// Depth/agent problems are the model's to fix — surface as tool errors it can read
|
||||
// and act on, not hard failures.
|
||||
SpawnError::DepthExceeded => ToolError::Other(err.to_string()),
|
||||
SpawnError::InvalidAgent(_) => ToolError::Invalid(err.to_string()),
|
||||
SpawnError::ReuseNotFound(_) => ToolError::Invalid(err.to_string()),
|
||||
SpawnError::Other(msg) => ToolError::Other(msg),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
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});
|
||||
crate::diagnostics::append_diagnostics(&ctx, &path, ¶ms.file_path, &mut output).await;
|
||||
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,
|
||||
spawner: None,
|
||||
context_reporter: None,
|
||||
diagnostics: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[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,23 @@ path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
harness-app = { workspace = true }
|
||||
harness-core = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
ratatui = { workspace = true }
|
||||
crossterm = { workspace = true }
|
||||
tui-textarea = { workspace = true }
|
||||
pulldown-cmark = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tracing-subscriber = { workspace = true }
|
||||
tracing-appender = { workspace = true }
|
||||
dirs = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
tokio-util = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
insta = "1"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
use std::io;
|
||||
use std::path::PathBuf;
|
||||
|
||||
use crossterm::event::EventStream;
|
||||
use futures::StreamExt;
|
||||
use harness_app::EngineHandle;
|
||||
use harness_core::event::AppEvent;
|
||||
use ratatui::backend::CrosstermBackend;
|
||||
use ratatui::Terminal;
|
||||
use tokio::sync::broadcast;
|
||||
use tokio::time::{interval, Duration};
|
||||
|
||||
use crate::input::{apply_action, handle_event};
|
||||
use crate::render::render;
|
||||
use crate::state::AppState;
|
||||
use crate::terminal::TerminalGuard;
|
||||
|
||||
const DEFAULT_MODEL: &str = "anthropic/claude-sonnet-4-5";
|
||||
const DEFAULT_AGENT: &str = "orchestrator";
|
||||
const RENDER_TICK_MS: u64 = 33;
|
||||
|
||||
pub struct App {
|
||||
engine: EngineHandle,
|
||||
state: AppState,
|
||||
_terminal_guard: TerminalGuard,
|
||||
terminal: Terminal<CrosstermBackend<io::Stdout>>,
|
||||
events: EventStream,
|
||||
bus_rx: broadcast::Receiver<AppEvent>,
|
||||
render_interval: tokio::time::Interval,
|
||||
}
|
||||
|
||||
impl App {
|
||||
pub async fn new(cwd: PathBuf) -> anyhow::Result<Self> {
|
||||
let engine = EngineHandle::init(cwd).await?;
|
||||
let bus_rx = engine.bus().subscribe();
|
||||
|
||||
let config = engine.config();
|
||||
let model_ref = config
|
||||
.model
|
||||
.clone()
|
||||
.unwrap_or_else(|| DEFAULT_MODEL.to_string());
|
||||
let agent = DEFAULT_AGENT.to_string();
|
||||
|
||||
let mut state = AppState::new(model_ref, agent);
|
||||
|
||||
// Create a default session so the user can start typing immediately.
|
||||
match engine.new_session(&state.agent, &state.model_ref).await {
|
||||
Ok(session_id) => {
|
||||
state.session_id = Some(session_id);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "failed to create default session");
|
||||
}
|
||||
}
|
||||
|
||||
let terminal_guard = TerminalGuard::enter()?;
|
||||
let backend = CrosstermBackend::new(io::stdout());
|
||||
let terminal = Terminal::new(backend)?;
|
||||
let events = EventStream::new();
|
||||
let render_interval = interval(Duration::from_millis(RENDER_TICK_MS));
|
||||
|
||||
Ok(Self {
|
||||
engine,
|
||||
state,
|
||||
_terminal_guard: terminal_guard,
|
||||
terminal,
|
||||
events,
|
||||
bus_rx,
|
||||
render_interval,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn run(&mut self) -> anyhow::Result<()> {
|
||||
self.state.dirty = true;
|
||||
while !self.state.quit {
|
||||
tokio::select! {
|
||||
biased;
|
||||
|
||||
maybe_event = self.events.next() => {
|
||||
match maybe_event {
|
||||
Some(Ok(event)) => {
|
||||
let action = handle_event(event, &mut self.state, &self.engine);
|
||||
apply_action(action, &mut self.state, &self.engine).await;
|
||||
}
|
||||
Some(Err(e)) => {
|
||||
tracing::error!(error = %e, "input error");
|
||||
}
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
|
||||
Ok(event) = self.bus_rx.recv() => {
|
||||
self.state.apply_event(event);
|
||||
}
|
||||
|
||||
_ = self.render_interval.tick() => {
|
||||
if self.state.dirty {
|
||||
self.terminal.draw(|frame| render(frame, &mut self.state))?;
|
||||
self.state.dirty = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,388 @@
|
||||
use crossterm::event::{Event, KeyCode, KeyEvent, KeyModifiers};
|
||||
use harness_app::EngineHandle;
|
||||
use harness_core::permission::PermissionReply;
|
||||
|
||||
use crate::modal;
|
||||
use crate::state::{AppState, ModalState};
|
||||
|
||||
/// Possible high-level actions produced by key input.
|
||||
#[derive(Debug)]
|
||||
pub enum InputAction {
|
||||
None,
|
||||
Submit {
|
||||
text: String,
|
||||
},
|
||||
/// A user-defined slash command: `name` (no leading `/`), its `args`, and the `raw` input
|
||||
/// to fall back to submitting verbatim if no such command is defined.
|
||||
RunCommand {
|
||||
name: String,
|
||||
args: String,
|
||||
raw: String,
|
||||
},
|
||||
Abort,
|
||||
Quit,
|
||||
LoadSessions,
|
||||
NewSession,
|
||||
SetModel(String),
|
||||
SetAgent(String),
|
||||
CloseModal,
|
||||
ScrollUp(u16),
|
||||
ScrollDown(u16),
|
||||
PermissionReply(PermissionReply),
|
||||
SessionPickerUp,
|
||||
SessionPickerDown,
|
||||
SessionPickerSelect,
|
||||
OpenJobs,
|
||||
JobsUp,
|
||||
JobsDown,
|
||||
JobsDrillIn,
|
||||
}
|
||||
|
||||
/// Translate a crossterm event into an action and/or mutate `state` directly.
|
||||
pub fn handle_event(event: Event, state: &mut AppState, _engine: &EngineHandle) -> InputAction {
|
||||
match event {
|
||||
Event::Key(key) => {
|
||||
let action = handle_key(key, state);
|
||||
// Any keypress can change the input buffer or cursor; force a redraw so typed
|
||||
// characters appear immediately rather than only when some other event sets dirty.
|
||||
state.dirty = true;
|
||||
action
|
||||
}
|
||||
Event::Resize(_, _) => {
|
||||
state.dirty = true;
|
||||
InputAction::None
|
||||
}
|
||||
Event::Mouse(_) | Event::FocusGained | Event::FocusLost | Event::Paste(_) => {
|
||||
InputAction::None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_key(key: KeyEvent, state: &mut AppState) -> InputAction {
|
||||
match &state.modal {
|
||||
ModalState::Permission { .. } => handle_permission_key(key, state),
|
||||
ModalState::SessionPicker { .. } => handle_session_picker_key(key, state),
|
||||
ModalState::JobsPane { .. } => handle_jobs_pane_key(key, state),
|
||||
ModalState::None => handle_normal_key(key, state),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_jobs_pane_key(key: KeyEvent, _state: &mut AppState) -> InputAction {
|
||||
match key.code {
|
||||
KeyCode::Up => InputAction::JobsUp,
|
||||
KeyCode::Down => InputAction::JobsDown,
|
||||
KeyCode::Enter => InputAction::JobsDrillIn,
|
||||
KeyCode::Esc => InputAction::CloseModal,
|
||||
_ => InputAction::None,
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_permission_key(key: KeyEvent, _state: &mut AppState) -> InputAction {
|
||||
match key.code {
|
||||
KeyCode::Char('y') | KeyCode::Char('Y') => {
|
||||
InputAction::PermissionReply(PermissionReply::Once)
|
||||
}
|
||||
KeyCode::Char('a') | KeyCode::Char('A') => {
|
||||
InputAction::PermissionReply(PermissionReply::Always)
|
||||
}
|
||||
KeyCode::Char('n') | KeyCode::Char('N') => {
|
||||
InputAction::PermissionReply(PermissionReply::Reject)
|
||||
}
|
||||
// Esc rejects the pending request rather than merely hiding the modal — closing
|
||||
// without a reply would leave the tool call blocked on its oneshot forever.
|
||||
KeyCode::Esc => InputAction::PermissionReply(PermissionReply::Reject),
|
||||
_ => InputAction::None,
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_session_picker_key(key: KeyEvent, _state: &mut AppState) -> InputAction {
|
||||
match key.code {
|
||||
KeyCode::Up => InputAction::SessionPickerUp,
|
||||
KeyCode::Down => InputAction::SessionPickerDown,
|
||||
KeyCode::Enter => InputAction::SessionPickerSelect,
|
||||
KeyCode::Esc => InputAction::CloseModal,
|
||||
_ => InputAction::None,
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_normal_key(key: KeyEvent, state: &mut AppState) -> InputAction {
|
||||
let ctrl = key.modifiers.contains(KeyModifiers::CONTROL);
|
||||
|
||||
// Reset the double-Ctrl+C guard on any key that isn't another Ctrl+C.
|
||||
if !(matches!(key.code, KeyCode::Char('c') | KeyCode::Char('C')) && ctrl) {
|
||||
state.ctrl_c_pressed = false;
|
||||
}
|
||||
|
||||
match key.code {
|
||||
KeyCode::Enter => {
|
||||
let text = state.input.lines().join("\n");
|
||||
state.input = tui_textarea::TextArea::default();
|
||||
state
|
||||
.input
|
||||
.set_cursor_line_style(ratatui::style::Style::default());
|
||||
if let Some(action) = parse_slash_command(&text) {
|
||||
action
|
||||
} else {
|
||||
InputAction::Submit { text }
|
||||
}
|
||||
}
|
||||
KeyCode::Char('c') if ctrl => {
|
||||
if state.ctrl_c_pressed {
|
||||
InputAction::Quit
|
||||
} else {
|
||||
state.ctrl_c_pressed = true;
|
||||
state.dirty = true;
|
||||
InputAction::None
|
||||
}
|
||||
}
|
||||
KeyCode::Esc => {
|
||||
if state.running {
|
||||
InputAction::Abort
|
||||
} else {
|
||||
InputAction::None
|
||||
}
|
||||
}
|
||||
KeyCode::Char('s') if ctrl => InputAction::LoadSessions,
|
||||
KeyCode::Char('j') if ctrl => InputAction::OpenJobs,
|
||||
KeyCode::Up => {
|
||||
let (row, _) = state.input.cursor();
|
||||
if row == 0 {
|
||||
InputAction::ScrollUp(3)
|
||||
} else {
|
||||
state.input.input(key);
|
||||
InputAction::None
|
||||
}
|
||||
}
|
||||
KeyCode::Down => {
|
||||
let (row, _) = state.input.cursor();
|
||||
let last_line = state.input.lines().len().saturating_sub(1);
|
||||
if row == last_line {
|
||||
InputAction::ScrollDown(3)
|
||||
} else {
|
||||
state.input.input(key);
|
||||
InputAction::None
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
state.input.input(key);
|
||||
InputAction::None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_slash_command(text: &str) -> Option<InputAction> {
|
||||
let trimmed = text.trim();
|
||||
if !trimmed.starts_with('/') {
|
||||
return None;
|
||||
}
|
||||
let mut parts = trimmed.split_whitespace();
|
||||
let cmd = parts.next()?;
|
||||
let rest: String = parts.collect::<Vec<_>>().join(" ");
|
||||
match cmd {
|
||||
"/new" => Some(InputAction::NewSession),
|
||||
"/model" if !rest.is_empty() => Some(InputAction::SetModel(rest)),
|
||||
"/agent" if !rest.is_empty() => Some(InputAction::SetAgent(rest)),
|
||||
"/sessions" => Some(InputAction::LoadSessions),
|
||||
"/jobs" => Some(InputAction::OpenJobs),
|
||||
"/quit" => Some(InputAction::Quit),
|
||||
// Any other `/word` is treated as a user-defined command, resolved against the engine's
|
||||
// loaded commands when applied; if none matches, the raw text is submitted as-is.
|
||||
other => Some(InputAction::RunCommand {
|
||||
name: other.trim_start_matches('/').to_string(),
|
||||
args: rest,
|
||||
raw: trimmed.to_string(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Apply an action that needs async engine calls. Non-async decisions are applied inline.
|
||||
pub async fn apply_action(action: InputAction, state: &mut AppState, engine: &EngineHandle) {
|
||||
match action {
|
||||
InputAction::Submit { text } => {
|
||||
if let Some(session_id) = state.session_id.clone() {
|
||||
if let Err(e) = engine.prompt(session_id, text, &state.model_ref).await {
|
||||
tracing::error!(error = %e, "prompt failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
InputAction::RunCommand { name, args, raw } => {
|
||||
let Some(session_id) = state.session_id.clone() else {
|
||||
return;
|
||||
};
|
||||
match engine.command(&name) {
|
||||
Some(cmd) => {
|
||||
let text = cmd.expand(&args);
|
||||
// A command may switch model/agent for this one run only.
|
||||
let model_ref = cmd.model.clone().unwrap_or_else(|| state.model_ref.clone());
|
||||
if let Err(e) = engine
|
||||
.prompt_with(session_id, text, &model_ref, cmd.agent.as_deref())
|
||||
.await
|
||||
{
|
||||
tracing::error!(error = %e, "command prompt failed");
|
||||
}
|
||||
}
|
||||
// Unknown command: submit the original text as an ordinary message.
|
||||
None => {
|
||||
if let Err(e) = engine.prompt(session_id, raw, &state.model_ref).await {
|
||||
tracing::error!(error = %e, "prompt failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
InputAction::Abort => {
|
||||
if let Some(session_id) = state.session_id.clone() {
|
||||
engine.abort(&session_id);
|
||||
}
|
||||
}
|
||||
InputAction::LoadSessions => match engine.list_sessions().await {
|
||||
Ok(sessions) => modal::open_session_picker(state, sessions),
|
||||
Err(e) => tracing::error!(error = %e, "failed to list sessions"),
|
||||
},
|
||||
InputAction::NewSession => match engine.new_session(&state.agent, &state.model_ref).await {
|
||||
Ok(id) => {
|
||||
state.session_id = Some(id);
|
||||
state.messages.clear();
|
||||
state.scroll_offset = 0;
|
||||
state.session_cost = 0.0;
|
||||
state.session_tokens = 0;
|
||||
state.dirty = true;
|
||||
}
|
||||
Err(e) => tracing::error!(error = %e, "failed to create session"),
|
||||
},
|
||||
InputAction::SetModel(model_ref) => {
|
||||
state.model_ref = model_ref;
|
||||
state.dirty = true;
|
||||
}
|
||||
InputAction::SetAgent(agent) => {
|
||||
state.agent = agent;
|
||||
state.dirty = true;
|
||||
}
|
||||
InputAction::Quit => state.quit = true,
|
||||
InputAction::CloseModal => state.close_modal(),
|
||||
InputAction::ScrollUp(n) => state.scroll_up(n),
|
||||
InputAction::ScrollDown(n) => state.scroll_down(n),
|
||||
InputAction::PermissionReply(reply) => {
|
||||
modal::resolve_permission(state, engine, reply);
|
||||
}
|
||||
InputAction::SessionPickerUp => {
|
||||
if let ModalState::SessionPicker {
|
||||
selected,
|
||||
sessions: _,
|
||||
} = &mut state.modal
|
||||
{
|
||||
*selected = selected.saturating_sub(1);
|
||||
}
|
||||
state.dirty = true;
|
||||
}
|
||||
InputAction::SessionPickerDown => {
|
||||
if let ModalState::SessionPicker { selected, sessions } = &mut state.modal {
|
||||
if *selected + 1 < sessions.len() {
|
||||
*selected += 1;
|
||||
}
|
||||
}
|
||||
state.dirty = true;
|
||||
}
|
||||
InputAction::SessionPickerSelect => {
|
||||
if let ModalState::SessionPicker { sessions, selected } = &state.modal {
|
||||
if let Some(session) = sessions.get(*selected).cloned() {
|
||||
if let Err(e) = modal::select_session(state, engine, session).await {
|
||||
tracing::error!(error = %e, "failed to load session");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
InputAction::OpenJobs => {
|
||||
if let Some(session_id) = state.session_id.clone() {
|
||||
match engine.jobs(session_id).await {
|
||||
Ok(jobs) => {
|
||||
state.set_jobs(jobs);
|
||||
state.open_jobs_pane();
|
||||
}
|
||||
Err(e) => tracing::error!(error = %e, "failed to load jobs"),
|
||||
}
|
||||
}
|
||||
}
|
||||
InputAction::JobsUp => {
|
||||
if let ModalState::JobsPane { selected } = &mut state.modal {
|
||||
*selected = selected.saturating_sub(1);
|
||||
}
|
||||
state.dirty = true;
|
||||
}
|
||||
InputAction::JobsDown => {
|
||||
let job_count = state.jobs.len();
|
||||
if let ModalState::JobsPane { selected } = &mut state.modal {
|
||||
if *selected + 1 < job_count {
|
||||
*selected += 1;
|
||||
}
|
||||
}
|
||||
state.dirty = true;
|
||||
}
|
||||
InputAction::JobsDrillIn => {
|
||||
if let Some(child) = state.selected_job_child() {
|
||||
match engine.get_session(child).await {
|
||||
Ok(Some(session)) => {
|
||||
if let Err(e) = modal::select_session(state, engine, session).await {
|
||||
tracing::error!(error = %e, "failed to open subtask session");
|
||||
}
|
||||
}
|
||||
Ok(None) => tracing::warn!("subtask session no longer exists"),
|
||||
Err(e) => tracing::error!(error = %e, "failed to load subtask session"),
|
||||
}
|
||||
}
|
||||
}
|
||||
InputAction::None => {}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::state::AppState;
|
||||
use crossterm::event::KeyEvent;
|
||||
|
||||
#[tokio::test]
|
||||
async fn typing_marks_frame_dirty_and_updates_input() {
|
||||
let engine = EngineHandle::init_in_memory(std::env::temp_dir()).unwrap();
|
||||
let mut state = AppState::new("anthropic/claude".into(), "orchestrator".into());
|
||||
state.dirty = false;
|
||||
|
||||
let action = handle_event(
|
||||
Event::Key(KeyEvent::from(KeyCode::Char('x'))),
|
||||
&mut state,
|
||||
&engine,
|
||||
);
|
||||
|
||||
assert!(matches!(action, InputAction::None));
|
||||
assert!(state.dirty, "a keystroke must request a redraw");
|
||||
assert_eq!(state.input.lines().join("\n"), "x");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builtin_slash_commands_still_parse() {
|
||||
assert!(matches!(
|
||||
parse_slash_command("/new"),
|
||||
Some(InputAction::NewSession)
|
||||
));
|
||||
assert!(matches!(
|
||||
parse_slash_command("/model openai/gpt-5"),
|
||||
Some(InputAction::SetModel(m)) if m == "openai/gpt-5"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_slash_becomes_a_run_command() {
|
||||
match parse_slash_command("/deploy prod now") {
|
||||
Some(InputAction::RunCommand { name, args, raw }) => {
|
||||
assert_eq!(name, "deploy");
|
||||
assert_eq!(args, "prod now");
|
||||
assert_eq!(raw, "/deploy prod now");
|
||||
}
|
||||
other => panic!("expected RunCommand, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_slash_text_is_not_a_command() {
|
||||
assert!(parse_slash_command("hello world").is_none());
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,156 @@
|
||||
fn main() {
|
||||
println!("harness {}", env!("CARGO_PKG_VERSION"));
|
||||
mod app;
|
||||
mod input;
|
||||
mod markdown;
|
||||
mod modal;
|
||||
mod render;
|
||||
mod state;
|
||||
mod terminal;
|
||||
|
||||
use std::path::PathBuf;
|
||||
|
||||
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]] | [tui]");
|
||||
}
|
||||
|
||||
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_headless(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).await {
|
||||
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
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn log_dir() -> PathBuf {
|
||||
dirs::data_dir()
|
||||
.unwrap_or_else(std::env::temp_dir)
|
||||
.join("ai-harness")
|
||||
.join("log")
|
||||
}
|
||||
|
||||
fn setup_tracing() -> anyhow::Result<tracing_appender::non_blocking::WorkerGuard> {
|
||||
let log_dir = log_dir();
|
||||
std::fs::create_dir_all(&log_dir)?;
|
||||
let file_appender = tracing_appender::rolling::daily(log_dir, "harness-tui.log");
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
|
||||
tracing_subscriber::fmt()
|
||||
.with_writer(non_blocking)
|
||||
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
|
||||
.init();
|
||||
Ok(guard)
|
||||
}
|
||||
|
||||
async fn run_tui() -> i32 {
|
||||
let cwd = match std::env::current_dir() {
|
||||
Ok(cwd) => cwd,
|
||||
Err(e) => {
|
||||
eprintln!("error: could not read cwd: {e}");
|
||||
return 1;
|
||||
}
|
||||
};
|
||||
|
||||
let _guard = match setup_tracing() {
|
||||
Ok(guard) => guard,
|
||||
Err(e) => {
|
||||
eprintln!("error: could not initialize logging: {e}");
|
||||
return 1;
|
||||
}
|
||||
};
|
||||
|
||||
let mut app = match crate::app::App::new(cwd).await {
|
||||
Ok(app) => app,
|
||||
Err(e) => {
|
||||
tracing::error!(error = %e, "failed to start TUI");
|
||||
eprintln!("error: {e}");
|
||||
return 1;
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(e) = app.run().await {
|
||||
tracing::error!(error = %e, "TUI error");
|
||||
eprintln!("error: {e}");
|
||||
return 1;
|
||||
}
|
||||
0
|
||||
}
|
||||
|
||||
#[tokio::main(flavor = "current_thread")]
|
||||
async fn main() {
|
||||
let args: Vec<String> = std::env::args().skip(1).collect();
|
||||
let exit_code = match args.first().map(String::as_str) {
|
||||
Some("run") => run_headless(&args[1..]).await,
|
||||
Some("tui") | None => run_tui().await,
|
||||
Some("help") | Some("--help") | Some("-h") => {
|
||||
print_usage();
|
||||
0
|
||||
}
|
||||
Some(cmd) => {
|
||||
eprintln!("unknown command: {cmd}");
|
||||
print_usage();
|
||||
2
|
||||
}
|
||||
};
|
||||
std::process::exit(exit_code);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
use pulldown_cmark::{CodeBlockKind, Event, HeadingLevel, Parser, Tag, TagEnd};
|
||||
use ratatui::style::{Color, Modifier, Style};
|
||||
use ratatui::text::{Line, Span};
|
||||
|
||||
/// Render a markdown string into wrapped ratatui lines.
|
||||
pub fn render_markdown(text: &str, width: u16) -> Vec<Line<'static>> {
|
||||
let mut lines: Vec<Line<'static>> = Vec::new();
|
||||
let mut current_spans: Vec<Span<'static>> = Vec::new();
|
||||
let mut current_text = String::new();
|
||||
|
||||
let mut in_bold = false;
|
||||
let mut in_italic = false;
|
||||
let mut in_code_block = false;
|
||||
let mut code_block_language = String::new();
|
||||
let mut list_stack: Vec<u64> = Vec::new();
|
||||
|
||||
let flush = |current: &mut String, spans: &mut Vec<Span<'static>>, bold, italic| {
|
||||
if current.is_empty() {
|
||||
return;
|
||||
}
|
||||
let style = base_style(bold, italic);
|
||||
let span = Span::styled(std::mem::take(current), style);
|
||||
spans.push(span);
|
||||
};
|
||||
|
||||
for event in Parser::new(text) {
|
||||
match event {
|
||||
Event::Start(tag) => match tag {
|
||||
Tag::Paragraph => {
|
||||
if !lines.is_empty() && !current_text.is_empty() {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
wrap_spans(&mut lines, &mut current_spans, width);
|
||||
}
|
||||
}
|
||||
Tag::Heading { level, .. } => {
|
||||
let level_num = heading_level_from(level);
|
||||
current_text.push_str(&"#".repeat(level_num as usize));
|
||||
current_text.push(' ');
|
||||
}
|
||||
Tag::Strong => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
in_bold = true;
|
||||
}
|
||||
Tag::Emphasis => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
in_italic = true;
|
||||
}
|
||||
Tag::List(start) => {
|
||||
list_stack.push(start.unwrap_or(1));
|
||||
}
|
||||
Tag::Item => {
|
||||
let prefix = if let Some(n) = list_stack.last_mut() {
|
||||
let p = format!("{n}. ");
|
||||
*n += 1;
|
||||
p
|
||||
} else {
|
||||
"• ".to_string()
|
||||
};
|
||||
current_text.push_str(&prefix);
|
||||
}
|
||||
Tag::CodeBlock(lang) => {
|
||||
in_code_block = true;
|
||||
code_block_language = match lang {
|
||||
CodeBlockKind::Fenced(name) => name.to_string(),
|
||||
CodeBlockKind::Indented => String::new(),
|
||||
};
|
||||
if !current_text.is_empty() || !current_spans.is_empty() {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
wrap_spans(&mut lines, &mut current_spans, width);
|
||||
}
|
||||
let border = if code_block_language.is_empty() {
|
||||
"┌────".to_string()
|
||||
} else {
|
||||
format!("┌──── {code_block_language}")
|
||||
};
|
||||
lines.push(Line::from(Span::styled(border, code_style())));
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
Event::End(tag_end) => match tag_end {
|
||||
TagEnd::Heading(_) => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
let mut spans = std::mem::take(&mut current_spans);
|
||||
for span in &mut spans {
|
||||
span.style = span.style.add_modifier(Modifier::BOLD);
|
||||
}
|
||||
wrap_spans(&mut lines, &mut spans, width);
|
||||
}
|
||||
TagEnd::Paragraph => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
wrap_spans(&mut lines, &mut current_spans, width);
|
||||
}
|
||||
TagEnd::Strong => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
in_bold = false;
|
||||
}
|
||||
TagEnd::Emphasis => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
in_italic = false;
|
||||
}
|
||||
TagEnd::List(_) => {
|
||||
list_stack.pop();
|
||||
}
|
||||
TagEnd::CodeBlock => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
if !current_spans.is_empty() {
|
||||
for line in wrap_spans_to_lines(&mut current_spans, width) {
|
||||
lines.push(prefix_code_block_line(line));
|
||||
}
|
||||
}
|
||||
lines.push(Line::from(Span::styled("└────", code_style())));
|
||||
in_code_block = false;
|
||||
code_block_language.clear();
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
Event::Text(t) => {
|
||||
if in_code_block {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
current_spans.push(Span::styled(t.to_string(), code_style()));
|
||||
} else {
|
||||
current_text.push_str(&t);
|
||||
}
|
||||
}
|
||||
Event::Code(c) => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
current_spans.push(Span::styled(c.to_string(), code_style()));
|
||||
}
|
||||
Event::Html(h) | Event::InlineHtml(h) => {
|
||||
current_text.push_str(&h);
|
||||
}
|
||||
Event::SoftBreak | Event::HardBreak => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
wrap_spans(&mut lines, &mut current_spans, width);
|
||||
}
|
||||
Event::Rule => {
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
wrap_spans(&mut lines, &mut current_spans, width);
|
||||
lines.push(Line::from("─".repeat(width as usize)));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
flush(&mut current_text, &mut current_spans, in_bold, in_italic);
|
||||
if in_code_block {
|
||||
for line in wrap_spans_to_lines(&mut current_spans, width) {
|
||||
lines.push(prefix_code_block_line(line));
|
||||
}
|
||||
lines.push(Line::from(Span::styled("└────", code_style())));
|
||||
} else {
|
||||
wrap_spans(&mut lines, &mut current_spans, width);
|
||||
}
|
||||
|
||||
if lines.is_empty() {
|
||||
lines.push(Line::default());
|
||||
}
|
||||
lines
|
||||
}
|
||||
|
||||
fn heading_level_from(level: HeadingLevel) -> u8 {
|
||||
match level {
|
||||
HeadingLevel::H1 => 1,
|
||||
HeadingLevel::H2 => 2,
|
||||
HeadingLevel::H3 => 3,
|
||||
HeadingLevel::H4 => 4,
|
||||
HeadingLevel::H5 => 5,
|
||||
HeadingLevel::H6 => 6,
|
||||
}
|
||||
}
|
||||
|
||||
fn base_style(bold: bool, italic: bool) -> Style {
|
||||
let mut style = Style::new();
|
||||
if bold {
|
||||
style = style.add_modifier(Modifier::BOLD);
|
||||
}
|
||||
if italic {
|
||||
style = style.add_modifier(Modifier::ITALIC);
|
||||
}
|
||||
style
|
||||
}
|
||||
|
||||
fn code_style() -> Style {
|
||||
Style::new().fg(Color::Yellow)
|
||||
}
|
||||
|
||||
fn prefix_code_block_line(line: Line<'static>) -> Line<'static> {
|
||||
let mut spans = vec![Span::styled("│ ", code_style())];
|
||||
spans.extend(line.spans);
|
||||
Line::from(spans)
|
||||
}
|
||||
|
||||
/// Flush accumulated spans into `lines`, wrapping to `width`.
|
||||
fn wrap_spans(lines: &mut Vec<Line<'static>>, spans: &mut Vec<Span<'static>>, width: u16) {
|
||||
if spans.is_empty() {
|
||||
return;
|
||||
}
|
||||
for line in wrap_spans_to_lines(spans, width) {
|
||||
lines.push(line);
|
||||
}
|
||||
spans.clear();
|
||||
}
|
||||
|
||||
fn wrap_spans_to_lines(spans: &mut Vec<Span<'static>>, width: u16) -> Vec<Line<'static>> {
|
||||
let width = width.max(1) as usize;
|
||||
let mut out: Vec<Line<'static>> = Vec::new();
|
||||
let mut current_line: Vec<Span<'static>> = Vec::new();
|
||||
let mut current_width = 0usize;
|
||||
|
||||
for span in spans.drain(..) {
|
||||
for word in span.content.split(' ') {
|
||||
let word_width = word.chars().count();
|
||||
let sep = if current_width == 0 { 0 } else { 1 };
|
||||
if current_width + sep + word_width > width && current_width > 0 {
|
||||
out.push(Line::from(std::mem::take(&mut current_line)));
|
||||
current_width = 0;
|
||||
}
|
||||
if current_width > 0 {
|
||||
current_line.push(Span::styled(" ", span.style));
|
||||
current_width += 1;
|
||||
}
|
||||
current_line.push(Span::styled(word.to_string(), span.style));
|
||||
current_width += word_width;
|
||||
}
|
||||
}
|
||||
|
||||
if !current_line.is_empty() {
|
||||
out.push(Line::from(current_line));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn renders_heading_and_bold() {
|
||||
let lines = render_markdown("# Hello\n\n**bold** text", 40);
|
||||
assert!(!lines.is_empty());
|
||||
let first: String = lines[0]
|
||||
.spans
|
||||
.iter()
|
||||
.map(|s| s.content.to_string())
|
||||
.collect();
|
||||
assert!(first.contains("# Hello"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wraps_long_paragraphs() {
|
||||
let text = "a ".repeat(50);
|
||||
let lines = render_markdown(&text, 20);
|
||||
assert!(lines.len() > 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
//! Modal helpers for permission dialogs and the session picker.
|
||||
//!
|
||||
//! Rendering lives in `render.rs`; key handling lives in `input.rs`. This module
|
||||
//! provides shared utilities for modal state transitions.
|
||||
|
||||
use harness_app::EngineHandle;
|
||||
use harness_core::permission::PermissionReply;
|
||||
use harness_core::types::Session;
|
||||
|
||||
use crate::state::{AppState, ModalState};
|
||||
|
||||
/// Resolve the current permission request with `reply`.
|
||||
pub fn resolve_permission(state: &mut AppState, engine: &EngineHandle, reply: PermissionReply) {
|
||||
state.handle_permission_reply(reply, engine);
|
||||
}
|
||||
|
||||
/// Open the session picker with the supplied sessions.
|
||||
pub fn open_session_picker(state: &mut AppState, sessions: Vec<Session>) {
|
||||
state.open_session_picker(sessions);
|
||||
}
|
||||
|
||||
/// Load a session's messages and parts into state.
|
||||
pub async fn select_session(
|
||||
state: &mut AppState,
|
||||
engine: &EngineHandle,
|
||||
session: Session,
|
||||
) -> Result<(), harness_app::AppError> {
|
||||
let session_id = session.id.clone();
|
||||
// Repair a session that crashed mid-run before loading it, so the transcript shows the
|
||||
// aborted turn rather than a tool call frozen "running".
|
||||
engine.repair_session(&session_id).await?;
|
||||
let messages = engine.session_messages(session_id.clone()).await?;
|
||||
let mut all_parts = Vec::new();
|
||||
for message in &messages {
|
||||
let parts = engine.message_parts(message.id.clone()).await?;
|
||||
all_parts.extend(parts);
|
||||
}
|
||||
state.set_session(session, messages, all_parts);
|
||||
// Surface any subtasks this session spawned (drill-in and session-switch both land here).
|
||||
if let Ok(jobs) = engine.jobs(session_id).await {
|
||||
state.set_jobs(jobs);
|
||||
}
|
||||
state.modal = ModalState::None;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,647 @@
|
||||
use ratatui::layout::{Constraint, Direction, Layout, Rect};
|
||||
use ratatui::style::{Color, Modifier, Style};
|
||||
use ratatui::text::{Line, Span};
|
||||
use ratatui::widgets::{Block, Borders, Clear, Paragraph, Wrap};
|
||||
use ratatui::Frame;
|
||||
|
||||
use crate::markdown::render_markdown;
|
||||
use crate::state::{AppState, ModalState, PartView, ToolStateView};
|
||||
use harness_core::types::Role;
|
||||
|
||||
/// Render the full UI into the terminal frame.
|
||||
pub fn render(frame: &mut Frame, state: &mut AppState) {
|
||||
let area = frame.area();
|
||||
let chunks = Layout::default()
|
||||
.direction(Direction::Vertical)
|
||||
.constraints([
|
||||
Constraint::Length(1),
|
||||
Constraint::Min(3),
|
||||
Constraint::Length(3),
|
||||
Constraint::Length(1),
|
||||
])
|
||||
.split(area);
|
||||
|
||||
render_header(frame, state, chunks[0]);
|
||||
render_chat(frame, state, chunks[1]);
|
||||
render_input(frame, state, chunks[2]);
|
||||
render_status(frame, state, chunks[3]);
|
||||
|
||||
match &state.modal {
|
||||
ModalState::Permission { requests } => {
|
||||
render_permission_modal(frame, &requests[0], area);
|
||||
}
|
||||
ModalState::SessionPicker { sessions, selected } => {
|
||||
render_session_picker(frame, sessions, *selected, area);
|
||||
}
|
||||
ModalState::JobsPane { selected } => {
|
||||
render_jobs_pane(frame, &state.jobs, *selected, area);
|
||||
}
|
||||
ModalState::None => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn render_header(frame: &mut Frame, state: &AppState, area: Rect) {
|
||||
let title = if state.session_title.is_empty() {
|
||||
"new session"
|
||||
} else {
|
||||
state.session_title.as_str()
|
||||
};
|
||||
let text = format!("{} · {} · {}", title, state.agent, state.model_ref);
|
||||
let paragraph = Paragraph::new(text)
|
||||
.style(Style::new().add_modifier(Modifier::BOLD))
|
||||
.alignment(ratatui::layout::Alignment::Center);
|
||||
frame.render_widget(paragraph, area);
|
||||
}
|
||||
|
||||
fn render_chat(frame: &mut Frame, state: &mut AppState, area: Rect) {
|
||||
let effective_width = area.width.saturating_sub(4).max(1);
|
||||
// Cached lines are wrapped to a specific width, so a resize must drop them.
|
||||
if state.render_width != effective_width {
|
||||
state.invalidate_caches();
|
||||
state.render_width = effective_width;
|
||||
}
|
||||
|
||||
let mut lines: Vec<Line<'static>> = Vec::new();
|
||||
for message in &mut state.messages {
|
||||
let role_style = match message.role {
|
||||
Role::User => Style::new().fg(Color::Cyan),
|
||||
Role::Assistant => Style::new(),
|
||||
};
|
||||
let prefix = match message.role {
|
||||
Role::User => "> ",
|
||||
Role::Assistant => "",
|
||||
};
|
||||
let mut first = true;
|
||||
for part in &mut message.parts {
|
||||
let part_lines = render_part(part, effective_width, role_style, prefix, first);
|
||||
lines.extend(part_lines);
|
||||
first = false;
|
||||
}
|
||||
lines.push(Line::default());
|
||||
}
|
||||
|
||||
let paragraph = Paragraph::new(lines)
|
||||
.block(Block::default().borders(Borders::ALL).title(" chat "))
|
||||
.scroll((state.scroll_offset, 0));
|
||||
frame.render_widget(paragraph, area);
|
||||
}
|
||||
|
||||
fn render_part(
|
||||
part: &mut PartView,
|
||||
effective_width: u16,
|
||||
role_style: Style,
|
||||
prefix: &'static str,
|
||||
first: bool,
|
||||
) -> Vec<Line<'static>> {
|
||||
match part {
|
||||
PartView::Text {
|
||||
text, cached_lines, ..
|
||||
} => {
|
||||
// Cache the prefix-free markdown render; the `> ` prefix is cheap to re-apply
|
||||
// to the clone each frame and would otherwise poison a shared cache.
|
||||
let base = cached_lines.get_or_insert_with(|| render_markdown(text, effective_width));
|
||||
let mut rendered = base.clone();
|
||||
if first && !prefix.is_empty() {
|
||||
if let Some(first_line) = rendered.first_mut() {
|
||||
let mut spans = vec![Span::styled(prefix, role_style)];
|
||||
spans.extend(first_line.spans.clone());
|
||||
*first_line = Line::from(spans);
|
||||
}
|
||||
}
|
||||
rendered
|
||||
}
|
||||
PartView::Reasoning { text, .. } => {
|
||||
let style = Style::new()
|
||||
.fg(Color::DarkGray)
|
||||
.add_modifier(Modifier::ITALIC);
|
||||
render_markdown(text, effective_width)
|
||||
.into_iter()
|
||||
.map(|line| {
|
||||
let mut spans = vec![Span::styled("🧠 ", style)];
|
||||
spans.extend(line.spans);
|
||||
Line::from(spans)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
PartView::Tool {
|
||||
name,
|
||||
state,
|
||||
cached_lines,
|
||||
..
|
||||
} => {
|
||||
if let Some(cached) = cached_lines {
|
||||
return cached.clone();
|
||||
}
|
||||
let status = state.status_label();
|
||||
let icon = match state {
|
||||
ToolStateView::Pending { .. } => "⏳",
|
||||
ToolStateView::Running { .. } => "⚙",
|
||||
ToolStateView::Completed { .. } => "✓",
|
||||
ToolStateView::Error { .. } => "✗",
|
||||
};
|
||||
let title = format!("{icon} {name} — {status}");
|
||||
let mut lines = vec![Line::from(title)];
|
||||
match state {
|
||||
ToolStateView::Pending { partial_input } => {
|
||||
lines.extend(render_markdown(partial_input, effective_width));
|
||||
}
|
||||
ToolStateView::Completed { output, .. } => {
|
||||
lines.extend(render_markdown(output, effective_width));
|
||||
}
|
||||
ToolStateView::Error { error } => {
|
||||
lines.extend(render_markdown(error, effective_width));
|
||||
}
|
||||
ToolStateView::Running { .. } => {}
|
||||
}
|
||||
*cached_lines = Some(lines.clone());
|
||||
lines
|
||||
}
|
||||
PartView::StepFinish { .. } => vec![Line::from("─── step ───")],
|
||||
}
|
||||
}
|
||||
|
||||
fn render_input(frame: &mut Frame, state: &mut AppState, area: Rect) {
|
||||
let block = Block::default().borders(Borders::ALL).title(" input ");
|
||||
state.input.set_block(block);
|
||||
frame.render_widget(&state.input, area);
|
||||
}
|
||||
|
||||
fn render_status(frame: &mut Frame, state: &AppState, area: Rect) {
|
||||
let spinner = if state.running { "⠋ " } else { "" };
|
||||
let status = if state.running { "running" } else { "idle" };
|
||||
let hints = if state.ctrl_c_pressed {
|
||||
"Press Ctrl+C again to quit"
|
||||
} else {
|
||||
"Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C: quit"
|
||||
};
|
||||
let usage = format!(
|
||||
"{} · ${:.4}",
|
||||
format_tokens(state.session_tokens),
|
||||
state.session_cost
|
||||
);
|
||||
let text = format!("{spinner}{status} · {usage} · {hints}");
|
||||
let paragraph = Paragraph::new(text).style(Style::new().fg(Color::Gray));
|
||||
frame.render_widget(paragraph, area);
|
||||
}
|
||||
|
||||
/// Compact token count: `1234` → `1.2k`, `2_000_000` → `2.0M`.
|
||||
fn format_tokens(n: u64) -> String {
|
||||
match n {
|
||||
0..=999 => format!("{n} tok"),
|
||||
1_000..=999_999 => format!("{:.1}k tok", n as f64 / 1_000.0),
|
||||
_ => format!("{:.1}M tok", n as f64 / 1_000_000.0),
|
||||
}
|
||||
}
|
||||
|
||||
fn render_permission_modal(
|
||||
frame: &mut Frame,
|
||||
request: &harness_core::event::PermissionRequest,
|
||||
area: Rect,
|
||||
) {
|
||||
let popup = centered_rect(60, 60, area);
|
||||
frame.render_widget(Clear, popup);
|
||||
|
||||
let text = vec![
|
||||
Line::from("Permission required").style(Style::new().add_modifier(Modifier::BOLD)),
|
||||
Line::default(),
|
||||
Line::from(format!("permission: {}", request.permission)),
|
||||
Line::from(format!("pattern: {}", request.pattern)),
|
||||
Line::default(),
|
||||
Line::from("y: allow once a: allow always n: reject"),
|
||||
];
|
||||
|
||||
let block = Block::default().borders(Borders::ALL).title(" permission ");
|
||||
let paragraph = Paragraph::new(text).block(block).wrap(Wrap { trim: true });
|
||||
frame.render_widget(paragraph, popup);
|
||||
}
|
||||
|
||||
fn render_session_picker(
|
||||
frame: &mut Frame,
|
||||
sessions: &[harness_core::types::Session],
|
||||
selected: usize,
|
||||
area: Rect,
|
||||
) {
|
||||
let popup = centered_rect(60, 60, area);
|
||||
frame.render_widget(Clear, popup);
|
||||
|
||||
let mut text: Vec<Line<'static>> =
|
||||
vec![Line::from("Select session").style(Style::new().add_modifier(Modifier::BOLD))];
|
||||
for (i, session) in sessions.iter().enumerate() {
|
||||
let marker = if i == selected { "> " } else { " " };
|
||||
let line = format!(
|
||||
"{}{} · {} · {}/{}",
|
||||
marker, session.id, session.agent, session.model.provider_id, session.model.model_id
|
||||
);
|
||||
let style = if i == selected {
|
||||
Style::new().bg(Color::Blue).fg(Color::White)
|
||||
} else {
|
||||
Style::new()
|
||||
};
|
||||
text.push(Line::from(Span::styled(line, style)));
|
||||
}
|
||||
|
||||
let block = Block::default().borders(Borders::ALL).title(" sessions ");
|
||||
let paragraph = Paragraph::new(text).block(block).wrap(Wrap { trim: false });
|
||||
frame.render_widget(paragraph, popup);
|
||||
}
|
||||
|
||||
fn render_jobs_pane(
|
||||
frame: &mut Frame,
|
||||
jobs: &[harness_core::engine::JobRecord],
|
||||
selected: usize,
|
||||
area: Rect,
|
||||
) {
|
||||
let popup = centered_rect(70, 70, area);
|
||||
frame.render_widget(Clear, popup);
|
||||
|
||||
let mut text: Vec<Line<'static>> =
|
||||
vec![Line::from("Background jobs").style(Style::new().add_modifier(Modifier::BOLD))];
|
||||
|
||||
if jobs.is_empty() {
|
||||
text.push(Line::default());
|
||||
text.push(Line::from("No subtasks spawned yet.").style(Style::new().fg(Color::DarkGray)));
|
||||
} else {
|
||||
for (i, job) in jobs.iter().enumerate() {
|
||||
let marker = if i == selected { "> " } else { " " };
|
||||
let (icon, color) = job_state_style(job.state);
|
||||
let header = format!(
|
||||
"{marker}{icon} {} · {} · {}",
|
||||
job.alias,
|
||||
job.agent,
|
||||
job_state_label(job.state),
|
||||
);
|
||||
let style = if i == selected {
|
||||
Style::new().bg(Color::Blue).fg(Color::White)
|
||||
} else {
|
||||
Style::new().fg(color)
|
||||
};
|
||||
text.push(Line::from(Span::styled(header, style)));
|
||||
if let Some(objective) = &job.objective {
|
||||
text.push(
|
||||
Line::from(format!(" {objective}")).style(Style::new().fg(Color::Gray)),
|
||||
);
|
||||
}
|
||||
if !job.context_files.is_empty() {
|
||||
let files: Vec<&str> = job
|
||||
.context_files
|
||||
.iter()
|
||||
.take(8)
|
||||
.map(|f| f.path.as_str())
|
||||
.collect();
|
||||
text.push(
|
||||
Line::from(format!(" read: {}", files.join(", ")))
|
||||
.style(Style::new().fg(Color::DarkGray)),
|
||||
);
|
||||
}
|
||||
}
|
||||
text.push(Line::default());
|
||||
text.push(
|
||||
Line::from("Enter: open subtask · Esc: close").style(Style::new().fg(Color::DarkGray)),
|
||||
);
|
||||
}
|
||||
|
||||
let block = Block::default().borders(Borders::ALL).title(" jobs ");
|
||||
let paragraph = Paragraph::new(text).block(block).wrap(Wrap { trim: false });
|
||||
frame.render_widget(paragraph, popup);
|
||||
}
|
||||
|
||||
fn job_state_label(state: harness_core::engine::JobState) -> &'static str {
|
||||
use harness_core::engine::JobState;
|
||||
match state {
|
||||
JobState::Running => "running",
|
||||
JobState::Completed => "completed",
|
||||
JobState::Error => "error",
|
||||
JobState::Cancelled => "cancelled",
|
||||
}
|
||||
}
|
||||
|
||||
fn job_state_style(state: harness_core::engine::JobState) -> (&'static str, Color) {
|
||||
use harness_core::engine::JobState;
|
||||
match state {
|
||||
JobState::Running => ("⚙", Color::Yellow),
|
||||
JobState::Completed => ("✓", Color::Green),
|
||||
JobState::Error => ("✗", Color::Red),
|
||||
JobState::Cancelled => ("∅", Color::DarkGray),
|
||||
}
|
||||
}
|
||||
|
||||
fn centered_rect(percent_x: u16, percent_y: u16, r: Rect) -> Rect {
|
||||
let popup_layout = Layout::default()
|
||||
.direction(Direction::Vertical)
|
||||
.constraints([
|
||||
Constraint::Percentage((100 - percent_y) / 2),
|
||||
Constraint::Percentage(percent_y),
|
||||
Constraint::Percentage((100 - percent_y) / 2),
|
||||
])
|
||||
.split(r);
|
||||
|
||||
Layout::default()
|
||||
.direction(Direction::Horizontal)
|
||||
.constraints([
|
||||
Constraint::Percentage((100 - percent_x) / 2),
|
||||
Constraint::Percentage(percent_x),
|
||||
Constraint::Percentage((100 - percent_x) / 2),
|
||||
])
|
||||
.split(popup_layout[1])[1]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use ratatui::backend::TestBackend;
|
||||
use ratatui::Terminal;
|
||||
|
||||
use super::*;
|
||||
use crate::state::{AppState, MessageView, PartView, ToolStateView};
|
||||
use harness_core::event::PermissionRequest;
|
||||
use harness_core::types::{MessageId, ModelRef, PartId, Role, Session, SessionId};
|
||||
|
||||
fn buffer_to_string(backend: &TestBackend) -> String {
|
||||
let buffer = backend.buffer();
|
||||
let area = buffer.area;
|
||||
let mut result = String::new();
|
||||
for y in 0..area.height {
|
||||
for x in 0..area.width {
|
||||
result.push_str(buffer[(x, y)].symbol());
|
||||
}
|
||||
while result.ends_with(' ') {
|
||||
result.pop();
|
||||
}
|
||||
result.push('\n');
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_empty_session() {
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new(
|
||||
"anthropic/claude-sonnet-4-5".to_string(),
|
||||
"orchestrator".to_string(),
|
||||
);
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
insta::assert_snapshot!(buffer_to_string(terminal.backend()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_user_and_assistant_messages() {
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new(
|
||||
"anthropic/claude-sonnet-4-5".to_string(),
|
||||
"orchestrator".to_string(),
|
||||
);
|
||||
state.session_id = Some(SessionId("ses_test_001".to_string()));
|
||||
state.messages.push(MessageView {
|
||||
id: MessageId("msg_test_001".to_string()),
|
||||
role: Role::User,
|
||||
parts: vec![PartView::Text {
|
||||
id: PartId("prt_test_001".to_string()),
|
||||
text: "Hello, can you read foo.txt?".to_string(),
|
||||
cached_lines: None,
|
||||
}],
|
||||
});
|
||||
state.messages.push(MessageView {
|
||||
id: MessageId("msg_test_002".to_string()),
|
||||
role: Role::Assistant,
|
||||
parts: vec![PartView::Text {
|
||||
id: PartId("prt_test_002".to_string()),
|
||||
text: "I'll read that file for you.\n\n```rust\nlet x = 42;\n```".to_string(),
|
||||
cached_lines: None,
|
||||
}],
|
||||
});
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
insta::assert_snapshot!(buffer_to_string(terminal.backend()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_tool_card_completed() {
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new(
|
||||
"anthropic/claude-sonnet-4-5".to_string(),
|
||||
"orchestrator".to_string(),
|
||||
);
|
||||
state.session_id = Some(SessionId("ses_test_001".to_string()));
|
||||
state.messages.push(MessageView {
|
||||
id: MessageId("msg_test_003".to_string()),
|
||||
role: Role::Assistant,
|
||||
parts: vec![PartView::Tool {
|
||||
id: PartId("prt_test_003".to_string()),
|
||||
name: "read".to_string(),
|
||||
state: ToolStateView::Completed {
|
||||
title: "read foo.txt".to_string(),
|
||||
output: "file contents here".to_string(),
|
||||
},
|
||||
cached_lines: None,
|
||||
}],
|
||||
});
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
insta::assert_snapshot!(buffer_to_string(terminal.backend()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_tool_card_running() {
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new(
|
||||
"anthropic/claude-sonnet-4-5".to_string(),
|
||||
"orchestrator".to_string(),
|
||||
);
|
||||
state.session_id = Some(SessionId("ses_test_001".to_string()));
|
||||
state.messages.push(MessageView {
|
||||
id: MessageId("msg_test_004".to_string()),
|
||||
role: Role::Assistant,
|
||||
parts: vec![PartView::Tool {
|
||||
id: PartId("prt_test_004".to_string()),
|
||||
name: "read".to_string(),
|
||||
state: ToolStateView::Running {
|
||||
title: Some("read foo.txt".to_string()),
|
||||
},
|
||||
cached_lines: None,
|
||||
}],
|
||||
});
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
insta::assert_snapshot!(buffer_to_string(terminal.backend()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_tool_card_error() {
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new(
|
||||
"anthropic/claude-sonnet-4-5".to_string(),
|
||||
"orchestrator".to_string(),
|
||||
);
|
||||
state.session_id = Some(SessionId("ses_test_001".to_string()));
|
||||
state.messages.push(MessageView {
|
||||
id: MessageId("msg_test_005".to_string()),
|
||||
role: Role::Assistant,
|
||||
parts: vec![PartView::Tool {
|
||||
id: PartId("prt_test_005".to_string()),
|
||||
name: "read".to_string(),
|
||||
state: ToolStateView::Error {
|
||||
error: "file not found".to_string(),
|
||||
},
|
||||
cached_lines: None,
|
||||
}],
|
||||
});
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
insta::assert_snapshot!(buffer_to_string(terminal.backend()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_permission_modal() {
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new(
|
||||
"anthropic/claude-sonnet-4-5".to_string(),
|
||||
"orchestrator".to_string(),
|
||||
);
|
||||
let request = PermissionRequest {
|
||||
id: "test".to_string(),
|
||||
session_id: SessionId("ses_test_001".to_string()),
|
||||
permission: "bash".to_string(),
|
||||
pattern: "rm -rf /".to_string(),
|
||||
always_pattern: "rm *".to_string(),
|
||||
metadata: serde_json::Value::Null,
|
||||
};
|
||||
state.modal = ModalState::Permission {
|
||||
requests: vec![request],
|
||||
};
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
insta::assert_snapshot!(buffer_to_string(terminal.backend()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_session_picker_modal() {
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new(
|
||||
"anthropic/claude-sonnet-4-5".to_string(),
|
||||
"orchestrator".to_string(),
|
||||
);
|
||||
let sessions = vec![
|
||||
Session {
|
||||
id: SessionId("ses_test_001".to_string()),
|
||||
parent_id: None,
|
||||
depth: 0,
|
||||
title: "Session One".to_string(),
|
||||
agent: "orchestrator".to_string(),
|
||||
model: ModelRef::new("anthropic", "claude-sonnet-4-5"),
|
||||
usage: harness_core::types::TokenUsage::default(),
|
||||
cost: 0.0,
|
||||
extra_rules: Vec::new(),
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
},
|
||||
Session {
|
||||
id: SessionId("ses_test_002".to_string()),
|
||||
parent_id: None,
|
||||
depth: 0,
|
||||
title: "Session Two".to_string(),
|
||||
agent: "coder".to_string(),
|
||||
model: ModelRef::new("openai", "gpt-4o"),
|
||||
usage: harness_core::types::TokenUsage::default(),
|
||||
cost: 0.0,
|
||||
extra_rules: Vec::new(),
|
||||
created_at: 2,
|
||||
updated_at: 2,
|
||||
},
|
||||
Session {
|
||||
id: SessionId("ses_test_003".to_string()),
|
||||
parent_id: None,
|
||||
depth: 0,
|
||||
title: "Session Three".to_string(),
|
||||
agent: "reviewer".to_string(),
|
||||
model: ModelRef::new("google", "gemini-2.5"),
|
||||
usage: harness_core::types::TokenUsage::default(),
|
||||
cost: 0.0,
|
||||
extra_rules: Vec::new(),
|
||||
created_at: 3,
|
||||
updated_at: 3,
|
||||
},
|
||||
];
|
||||
state.modal = ModalState::SessionPicker {
|
||||
sessions,
|
||||
selected: 1,
|
||||
};
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
insta::assert_snapshot!(buffer_to_string(terminal.backend()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_jobs_pane() {
|
||||
use harness_core::engine::{ContextFile, JobRecord, JobState};
|
||||
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new(
|
||||
"anthropic/claude-sonnet-4-5".to_string(),
|
||||
"orchestrator".to_string(),
|
||||
);
|
||||
state.session_id = Some(SessionId("ses_test_001".to_string()));
|
||||
state.jobs = vec![
|
||||
JobRecord {
|
||||
task_id: "t1".to_string(),
|
||||
alias: "exp-1".to_string(),
|
||||
parent_session: SessionId("ses_test_001".to_string()),
|
||||
child_session: SessionId("ses_child_001".to_string()),
|
||||
agent: "explorer".to_string(),
|
||||
description: "map auth".to_string(),
|
||||
objective: Some("map the auth flow".to_string()),
|
||||
state: JobState::Completed,
|
||||
reconciled: true,
|
||||
result_summary: Some("done".to_string()),
|
||||
context_files: vec![ContextFile {
|
||||
path: "src/auth.rs".to_string(),
|
||||
lines: 42,
|
||||
}],
|
||||
launched_at: 1,
|
||||
updated_at: 2,
|
||||
last_used_at: 2,
|
||||
},
|
||||
JobRecord {
|
||||
task_id: "t2".to_string(),
|
||||
alias: "fix-1".to_string(),
|
||||
parent_session: SessionId("ses_test_001".to_string()),
|
||||
child_session: SessionId("ses_child_002".to_string()),
|
||||
agent: "fixer".to_string(),
|
||||
description: "patch bug".to_string(),
|
||||
objective: Some("fix the null deref".to_string()),
|
||||
state: JobState::Running,
|
||||
reconciled: false,
|
||||
result_summary: None,
|
||||
context_files: Vec::new(),
|
||||
launched_at: 3,
|
||||
updated_at: 3,
|
||||
last_used_at: 3,
|
||||
},
|
||||
];
|
||||
state.modal = ModalState::JobsPane { selected: 1 };
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
insta::assert_snapshot!(buffer_to_string(terminal.backend()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_empty_state_produces_frame() {
|
||||
let backend = TestBackend::new(80, 24);
|
||||
let mut terminal = Terminal::new(backend).unwrap();
|
||||
let mut state = AppState::new("anthropic/claude".to_string(), "orchestrator".to_string());
|
||||
terminal.draw(|frame| render(frame, &mut state)).unwrap();
|
||||
let buffer = terminal.backend().buffer().clone();
|
||||
assert_eq!(buffer.area.width, 80);
|
||||
assert_eq!(buffer.area.height, 24);
|
||||
// Header should contain the agent/model line.
|
||||
let header_row: String = buffer
|
||||
.content
|
||||
.chunks(80)
|
||||
.next()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|c| c.symbol())
|
||||
.collect();
|
||||
assert!(header_row.contains("orchestrator"));
|
||||
assert!(header_row.contains("anthropic/claude"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
---
|
||||
source: crates/harness-tui/src/render.rs
|
||||
assertion_line: 381
|
||||
expression: buffer_to_string(terminal.backend())
|
||||
---
|
||||
new session · orchestrator · anthropic/claude-sonnet-4-5
|
||||
┌ chat ────────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
┌ input ───────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
idle · 0 tok · $0.0000 · Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C:
|
||||
@@ -0,0 +1,29 @@
|
||||
---
|
||||
source: crates/harness-tui/src/render.rs
|
||||
assertion_line: 621
|
||||
expression: buffer_to_string(terminal.backend())
|
||||
---
|
||||
new session · orchestrator · anthropic/claude-sonnet-4-5
|
||||
┌ chat ────────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
│ │
|
||||
│ ┌ jobs ────────────────────────────────────────────────┐ │
|
||||
│ │Background jobs │ │
|
||||
│ │ ✓ exp-1 · explorer · completed │ │
|
||||
│ │ map the auth flow │ │
|
||||
│ │ read: src/auth.rs │ │
|
||||
│ │> ⚙ fix-1 · fixer · running │ │
|
||||
│ │ fix the null deref │ │
|
||||
│ │ │ │
|
||||
│ │Enter: open subtask · Esc: close │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
└───────────└──────────────────────────────────────────────────────┘───────────┘
|
||||
┌ input ───────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
idle · 0 tok · $0.0000 · Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C:
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
---
|
||||
source: crates/harness-tui/src/render.rs
|
||||
assertion_line: 511
|
||||
expression: buffer_to_string(terminal.backend())
|
||||
---
|
||||
new session · orchestrator · anthropic/claude-sonnet-4-5
|
||||
┌ chat ────────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ ┌ permission ──────────────────────────────────┐ │
|
||||
│ │Permission required │ │
|
||||
│ │ │ │
|
||||
│ │permission: bash │ │
|
||||
│ │pattern: rm -rf / │ │
|
||||
│ │ │ │
|
||||
│ │y: allow once a: allow always n: reject │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ └──────────────────────────────────────────────┘ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
┌ input ───────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
idle · 0 tok · $0.0000 · Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C:
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
---
|
||||
source: crates/harness-tui/src/render.rs
|
||||
assertion_line: 568
|
||||
expression: buffer_to_string(terminal.backend())
|
||||
---
|
||||
new session · orchestrator · anthropic/claude-sonnet-4-5
|
||||
┌ chat ────────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ ┌ sessions ────────────────────────────────────┐ │
|
||||
│ │Select session │ │
|
||||
│ │ ses_test_001 · orchestrator · │ │
|
||||
│ │anthropic/claude-sonnet-4-5 │ │
|
||||
│ │> ses_test_002 · coder · openai/gpt-4o │ │
|
||||
│ │ ses_test_003 · reviewer · google/gemini-2.5 │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ │ │ │
|
||||
│ └──────────────────────────────────────────────┘ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
┌ input ───────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
idle · 0 tok · $0.0000 · Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C:
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
---
|
||||
source: crates/harness-tui/src/render.rs
|
||||
assertion_line: 438
|
||||
expression: buffer_to_string(terminal.backend())
|
||||
---
|
||||
new session · orchestrator · anthropic/claude-sonnet-4-5
|
||||
┌ chat ────────────────────────────────────────────────────────────────────────┐
|
||||
│✓ read — read foo.txt │
|
||||
│file contents here │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
┌ input ───────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
idle · 0 tok · $0.0000 · Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C:
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
---
|
||||
source: crates/harness-tui/src/render.rs
|
||||
assertion_line: 488
|
||||
expression: buffer_to_string(terminal.backend())
|
||||
---
|
||||
new session · orchestrator · anthropic/claude-sonnet-4-5
|
||||
┌ chat ────────────────────────────────────────────────────────────────────────┐
|
||||
│✗ read — error: file not found │
|
||||
│file not found │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
┌ input ───────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
idle · 0 tok · $0.0000 · Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C:
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
---
|
||||
source: crates/harness-tui/src/render.rs
|
||||
assertion_line: 463
|
||||
expression: buffer_to_string(terminal.backend())
|
||||
---
|
||||
new session · orchestrator · anthropic/claude-sonnet-4-5
|
||||
┌ chat ────────────────────────────────────────────────────────────────────────┐
|
||||
│⚙ read — read foo.txt │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
┌ input ───────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
idle · 0 tok · $0.0000 · Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C:
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
---
|
||||
source: crates/harness-tui/src/render.rs
|
||||
assertion_line: 412
|
||||
expression: buffer_to_string(terminal.backend())
|
||||
---
|
||||
new session · orchestrator · anthropic/claude-sonnet-4-5
|
||||
┌ chat ────────────────────────────────────────────────────────────────────────┐
|
||||
│> Hello, can you read foo.txt? │
|
||||
│ │
|
||||
│I'll read that file for you. │
|
||||
│┌──── rust │
|
||||
││ let x = 42; │
|
||||
│└──── │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
┌ input ───────────────────────────────────────────────────────────────────────┐
|
||||
│ │
|
||||
└──────────────────────────────────────────────────────────────────────────────┘
|
||||
idle · 0 tok · $0.0000 · Ctrl+S: sessions | Ctrl+J: jobs | Esc: abort | Ctrl+C:
|
||||
@@ -0,0 +1,453 @@
|
||||
use harness_app::EngineHandle;
|
||||
use harness_core::engine::JobRecord;
|
||||
use harness_core::event::{AppEvent, PermissionRequest, RunOutcome};
|
||||
use harness_core::permission::PermissionReply;
|
||||
use harness_core::types::{Message, Part, PartBody, PartId, Role, Session, SessionId, ToolState};
|
||||
use ratatui::text::Line;
|
||||
|
||||
/// Cached view of a message in the chat viewport.
|
||||
#[derive(Debug)]
|
||||
pub struct MessageView {
|
||||
pub id: harness_core::types::MessageId,
|
||||
pub role: Role,
|
||||
pub parts: Vec<PartView>,
|
||||
}
|
||||
|
||||
/// Cached view of a part. The optional `cached_lines` field stores a pre-rendered markdown
|
||||
/// representation; it is invalidated whenever the underlying text changes.
|
||||
#[derive(Debug)]
|
||||
pub enum PartView {
|
||||
Text {
|
||||
id: PartId,
|
||||
text: String,
|
||||
cached_lines: Option<Vec<Line<'static>>>,
|
||||
},
|
||||
Reasoning {
|
||||
id: PartId,
|
||||
text: String,
|
||||
},
|
||||
Tool {
|
||||
id: PartId,
|
||||
name: String,
|
||||
state: ToolStateView,
|
||||
cached_lines: Option<Vec<Line<'static>>>,
|
||||
},
|
||||
StepFinish {
|
||||
id: PartId,
|
||||
},
|
||||
}
|
||||
|
||||
/// Simplified tool state for rendering.
|
||||
#[derive(Debug)]
|
||||
pub enum ToolStateView {
|
||||
Pending { partial_input: String },
|
||||
Running { title: Option<String> },
|
||||
Completed { title: String, output: String },
|
||||
Error { error: String },
|
||||
}
|
||||
|
||||
impl ToolStateView {
|
||||
fn from_tool_state(state: &ToolState) -> Self {
|
||||
match state {
|
||||
ToolState::Pending { partial_input } => Self::Pending {
|
||||
partial_input: partial_input.clone(),
|
||||
},
|
||||
ToolState::Running { title, .. } => Self::Running {
|
||||
title: title.clone(),
|
||||
},
|
||||
ToolState::Completed { title, output, .. } => Self::Completed {
|
||||
title: title.clone(),
|
||||
output: output.clone(),
|
||||
},
|
||||
ToolState::Error { error, .. } => Self::Error {
|
||||
error: error.clone(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn status_label(&self) -> String {
|
||||
match self {
|
||||
Self::Pending { .. } => "pending".to_string(),
|
||||
Self::Running { title } => title.clone().unwrap_or_else(|| "running".to_string()),
|
||||
Self::Completed { title, .. } => title.clone(),
|
||||
Self::Error { error } => format!("error: {error}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Active modal overlay state.
|
||||
#[derive(Debug, Default)]
|
||||
pub enum ModalState {
|
||||
#[default]
|
||||
None,
|
||||
Permission {
|
||||
requests: Vec<PermissionRequest>,
|
||||
},
|
||||
SessionPicker {
|
||||
sessions: Vec<Session>,
|
||||
selected: usize,
|
||||
},
|
||||
/// Background job board for the current session, with a cursor for drill-in.
|
||||
JobsPane {
|
||||
selected: usize,
|
||||
},
|
||||
}
|
||||
|
||||
/// All mutable UI state lives here.
|
||||
#[derive(Debug)]
|
||||
pub struct AppState {
|
||||
pub session_id: Option<SessionId>,
|
||||
pub session_title: String,
|
||||
pub messages: Vec<MessageView>,
|
||||
pub modal: ModalState,
|
||||
pub input: tui_textarea::TextArea<'static>,
|
||||
pub scroll_offset: u16,
|
||||
pub running: bool,
|
||||
pub model_ref: String,
|
||||
pub agent: String,
|
||||
pub quit: bool,
|
||||
pub dirty: bool,
|
||||
pub ctrl_c_pressed: bool,
|
||||
/// Width the currently cached part lines were wrapped to; a change invalidates them.
|
||||
pub render_width: u16,
|
||||
/// Accumulated session cost in USD, from `SessionCreated`/`SessionUpdated` events.
|
||||
pub session_cost: f64,
|
||||
/// Accumulated session tokens (input + output), for the status bar.
|
||||
pub session_tokens: u64,
|
||||
/// Background jobs spawned by the current session, newest activity last. Populated on
|
||||
/// session load and kept live via `JobUpdated` events; surfaced in the jobs pane.
|
||||
pub jobs: Vec<JobRecord>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
pub fn new(model_ref: String, agent: String) -> Self {
|
||||
let mut input = tui_textarea::TextArea::default();
|
||||
input.set_cursor_line_style(ratatui::style::Style::default());
|
||||
Self {
|
||||
session_id: None,
|
||||
session_title: String::new(),
|
||||
messages: Vec::new(),
|
||||
modal: ModalState::None,
|
||||
input,
|
||||
scroll_offset: 0,
|
||||
running: false,
|
||||
model_ref,
|
||||
agent,
|
||||
quit: false,
|
||||
dirty: true,
|
||||
ctrl_c_pressed: false,
|
||||
render_width: 0,
|
||||
session_cost: 0.0,
|
||||
session_tokens: 0,
|
||||
jobs: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop every part's cached render (e.g. after a resize changes the wrap width).
|
||||
pub fn invalidate_caches(&mut self) {
|
||||
for message in &mut self.messages {
|
||||
for part in &mut message.parts {
|
||||
match part {
|
||||
PartView::Text { cached_lines, .. } | PartView::Tool { cached_lines, .. } => {
|
||||
*cached_lines = None
|
||||
}
|
||||
PartView::Reasoning { .. } | PartView::StepFinish { .. } => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn set_session(&mut self, session: Session, messages: Vec<Message>, parts: Vec<Part>) {
|
||||
self.session_id = Some(session.id.clone());
|
||||
self.session_title = session.title;
|
||||
self.agent = session.agent;
|
||||
self.model_ref = format!("{}/{}", session.model.provider_id, session.model.model_id);
|
||||
self.session_cost = session.cost;
|
||||
self.session_tokens = session.usage.input + session.usage.output;
|
||||
self.messages.clear();
|
||||
// Jobs are reloaded for the newly-selected session by the caller.
|
||||
self.jobs.clear();
|
||||
|
||||
let mut messages = messages;
|
||||
messages.sort_by_key(|m| m.created_at);
|
||||
for message in messages {
|
||||
let message_parts: Vec<Part> = parts
|
||||
.iter()
|
||||
.filter(|p| p.message_id == message.id)
|
||||
.cloned()
|
||||
.collect();
|
||||
self.messages.push(message_view(message, message_parts));
|
||||
}
|
||||
self.scroll_offset = 0;
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
pub fn apply_event(&mut self, event: AppEvent) {
|
||||
match event {
|
||||
AppEvent::SessionCreated { session } | AppEvent::SessionUpdated { session } => {
|
||||
if self.session_id.as_ref() == Some(&session.id) {
|
||||
self.session_title = session.title;
|
||||
self.agent = session.agent;
|
||||
self.model_ref =
|
||||
format!("{}/{}", session.model.provider_id, session.model.model_id);
|
||||
self.session_cost = session.cost;
|
||||
self.session_tokens = session.usage.input + session.usage.output;
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
AppEvent::MessageCreated { message } => {
|
||||
if self.session_id.as_ref() == Some(&message.session_id) {
|
||||
self.messages.push(message_view(message, Vec::new()));
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
AppEvent::MessageUpdated { message } => {
|
||||
if let Some(view) = self.messages.iter_mut().find(|m| m.id == message.id) {
|
||||
view.role = message.role;
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
AppEvent::PartUpdated { part } => {
|
||||
if self.session_id.as_ref() == Some(&part.session_id) {
|
||||
self.update_or_insert_part(part);
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
AppEvent::PartDelta {
|
||||
part_id,
|
||||
message_id,
|
||||
delta,
|
||||
} => {
|
||||
self.apply_delta(part_id, message_id, delta);
|
||||
self.dirty = true;
|
||||
}
|
||||
AppEvent::RunStarted { session_id } => {
|
||||
if self.session_id.as_ref() == Some(&session_id) {
|
||||
self.running = true;
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
AppEvent::RunFinished {
|
||||
session_id,
|
||||
outcome,
|
||||
} => {
|
||||
if self.session_id.as_ref() == Some(&session_id) {
|
||||
self.running = false;
|
||||
if matches!(outcome, RunOutcome::Errored { .. }) {
|
||||
self.messages.push(MessageView {
|
||||
id: harness_core::types::MessageId::new(),
|
||||
role: Role::Assistant,
|
||||
parts: vec![PartView::Text {
|
||||
id: PartId::new(),
|
||||
text: match outcome {
|
||||
RunOutcome::Errored { message } => message,
|
||||
_ => String::new(),
|
||||
},
|
||||
cached_lines: None,
|
||||
}],
|
||||
});
|
||||
}
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
AppEvent::PermissionAsked { request } => {
|
||||
if let ModalState::Permission { requests } = &mut self.modal {
|
||||
requests.push(request);
|
||||
} else {
|
||||
self.modal = ModalState::Permission {
|
||||
requests: vec![request],
|
||||
};
|
||||
}
|
||||
self.dirty = true;
|
||||
}
|
||||
AppEvent::PermissionResolved { id } => {
|
||||
if let ModalState::Permission { requests } = &mut self.modal {
|
||||
requests.retain(|r| r.id != id);
|
||||
if requests.is_empty() {
|
||||
self.modal = ModalState::None;
|
||||
}
|
||||
}
|
||||
self.dirty = true;
|
||||
}
|
||||
AppEvent::JobUpdated { job } => {
|
||||
if let Ok(record) = serde_json::from_value::<JobRecord>(job.0) {
|
||||
if self.session_id.as_ref() == Some(&record.parent_session) {
|
||||
self.upsert_job(record);
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
AppEvent::AuthPrompt { .. } | AppEvent::ServerNotice { .. } => {
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Replaces this session's job list (e.g. after loading a session).
|
||||
pub fn set_jobs(&mut self, jobs: Vec<JobRecord>) {
|
||||
self.jobs = jobs;
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
/// Inserts or replaces a job by `task_id`, preserving list order for stable rendering.
|
||||
fn upsert_job(&mut self, record: JobRecord) {
|
||||
if let Some(existing) = self.jobs.iter_mut().find(|j| j.task_id == record.task_id) {
|
||||
*existing = record;
|
||||
} else {
|
||||
self.jobs.push(record);
|
||||
}
|
||||
}
|
||||
|
||||
/// Opens the jobs pane (cursor at the top).
|
||||
pub fn open_jobs_pane(&mut self) {
|
||||
self.modal = ModalState::JobsPane { selected: 0 };
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
/// Child session of the currently-selected job in the jobs pane, for drill-in.
|
||||
pub fn selected_job_child(&self) -> Option<SessionId> {
|
||||
if let ModalState::JobsPane { selected } = &self.modal {
|
||||
self.jobs.get(*selected).map(|j| j.child_session.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub fn scroll_up(&mut self, n: u16) {
|
||||
self.scroll_offset = self.scroll_offset.saturating_sub(n);
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
pub fn scroll_down(&mut self, n: u16) {
|
||||
self.scroll_offset = self.scroll_offset.saturating_add(n);
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
pub fn handle_permission_reply(&mut self, reply: PermissionReply, engine: &EngineHandle) {
|
||||
let request_id = if let ModalState::Permission { requests } = &self.modal {
|
||||
requests.first().map(|r| r.id.clone())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
if let Some(id) = request_id {
|
||||
engine.permission_reply(&id, reply);
|
||||
if let ModalState::Permission { requests } = &mut self.modal {
|
||||
requests.retain(|r| r.id != id);
|
||||
if requests.is_empty() {
|
||||
self.modal = ModalState::None;
|
||||
}
|
||||
}
|
||||
self.dirty = true;
|
||||
}
|
||||
}
|
||||
|
||||
pub fn open_session_picker(&mut self, sessions: Vec<Session>) {
|
||||
self.modal = ModalState::SessionPicker {
|
||||
sessions,
|
||||
selected: 0,
|
||||
};
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
pub fn close_modal(&mut self) {
|
||||
self.modal = ModalState::None;
|
||||
self.dirty = true;
|
||||
}
|
||||
|
||||
fn update_or_insert_part(&mut self, part: Part) {
|
||||
let target_id = part.id.clone();
|
||||
let Some(message_view) = self.messages.iter_mut().find(|m| m.id == part.message_id) else {
|
||||
return;
|
||||
};
|
||||
|
||||
let new_view = part_view(part);
|
||||
if let Some(existing) = message_view
|
||||
.parts
|
||||
.iter_mut()
|
||||
.find(|p| view_part_id(p) == target_id)
|
||||
{
|
||||
*existing = new_view;
|
||||
} else {
|
||||
message_view.parts.push(new_view);
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_delta(
|
||||
&mut self,
|
||||
target_id: PartId,
|
||||
message_id: harness_core::types::MessageId,
|
||||
delta: String,
|
||||
) {
|
||||
let Some(message_view) = self.messages.iter_mut().find(|m| m.id == message_id) else {
|
||||
return;
|
||||
};
|
||||
|
||||
for part in &mut message_view.parts {
|
||||
if view_part_id(part) == target_id {
|
||||
match part {
|
||||
PartView::Text {
|
||||
text, cached_lines, ..
|
||||
} => {
|
||||
text.push_str(&delta);
|
||||
*cached_lines = None;
|
||||
}
|
||||
PartView::Reasoning { text, .. } => {
|
||||
text.push_str(&delta);
|
||||
}
|
||||
PartView::Tool { cached_lines, .. } => {
|
||||
*cached_lines = None;
|
||||
}
|
||||
PartView::StepFinish { .. } => {}
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn message_view(message: Message, mut parts: Vec<Part>) -> MessageView {
|
||||
parts.sort_by_key(|p| p.idx);
|
||||
MessageView {
|
||||
id: message.id,
|
||||
role: message.role,
|
||||
parts: parts.into_iter().map(part_view).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn part_view(part: Part) -> PartView {
|
||||
match part.body {
|
||||
PartBody::Text { text, .. } => PartView::Text {
|
||||
id: part.id,
|
||||
text,
|
||||
cached_lines: None,
|
||||
},
|
||||
PartBody::Reasoning { text, .. } => PartView::Reasoning { id: part.id, text },
|
||||
PartBody::Tool { name, state, .. } => PartView::Tool {
|
||||
id: part.id,
|
||||
name,
|
||||
state: ToolStateView::from_tool_state(&state),
|
||||
cached_lines: None,
|
||||
},
|
||||
PartBody::StepStart | PartBody::StepFinish { .. } => PartView::StepFinish { id: part.id },
|
||||
PartBody::Subtask { description, .. } => PartView::Text {
|
||||
id: part.id,
|
||||
text: description,
|
||||
cached_lines: None,
|
||||
},
|
||||
PartBody::Compaction { summary, .. } => PartView::Text {
|
||||
id: part.id,
|
||||
text: format!("_Compaction summary: {summary}_"),
|
||||
cached_lines: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn view_part_id(part: &PartView) -> PartId {
|
||||
match part {
|
||||
PartView::Text { id, .. }
|
||||
| PartView::Reasoning { id, .. }
|
||||
| PartView::Tool { id, .. }
|
||||
| PartView::StepFinish { id } => id.clone(),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
use std::io;
|
||||
use std::panic;
|
||||
|
||||
use crossterm::cursor::Show;
|
||||
use crossterm::execute;
|
||||
use crossterm::terminal::{
|
||||
disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen,
|
||||
};
|
||||
|
||||
/// Restores the terminal when dropped or on panic.
|
||||
pub struct TerminalGuard;
|
||||
|
||||
impl TerminalGuard {
|
||||
/// Enables raw mode, enters the alternate screen, and installs a panic hook
|
||||
/// that restores the terminal before printing panic info.
|
||||
pub fn enter() -> io::Result<Self> {
|
||||
enable_raw_mode()?;
|
||||
let mut stdout = io::stdout();
|
||||
execute!(stdout, EnterAlternateScreen, Show)?;
|
||||
|
||||
let original_hook = panic::take_hook();
|
||||
panic::set_hook(Box::new(move |info| {
|
||||
let _ = restore();
|
||||
original_hook(info);
|
||||
}));
|
||||
|
||||
Ok(Self)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TerminalGuard {
|
||||
fn drop(&mut self) {
|
||||
let _ = restore();
|
||||
}
|
||||
}
|
||||
|
||||
fn restore() -> io::Result<()> {
|
||||
disable_raw_mode()?;
|
||||
let mut stdout = io::stdout();
|
||||
execute!(stdout, LeaveAlternateScreen, Show)
|
||||
}
|
||||
Reference in New Issue
Block a user