Something went wrong. Try again.
A safe, simple, extensible, and fast agent harness
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611// SPDX-License-Identifier: Apache-2.0// Copyright 2026 Mason Stallmo
//! Delegation of an investigation to a nested agent.//!//! A question like "which of these forty files defines the retry policy" is//! answered by reading a great deal and keeping almost none of it. In the//! session's own context that reading fills the window and ages out the//! results an in-progress edit depends on.//!//! [`TaskTool`] runs a nested [`Agent`] over the same workspace with a//! conversation of its own, and only its final message comes back.
use std::sync::Arc;use std::sync::atomic::AtomicBool;
use async_trait::async_trait;use ein_agent::model_clients::ModelClient;use ein_agent::tools::{Param, ReadOnlyToolSet, Tool, ToolDef, ToolEffect, ToolResult, ToolSet};use ein_agent::{Agent, budget_for_context_window};use ein_plugin::model_client::{Message, Role};use ein_wasm::{ModelClientSpec, PluginRuntime};use serde::Deserialize;use tracing::{info, warn};
/// The share of the session's context window a delegated task may use.////// A task given the whole window can read until it triggers its own eviction,/// after which it reports from a summary of what it read rather than from the/// reading.const WINDOW_FRACTION: f64 = 0.5;
/// The `Task` tool's arguments.////// `deny_unknown_fields` enforces the schema's `additionalProperties: false`,/// which a provider may not.#[derive(Debug, Deserialize)]#[serde(deny_unknown_fields)]struct TaskArgs { description: String, prompt: String,}
/// The `Task` tool, which runs a read-only nested agent and returns its/// answer.pub struct TaskTool { runtime: PluginRuntime, model_spec: ModelClientSpec, /// Loaded once and shared by every call, since loading the plugin /// directory instantiates every plugin in it. This is a separate load from /// the session's own set, which holds `Task`. tools: Arc<dyn ToolSet>, /// The session's cancel flag, so stopping the turn also stops a task /// running inside it. cancel: Arc<AtomicBool>, context_window: Option<u32>,}
impl TaskTool { pub fn new( runtime: PluginRuntime, model_spec: ModelClientSpec, tools: Arc<dyn ToolSet>, cancel: Arc<AtomicBool>, context_window: Option<u32>, ) -> Self { Self { runtime, model_spec, tools, cancel, context_window, } }}
#[async_trait]impl Tool for TaskTool { fn name(&self) -> &str { "Task" }
fn schema(&self) -> ToolDef { ToolDef::function( self.name(), "Delegate an investigation to a separate agent that reads the \ workspace and reports back. Use it when answering something means \ reading a lot and keeping little: locating where a behaviour is \ implemented, working out how a subsystem fits together, checking \ whether a pattern appears anywhere else. The agent's reading stays \ in its own context and never enters yours; only its written answer \ comes back.\n\n\ Do not use it for work you can do in one or two calls — a single \ known file, a single search — where delegating costs a whole model \ turn to save nothing. Do not use it to make changes: the agent can \ only read, and cannot write, edit, run commands, or reply to the \ user. Act on what it reports yourself.\n\n\ It cannot ask you anything and you cannot steer it once it starts, \ so `prompt` has to carry everything it needs: what to look for, \ where to start, and what a useful answer contains.", ) .with_param(Param::string( "description", "A few words naming the task, shown to the user while it runs.", )) .with_param(Param::string( "prompt", "The complete instruction for the agent: what to investigate, where \ to begin, and what its answer should tell you.", )) .build() }
fn primary_arg(&self) -> Option<&str> { Some("description") }
/// Returns `ReadOnly`, which its tool set enforces, so a batch of /// delegations runs concurrently. fn effect(&self) -> ToolEffect { ToolEffect::ReadOnly }
async fn call(&self, id: &str, args: &str) -> anyhow::Result<ToolResult> { let args: TaskArgs = serde_json::from_str(args)?;
info!("[task] delegating: {}", args.description);
let model = self .runtime .model_clients() .new_session(&self.model_spec) .await?;
let answer = delegate( model, self.tools.clone(), self.cancel.clone(), delegate_window(self.context_window), args.prompt, ) .await?;
Ok(report(id, &args.description, &answer)) }}
/// Runs one delegated turn and returns what it wrote.////// Generic rather than taking a [`TaskTool`], so it can be tested without a/// WASM runtime or a provider.////// The agent sees only the read-only tools in `tools` and has no event handler/// or message sink. Its events would otherwise reach the session's stream/// looking like the session's own: its text drawn as the session's reply, and/// its usage overwriting the parent's count.async fn delegate<MC: ModelClient, TS: ToolSet>( model: MC, tools: TS, cancel: Arc<AtomicBool>, window: Option<u32>, prompt: String,) -> anyhow::Result<String> { let mut builder = Agent::builder_with_tool_set(model, ReadOnlyToolSet::new(tools)) .with_system_preamble(vec![Message { role: Role::System, content: Some(preamble()), tool_calls: None, tool_call_id: None, }]) .with_cancel_handle(cancel);
if let Some(window) = window { builder = builder .with_context_window(window) .tool_result_budget(budget_for_context_window(window)); }
Ok(builder .build() .chat(prompt) .await? .content .unwrap_or_default())}
/// Returns the window a delegated task gets, [`WINDOW_FRACTION`] of the/// session's and at least 1.fn delegate_window(session: Option<u32>) -> Option<u32> { session.map(|window| ((f64::from(window) * WINDOW_FRACTION) as u32).max(1))}
/// Returns the task's answer as the delegating agent reads it.////// A task that wrote nothing is reported as that. An empty result reads as a/// tool that silently did nothing, which suggests calling it again.fn report(id: &str, description: &str, answer: &str) -> ToolResult { if !answer.trim().is_empty() { return ToolResult::new(id, answer.to_owned()); }
warn!("[task] '{description}' produced no answer");
ToolResult::new( id, format!( "The delegated agent for `{description}` ended without writing an \ answer. It may have spent its tool calls before concluding, or the \ turn may have been cancelled. Narrow the task, or investigate it \ directly." ), )}
/// Returns the delegated agent's system preamble.////// The session's preamble does not apply: the task has no core memory to write/// and no user to address. It is told what it cannot infer, that its context/// is discarded and only its reply survives.fn preamble() -> String { "You are a sub-agent. Another agent has delegated one investigation to you \ and is waiting on your answer.\n\ \n\ You can only read: the file and search tools, and nothing that changes \ anything. You cannot run commands, edit files, or reply to the user \ directly. Do not try to; report what you found instead.\n\ \n\ Your reply is the only thing that survives. Everything you read is \ discarded when you finish, and the agent waiting on you sees none of it — \ so a reply saying where to look, without saying what is there, is of no \ use to it. Quote the file paths, line numbers, symbol names and error \ text you found, exactly as they appear. An approximation of a path is \ worse than leaving it out, because it will be acted on.\n\ \n\ Answer in one message when you are done reading. Say plainly what you did \ not manage to establish rather than filling it in." .to_string()}
#[cfg(test)]mod tests { use super::*; use std::sync::Mutex; use std::sync::atomic::Ordering;
use ein_agent::tools::NativeToolSet; use ein_plugin::model_client::{ Choice, CompletionResponse, FinishReason, FunctionCall, ToolCall, };
/// A tool that counts its calls, so a test can assert it never ran. struct Recorder { name: &'static str, effect: ToolEffect, calls: Arc<Mutex<usize>>, }
impl Recorder { fn new(name: &'static str, effect: ToolEffect) -> Self { Self { name, effect, calls: Arc::new(Mutex::new(0)), } } }
#[async_trait] impl Tool for Recorder { fn name(&self) -> &str { self.name }
fn schema(&self) -> ToolDef { ToolDef::function(self.name, "a tool").build() }
fn effect(&self) -> ToolEffect { self.effect }
async fn call(&self, id: &str, _args: &str) -> anyhow::Result<ToolResult> { *self.calls.lock().unwrap() += 1;
Ok(ToolResult::new(id, format!("{} ran", self.name))) } }
/// A model client that plays a fixed script of turns and keeps the /// requests it was sent. struct Script { turns: Mutex<Vec<CompletionResponse>>, seen: Arc<Mutex<Vec<Vec<Message>>>>, }
impl Script { fn new(turns: Vec<CompletionResponse>) -> Self { Self { turns: Mutex::new(turns), seen: Arc::new(Mutex::new(Vec::new())), } } }
#[async_trait] impl ModelClient for Script { async fn complete( &mut self, messages: &[Message], _tools: &[ToolDef], ) -> anyhow::Result<CompletionResponse> { self.seen.lock().unwrap().push(messages.to_vec());
let mut turns = self.turns.lock().unwrap();
Ok(match turns.is_empty() { true => says("done"), false => turns.remove(0), }) } }
fn says(text: &str) -> CompletionResponse { CompletionResponse { choices: vec![Choice { index: None, finish_reason: FinishReason::Stop, message: Message { role: Role::Assistant, content: Some(text.to_owned()), tool_calls: None, tool_call_id: None, }, }], usage: None, error: None, provenance: None, } }
fn calls(name: &str) -> CompletionResponse { CompletionResponse { choices: vec![Choice { index: None, finish_reason: FinishReason::ToolCalls, message: Message { role: Role::Assistant, content: None, tool_calls: Some(vec![ToolCall::Function { id: "1".to_owned(), index: 0, function: FunctionCall { name: name.to_owned(), arguments: "{}".to_owned(), }, }]), tool_call_id: None, }, }], usage: None, error: None, provenance: None, } }
fn quiet() -> Arc<AtomicBool> { Arc::new(AtomicBool::new(false)) }
/// A delegated agent given a mutating tool cannot call it. /// /// Its changes would be ones the delegating agent never saw. #[tokio::test] async fn a_delegated_agent_cannot_reach_a_tool_that_changes_anything() { let editor = Recorder::new("Edit", ToolEffect::Mutating); let edits = editor.calls.clone(); let reader = Recorder::new("Read", ToolEffect::ReadOnly); let reads = reader.calls.clone();
let mut tools = NativeToolSet::default(); tools.insert(editor); tools.insert(reader);
let answer = delegate( Script::new(vec![calls("Read"), calls("Edit"), says("what I found")]), tools, quiet(), None, "look into it".to_owned(), ) .await .unwrap();
assert_eq!(*reads.lock().unwrap(), 1); assert_eq!(*edits.lock().unwrap(), 0, "a delegated task edited a file"); assert_eq!(answer, "what I found"); }
/// Only the task's final message comes back, not its tool results. #[tokio::test] async fn only_the_answer_crosses_back_not_the_reading() { let reader = Recorder::new("Read", ToolEffect::ReadOnly);
let mut tools = NativeToolSet::default(); tools.insert(reader);
let answer = delegate( Script::new(vec![ calls("Read"), calls("Read"), says("the retry policy is in retry.rs"), ]), tools, quiet(), None, "where is the retry policy".to_owned(), ) .await .unwrap();
assert_eq!(answer, "the retry policy is in retry.rs"); assert!(!answer.contains("ran"), "a tool result came back with it"); }
/// A tool that answers by delegating, so the parent's events can be /// watched. struct Delegating { answer: &'static str, }
#[async_trait] impl Tool for Delegating { fn name(&self) -> &str { "Task" }
fn schema(&self) -> ToolDef { ToolDef::function(self.name(), "delegate").build() }
fn effect(&self) -> ToolEffect { ToolEffect::Unknown }
async fn call(&self, id: &str, _args: &str) -> anyhow::Result<ToolResult> { let answer = delegate( Script::new(vec![says(self.answer)]), NativeToolSet::default(), quiet(), None, "look into it".to_owned(), ) .await?;
Ok(ToolResult::new(id, answer)) } }
/// Nothing a delegated agent writes reaches the session's event stream. /// /// `delegate` installing no event handler is all that keeps a task's text /// from being drawn as the session's reply. #[tokio::test] async fn a_delegated_turn_says_nothing_to_the_session() { const NESTED: &str = "the retry policy is in retry.rs";
let mut tools = NativeToolSet::default(); tools.insert(Delegating { answer: NESTED });
let parent = Script::new(vec![calls("Task"), says("done")]); let seen = parent.seen.clone();
let heard = Arc::new(Mutex::new(Vec::new())); let collect = heard.clone();
let answer = Agent::builder_with_tool_set(parent, tools) .with_event_handler(move |event| { let heard = collect.clone();
async move { heard.lock().unwrap().push(event); } }) .build() .chat("where is the retry policy".to_owned()) .await .unwrap() .content .unwrap_or_default();
for event in heard.lock().unwrap().iter() { if let ein_agent::AgentEvent::ContentDelta(text) | ein_agent::AgentEvent::ReasoningDelta(text) = event { assert!( !text.contains(NESTED), "a delegated turn's text reached the session's stream: {text}" ); } }
assert_eq!(answer, "done", "the session answered for itself");
// The answer did reach the parent, so the check above did not pass on // a task that produced nothing. let last = seen.lock().unwrap().last().cloned().unwrap_or_default(); assert!( last.iter() .any(|m| m.content.as_deref().unwrap_or_default().contains(NESTED)), "the task's answer never reached the agent that asked for it", ); }
/// A cancel raised before the task starts stops it, and stays raised. /// /// An agent clearing a flag it was given would swallow a cancel raised /// while the task started, and neither agent would stop. #[tokio::test] async fn a_cancel_raised_before_the_task_starts_still_stops_it() { let reader = Recorder::new("Read", ToolEffect::ReadOnly); let reads = reader.calls.clone();
let mut tools = NativeToolSet::default(); tools.insert(reader);
let cancel = Arc::new(AtomicBool::new(true));
delegate( Script::new(vec![calls("Read"), says("done")]), tools, cancel.clone(), None, "look into it".to_owned(), ) .await .unwrap();
assert_eq!(*reads.lock().unwrap(), 0, "the task ran after a cancel"); assert!( cancel.load(Ordering::Relaxed), "the delegated agent cleared the session's cancel flag" ); }
/// The delegated agent's preamble tells it its context is discarded. /// /// Otherwise it answers by pointing at what it read, which the delegating /// agent cannot see. #[tokio::test] async fn the_delegated_agent_is_told_its_context_is_discarded() { let script = Script::new(vec![says("ok")]); let seen = script.seen.clone();
delegate( script, NativeToolSet::default(), quiet(), None, "go".to_owned(), ) .await .unwrap();
let first = seen.lock().unwrap()[0][0].clone();
assert_eq!(first.role, Role::System); assert!( first.content.as_deref().unwrap().contains("discarded"), "{first:?}" ); }
/// A delegated task gets half the session's window. #[test] fn a_delegated_task_gets_less_of_the_window_than_the_session() { assert_eq!(delegate_window(Some(200_000)), Some(100_000)); assert_eq!(delegate_window(None), None); }
/// A window that would round to zero is kept at 1. /// /// Zero disables pressure management. #[test] fn a_tiny_window_still_leaves_something_to_manage() { assert_eq!(delegate_window(Some(1)), Some(1)); }
/// A task that wrote nothing is reported as that, not as an empty result. /// /// An empty result suggests calling the tool again, which costs another /// delegated turn. #[test] fn a_task_that_wrote_no_answer_says_so_rather_than_coming_back_empty() { for answer in ["", " \n "] { let result = report("1", "find the retry policy", answer);
assert!( result.content.contains("without writing an answer"), "{:?}", result.content ); assert!(result.content.contains("find the retry policy")); } }
#[test] fn an_answer_comes_back_as_it_was_written() { assert_eq!(report("1", "d", "the answer").content, "the answer"); }}