Files
ai-harness/crates/harness-core/src/engine/processor.rs
T
2026-07-09 07:25:14 +02:00

700 lines
23 KiB
Rust

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>>,
}
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,
})
}