use std::marker::PhantomData;
use std::sync::Arc;
use serde_json::Value;
use crate::agent::{AgentOutput, AgentRunner};
use crate::error::Error;
use super::ctx::WorkflowCtx;
use super::event::WorkflowEvent;
use super::journal;
const DEFAULT_TOOL_TURNS: usize = 10;
pub trait StructuredSchema: serde::de::DeserializeOwned + Send {
fn json_schema() -> Value;
}
pub struct NoSchema;
pub struct RawJson;
#[derive(Clone, Default)]
#[non_exhaustive]
pub struct AgentOpts {
pub label: Option<String>,
pub phase: Option<String>,
pub schema: Option<Value>,
pub tools: Option<Vec<Arc<dyn crate::tool::Tool>>>,
pub goal: Option<crate::agent::goal::GoalCondition>,
pub model: Option<String>,
pub isolation: Isolation,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum Isolation {
#[default]
None,
Worktree,
}
pub struct AgentCall<T = NoSchema> {
ctx: WorkflowCtx,
prompt: String,
opts: AgentOpts,
phase_snapshot: Option<Arc<str>>,
_t: PhantomData<fn() -> T>,
}
pub fn agent(ctx: &WorkflowCtx, prompt: impl Into<String>) -> AgentCall<NoSchema> {
let phase_snapshot = ctx.current_phase();
AgentCall {
ctx: ctx.clone(),
prompt: prompt.into(),
opts: AgentOpts::default(),
phase_snapshot,
_t: PhantomData,
}
}
impl<T> AgentCall<T> {
pub fn label(mut self, label: impl Into<String>) -> Self {
self.opts.label = Some(label.into());
self
}
pub fn phase(mut self, phase: impl Into<String>) -> Self {
self.opts.phase = Some(phase.into());
self
}
pub fn model(mut self, model: impl Into<String>) -> Self {
self.opts.model = Some(model.into());
self
}
pub fn isolation(mut self, isolation: Isolation) -> Self {
self.opts.isolation = isolation;
self
}
pub fn tools(mut self, tools: Vec<Arc<dyn crate::tool::Tool>>) -> Self {
self.opts.tools = Some(tools);
self
}
pub fn goal(mut self, goal: crate::agent::goal::GoalCondition) -> Self {
self.opts.goal = Some(goal);
self
}
fn effective_phase(&self) -> Option<String> {
self.opts
.phase
.clone()
.or_else(|| self.phase_snapshot.as_deref().map(str::to_owned))
}
async fn run_leaf(self) -> Result<Option<AgentOutput>, Error> {
let label = self
.opts
.label
.clone()
.unwrap_or_else(|| "agent".to_string());
if self.ctx.is_cancelled() {
self.ctx.emit(WorkflowEvent::AgentSkipped { label });
return Ok(None);
}
if self.opts.isolation == Isolation::None
&& let Some(journal) = self.ctx.journal_arc()
{
let hash = journal::content_hash(
&self.prompt,
self.opts.model.as_deref(),
self.opts.schema.as_ref(),
);
let occurrence = journal.next_occurrence(&hash);
let key = journal::CallKey {
content_hash: hash,
occurrence,
};
if let Some(cached) = journal.lookup(&key) {
self.ctx.emit(WorkflowEvent::AgentReplayed {
label,
usage: cached.tokens_used,
});
return Ok(Some(cached));
}
let result = self.run_live(label).await?;
if let Some(ref output) = result {
journal.append(&key, output)?;
}
return Ok(result);
}
self.run_live(label).await
}
async fn run_live(self, label: String) -> Result<Option<AgentOutput>, Error> {
let phase = self.effective_phase();
let _permit = self
.ctx
.semaphore()
.acquire_owned()
.await
.map_err(|_| Error::Agent("flow concurrency limiter closed".to_string()))?;
self.ctx.register_agent()?;
self.ctx.admit_budget()?;
self.ctx.emit(WorkflowEvent::AgentStarted {
label: label.clone(),
phase,
});
let guard = match self.opts.isolation {
Isolation::None => None,
Isolation::Worktree => {
let root = self.ctx.workspace().ok_or_else(|| {
Error::Config(
"isolation: Worktree requires WorkflowCtxBuilder::workspace(repo root)"
.into(),
)
})?;
let name = format!(
"hb-{}-{}",
super::worktree::sanitize_name(&label),
self.ctx.spawned()
);
Some(super::worktree::WorktreeGuard::create(&root, &name).await?)
}
};
let workspace_override = guard.as_ref().map(|g| g.path().to_path_buf());
let output = tokio::select! {
biased;
_ = self.ctx.cancel_token().cancelled() => {
self.ctx.emit(WorkflowEvent::AgentSkipped { label });
if let Some(g) = guard {
let _ = g.cleanup().await;
}
return Ok(None);
}
result = run_one(&self.ctx, &self.prompt, &self.opts, workspace_override.as_deref()) => {
if let Some(g) = guard {
match g.cleanup().await {
Ok(super::worktree::Disposition::Kept { branch }) => {
self.ctx.emit(WorkflowEvent::LogLine {
msg: format!(
"worktree of '{label}' had changes — preserved on branch {branch}"
),
});
}
Ok(super::worktree::Disposition::Pruned) => {}
Err(e) => {
self.ctx.emit(WorkflowEvent::LogLine {
msg: format!("worktree cleanup for '{label}' failed: {e}"),
});
}
}
}
result?
}
};
self.ctx.record_spend(&output.tokens_used);
self.ctx.emit(WorkflowEvent::AgentFinished {
label,
usage: output.tokens_used,
});
Ok(Some(output))
}
}
impl AgentCall<NoSchema> {
pub fn schema<S: StructuredSchema>(self) -> AgentCall<S> {
let mut opts = self.opts;
opts.schema = Some(S::json_schema());
AgentCall {
ctx: self.ctx,
prompt: self.prompt,
opts,
phase_snapshot: self.phase_snapshot,
_t: PhantomData,
}
}
pub fn schema_value(self, schema: Value) -> AgentCall<RawJson> {
let mut opts = self.opts;
opts.schema = Some(schema);
AgentCall {
ctx: self.ctx,
prompt: self.prompt,
opts,
phase_snapshot: self.phase_snapshot,
_t: PhantomData,
}
}
pub async fn run(self) -> Result<Option<String>, Error> {
Ok(self.run_leaf().await?.map(|o| o.result))
}
}
impl<T: StructuredSchema> AgentCall<T> {
pub async fn run(self) -> Result<Option<T>, Error> {
match self.run_leaf().await? {
None => Ok(None),
Some(output) => {
let value = output.structured.ok_or_else(|| {
Error::Agent("agent finished without structured output".to_string())
})?;
let typed = serde_json::from_value::<T>(value).map_err(Error::Json)?;
Ok(Some(typed))
}
}
}
}
impl AgentCall<RawJson> {
pub async fn run(self) -> Result<Option<Value>, Error> {
match self.run_leaf().await? {
None => Ok(None),
Some(output) => {
let value = output.structured.ok_or_else(|| {
Error::Agent("agent finished without structured output".to_string())
})?;
Ok(Some(value))
}
}
}
}
async fn run_one(
ctx: &WorkflowCtx,
prompt: &str,
opts: &AgentOpts,
workspace_override: Option<&std::path::Path>,
) -> Result<AgentOutput, Error> {
let provider = match &opts.model {
Some(role) => match ctx.provider_factory() {
Some(factory) => factory(role)?,
None => {
ctx.emit(super::event::WorkflowEvent::LogLine {
msg: format!(
"model '{role}' requested but no provider factory installed — \
using the default provider"
),
});
ctx.provider()
}
},
None => ctx.provider(),
};
let mut builder = AgentRunner::builder(provider);
if let Some(label) = &opts.label {
builder = builder.name(label.clone());
}
if let Some(ws) = workspace_override {
builder = builder.workspace(ws.to_path_buf());
}
if let Some(sink) = ctx.agent_events() {
builder = builder.on_event(sink);
}
if let Some(schema) = &opts.schema {
builder = builder.structured_schema(schema.clone());
}
if let Some(tools) = opts.tools.clone().or_else(|| ctx.base_tools()) {
builder = builder.tools(tools).max_turns(DEFAULT_TOOL_TURNS);
}
if let Some(goal) = &opts.goal {
let needed = (goal.max_continuations() as usize).saturating_add(1);
builder = builder
.goal(goal.clone())
.max_turns(needed.max(DEFAULT_TOOL_TURNS));
let remaining = ctx.remaining();
if remaining != u64::MAX {
builder = builder.max_total_tokens(remaining);
}
}
let runner = builder.build()?;
runner.execute(prompt).await
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
use super::*;
use crate::agent::test_helpers::MockProvider;
use crate::error::Error;
use crate::llm::types::TokenUsage;
use crate::llm::types::{CompletionRequest, CompletionResponse, ContentBlock, StopReason};
use crate::llm::{BoxedProvider, LlmProvider};
use super::super::parallel::{BoxThunk, parallel, thunk};
struct ConcurrencyTrackingProvider {
current: Arc<AtomicUsize>,
peak: Arc<AtomicUsize>,
}
impl LlmProvider for ConcurrencyTrackingProvider {
async fn complete(&self, _request: CompletionRequest) -> Result<CompletionResponse, Error> {
let now = self.current.fetch_add(1, Ordering::SeqCst) + 1;
self.peak.fetch_max(now, Ordering::SeqCst);
tokio::time::sleep(Duration::from_millis(50)).await;
self.current.fetch_sub(1, Ordering::SeqCst);
Ok(CompletionResponse {
content: vec![ContentBlock::Text {
text: "done".into(),
}],
stop_reason: StopReason::EndTurn,
reasoning: None,
usage: TokenUsage {
input_tokens: 1,
output_tokens: 1,
..Default::default()
},
model: None,
})
}
fn model_name(&self) -> Option<&str> {
Some("concurrency-mock")
}
}
#[tokio::test]
async fn text_terminal_returns_agent_output() {
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"mock text",
10,
5,
)]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.build()
.expect("build ctx");
let out = agent(&ctx, "do the thing").run().await.expect("run ok");
assert_eq!(out.as_deref(), Some("mock text"));
assert_eq!(mock.captured_requests.lock().expect("lock").len(), 1);
}
#[tokio::test]
async fn agent_event_sink_receives_inner_runner_events() {
use crate::agent::events::{AgentEvent, OnEvent};
use std::sync::Mutex;
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"hi", 1, 1,
)]));
let captured: Arc<Mutex<Vec<AgentEvent>>> = Arc::new(Mutex::new(Vec::new()));
let sink: Arc<OnEvent> = {
let captured = Arc::clone(&captured);
Arc::new(move |ev| captured.lock().expect("lock").push(ev))
};
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(mock)))
.on_agent_event(sink)
.build()
.expect("build ctx");
let out = agent(&ctx, "do it")
.label("lane:1")
.run()
.await
.expect("run ok");
assert_eq!(out.as_deref(), Some("hi"));
let events = captured.lock().expect("lock");
assert!(
events
.iter()
.any(|e| matches!(e, AgentEvent::RunStarted { agent, .. } if agent == "lane:1")),
"inner runner's RunStarted must reach the sink with the label as agent name"
);
assert!(
events
.iter()
.any(|e| matches!(e, AgentEvent::RunCompleted { agent, .. } if agent == "lane:1")),
"inner runner's RunCompleted must reach the sink"
);
}
#[tokio::test]
async fn per_call_model_resolves_through_the_factory() {
let default_mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"from-default",
1,
1,
)]));
let fast_mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"from-fast",
1,
1,
)]));
let fast_for_factory = Arc::clone(&fast_mock);
let factory: Arc<crate::agent::flow::ProviderFactory> = Arc::new(move |role: &str| {
assert_eq!(role, "fast");
Ok(Arc::new(BoxedProvider::from_arc(Arc::clone(
&fast_for_factory,
))))
});
let ctx =
WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&default_mock))))
.provider_factory(factory)
.build()
.expect("ctx");
let out = agent(&ctx, "go").model("fast").run().await.expect("run");
assert_eq!(out.as_deref(), Some("from-fast"));
assert_eq!(
fast_mock.captured_requests.lock().expect("lock").len(),
1,
"the fast provider must take the call"
);
assert!(
default_mock
.captured_requests
.lock()
.expect("lock")
.is_empty(),
"the default provider must NOT be used"
);
}
#[tokio::test]
async fn model_without_factory_degrades_to_default_with_a_log() {
use std::sync::Mutex;
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"ok", 1, 1,
)]));
let logs: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let sink = {
let logs = Arc::clone(&logs);
Arc::new(move |ev: super::super::event::WorkflowEvent| {
if let super::super::event::WorkflowEvent::LogLine { msg } = ev {
logs.lock().expect("lock").push(msg);
}
})
};
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.on_event(sink)
.build()
.expect("ctx");
let out = agent(&ctx, "go").model("fast").run().await.expect("run");
assert_eq!(out.as_deref(), Some("ok"), "degraded, never fatal");
let logs = logs.lock().expect("lock");
assert!(
logs.iter()
.any(|l| l.contains("fast") && l.contains("factory")),
"must log the degradation: {logs:?}"
);
}
async fn fixture_repo() -> (tempfile::TempDir, std::path::PathBuf) {
let dir = tempfile::tempdir().unwrap();
let root = dir.path().to_path_buf();
for args in [
vec!["init", "-q"],
vec!["config", "user.name", "t"],
vec!["config", "user.email", "t@t"],
vec!["commit", "--allow-empty", "-q", "-m", "init"],
] {
let ok = std::process::Command::new("git")
.current_dir(&root)
.args(&args)
.status()
.unwrap()
.success();
assert!(ok, "fixture git {args:?}");
}
(dir, root)
}
#[tokio::test(flavor = "multi_thread")]
async fn worktree_isolation_runs_the_leaf_in_its_own_worktree() {
let (_keep, root) = fixture_repo().await;
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"done", 1, 1,
)]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.workspace(root.clone())
.build()
.expect("ctx");
let out = agent(&ctx, "mutate files")
.label("iso-1")
.isolation(super::Isolation::Worktree)
.run()
.await
.expect("run ok");
assert_eq!(out.as_deref(), Some("done"));
let reqs = mock.captured_requests.lock().expect("lock");
let sys = &reqs[0].system;
assert!(
sys.contains("heartbit-worktrees")
&& !sys.contains(&root.to_string_lossy().to_string()),
"leaf must run against the worktree, got system: {sys}"
);
}
#[tokio::test]
async fn worktree_without_ctx_workspace_is_a_config_error() {
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"x", 1, 1,
)]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(mock)))
.build()
.expect("ctx");
let err = agent(&ctx, "go")
.isolation(super::Isolation::Worktree)
.run()
.await
.unwrap_err();
assert!(
err.to_string().contains("workspace"),
"must name the missing seam: {err}"
);
}
#[tokio::test(flavor = "multi_thread")]
async fn journal_refuses_isolated_calls() {
let (_keep, root) = fixture_repo().await;
let dir = tempfile::tempdir().unwrap();
let jpath = dir.path().join("j.jsonl");
let mock = Arc::new(MockProvider::new(vec![
MockProvider::text_response("a", 1, 1),
MockProvider::text_response("b", 1, 1),
]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.workspace(root)
.journal(&jpath, super::super::journal::ResumeMode::Resume)
.expect("journal")
.build()
.expect("ctx");
for _ in 0..2 {
let _ = agent(&ctx, "same prompt")
.label("iso-j")
.isolation(super::Isolation::Worktree)
.run()
.await
.expect("run");
}
assert_eq!(
mock.captured_requests.lock().expect("lock").len(),
2,
"isolated calls must NEVER replay from the journal (side effects \
are not restored)"
);
}
#[tokio::test]
async fn journal_hash_includes_the_model() {
let dir = tempfile::tempdir().unwrap();
let jpath = dir.path().join("j.jsonl");
let mock = Arc::new(MockProvider::new(vec![
MockProvider::text_response("a", 1, 1),
MockProvider::text_response("b", 1, 1),
]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.journal(&jpath, super::super::journal::ResumeMode::Resume)
.expect("journal")
.build()
.expect("ctx");
let one = agent(&ctx, "same").run().await.expect("run");
let two = agent(&ctx, "same")
.model("other-model")
.run()
.await
.expect("run");
assert_eq!(one.as_deref(), Some("a"));
assert_eq!(
two.as_deref(),
Some("b"),
"a different model must be a journal MISS, not a replay of 'a'"
);
}
#[tokio::test]
async fn no_agent_event_sink_means_no_events_and_no_failure() {
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"ok", 1, 1,
)]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(mock)))
.build()
.expect("build ctx");
let out = agent(&ctx, "go").run().await.expect("run ok");
assert_eq!(out.as_deref(), Some("ok"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrency_cap_is_respected_and_used() {
let current = Arc::new(AtomicUsize::new(0));
let peak = Arc::new(AtomicUsize::new(0));
let provider = ConcurrencyTrackingProvider {
current: Arc::clone(¤t),
peak: Arc::clone(&peak),
};
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::new(provider)))
.max_concurrency(2)
.max_agents(1000)
.build()
.expect("build ctx");
let thunks: Vec<BoxThunk<String>> = (0..10)
.map(|i| {
let ctx = ctx.clone();
thunk(move || async move {
Ok(agent(&ctx, format!("task {i}"))
.run()
.await?
.unwrap_or_default())
})
})
.collect();
let out = parallel(&ctx, thunks).await;
assert_eq!(out.iter().filter(|o| o.is_some()).count(), 10);
let observed = peak.load(Ordering::SeqCst);
assert!(observed <= 2, "peak {observed} exceeded cap 2");
assert_eq!(
observed, 2,
"cap should actually be reached, peak was {observed}"
);
}
#[tokio::test]
async fn backstop_rejects_excess_agents() {
let mock = Arc::new(MockProvider::new(vec![
MockProvider::text_response("a", 1, 1),
MockProvider::text_response("b", 1, 1),
]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.max_agents(2)
.build()
.expect("build ctx");
assert!(agent(&ctx, "1").run().await.is_ok());
assert!(agent(&ctx, "2").run().await.is_ok());
let third = agent(&ctx, "3").run().await;
match third {
Err(Error::AgentBudgetExceeded { limit }) => assert_eq!(limit, 2),
other => panic!("expected AgentBudgetExceeded {{ limit: 2 }}, got {other:?}"),
}
assert_eq!(mock.captured_requests.lock().expect("lock").len(), 2);
}
#[tokio::test]
async fn cancelled_run_skips_without_invoking_provider() {
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"should not run",
1,
1,
)]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.build()
.expect("build ctx");
ctx.cancellation_token().cancel();
let out = agent(&ctx, "task").run().await.expect("run ok");
assert!(out.is_none(), "cancelled run must yield Ok(None)");
assert!(
mock.captured_requests.lock().expect("lock").is_empty(),
"provider must not be invoked when cancelled"
);
}
use super::super::ctx::ControlBreach;
use super::super::parallel::{parallel as flow_parallel, thunk as flow_thunk};
#[tokio::test]
async fn leaf_records_weighted_cost_into_budget() {
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"x", 10, 5,
)]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(mock)))
.build()
.expect("build ctx");
assert_eq!(ctx.budget().spent(), 0);
agent(&ctx, "task").run().await.expect("run ok");
assert_eq!(ctx.budget().spent(), 15);
}
#[tokio::test]
async fn sequential_budget_ceiling_rejects_after_exhaustion() {
let mock = Arc::new(MockProvider::new(vec![
MockProvider::text_response("a", 10, 0),
MockProvider::text_response("b", 10, 0),
MockProvider::text_response("c", 10, 0),
]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.budget(25)
.build()
.expect("build ctx");
assert!(agent(&ctx, "1").run().await.is_ok()); assert!(agent(&ctx, "2").run().await.is_ok()); assert!(agent(&ctx, "3").run().await.is_ok()); let fourth = agent(&ctx, "4").run().await; match fourth {
Err(Error::BudgetExceeded { used, limit }) => {
assert_eq!(used, 30);
assert_eq!(limit, 25);
}
other => panic!("expected BudgetExceeded, got {other:?}"),
}
assert_eq!(mock.captured_requests.lock().expect("lock").len(), 3);
}
#[tokio::test]
async fn unbounded_budget_leaves_runs_unthrottled() {
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"ok", 1_000, 1_000,
)]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(mock)))
.build() .expect("build ctx");
assert!(agent(&ctx, "task").run().await.is_ok());
assert_eq!(ctx.budget().spent(), 2_000);
assert_eq!(ctx.remaining(), u64::MAX);
assert!(!ctx.is_cancelled());
assert!(ctx.control_breach().is_none());
}
#[tokio::test]
async fn budget_breach_in_parallel_is_sticky_and_trips_cancellation() {
let mock = Arc::new(MockProvider::new(vec![
MockProvider::text_response("x", 1, 1),
MockProvider::text_response("y", 1, 1),
]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(Arc::clone(&mock))))
.budget(1)
.build()
.expect("build ctx");
ctx.budget().record(&crate::llm::types::TokenUsage {
input_tokens: 5,
..Default::default()
});
let thunks: Vec<BoxThunk<Option<String>>> = (0..2)
.map(|i| {
let ctx = ctx.clone();
flow_thunk(move || async move { agent(&ctx, format!("p{i}")).run().await })
})
.collect();
let out = flow_parallel(&ctx, thunks).await;
assert!(
out.iter().all(|slot| !matches!(slot, Some(Some(_)))),
"a budget-exhausted run must yield no real output, got {out:?}"
);
assert!(
ctx.is_cancelled(),
"a control breach must fire run-wide cancel"
);
assert!(matches!(
ctx.control_breach(),
Some(ControlBreach::Budget { limit: 1, .. })
));
assert!(mock.captured_requests.lock().expect("lock").is_empty());
}
#[tokio::test]
async fn agent_domain_error_in_parallel_is_not_sticky() {
let mock = Arc::new(MockProvider::new(vec![]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(mock)))
.build()
.expect("build ctx");
let thunks: Vec<BoxThunk<String>> = (0..2)
.map(|i| {
let ctx = ctx.clone();
flow_thunk(move || async move {
Ok(agent(&ctx, format!("p{i}"))
.run()
.await?
.unwrap_or_default())
})
})
.collect();
let out = flow_parallel(&ctx, thunks).await;
assert!(out.iter().all(Option::is_none));
assert!(!ctx.is_cancelled(), "a domain error must NOT fire cancel");
assert!(ctx.control_breach().is_none());
}
struct FixedCostProvider {
input: u32,
output: u32,
}
impl LlmProvider for FixedCostProvider {
async fn complete(&self, _request: CompletionRequest) -> Result<CompletionResponse, Error> {
Ok(CompletionResponse {
content: vec![ContentBlock::Text { text: "ok".into() }],
stop_reason: StopReason::EndTurn,
reasoning: None,
usage: TokenUsage {
input_tokens: self.input,
output_tokens: self.output,
..Default::default()
},
model: None,
})
}
fn model_name(&self) -> Option<&str> {
Some("fixed-cost-mock")
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_spend_accounting_loses_nothing() {
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::new(FixedCostProvider {
input: 1,
output: 1,
})))
.max_concurrency(8)
.build()
.expect("build ctx");
let thunks: Vec<BoxThunk<String>> = (0..64)
.map(|i| {
let ctx = ctx.clone();
thunk(move || async move {
Ok(agent(&ctx, format!("t{i}"))
.run()
.await?
.unwrap_or_default())
})
})
.collect();
let out = parallel(&ctx, thunks).await;
assert_eq!(out.iter().filter(|o| o.is_some()).count(), 64);
assert_eq!(ctx.budget().spent(), 128, "atomic accounting lost a record");
}
#[tokio::test]
async fn loop_until_budget_terminates() {
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::new(FixedCostProvider {
input: 10,
output: 0,
})))
.budget(100)
.build()
.expect("build ctx");
let mut ran = 0u32;
while ctx.remaining() >= 10 {
agent(&ctx, "work").run().await.expect("run ok");
ran += 1;
assert!(ran <= 100, "loop-until-budget failed to terminate");
}
assert_eq!(ran, 10);
assert_eq!(ctx.budget().spent(), 100);
assert_eq!(ctx.remaining(), 0);
}
use super::{RawJson, StructuredSchema};
use crate::llm::types::RESPOND_TOOL_NAME;
struct RespondProvider {
payload: serde_json::Value,
}
impl LlmProvider for RespondProvider {
async fn complete(&self, _request: CompletionRequest) -> Result<CompletionResponse, Error> {
Ok(CompletionResponse {
content: vec![ContentBlock::ToolUse {
id: "resp-1".into(),
name: RESPOND_TOOL_NAME.into(),
input: self.payload.clone(),
}],
stop_reason: StopReason::ToolUse,
reasoning: None,
usage: TokenUsage {
input_tokens: 5,
output_tokens: 5,
..Default::default()
},
model: None,
})
}
fn model_name(&self) -> Option<&str> {
Some("respond-mock")
}
}
#[derive(serde::Deserialize, Debug, PartialEq)]
struct Finding {
title: String,
severity: u8,
}
impl StructuredSchema for Finding {
fn json_schema() -> serde_json::Value {
serde_json::json!({
"type": "object",
"required": ["title", "severity"],
"properties": {
"title": { "type": "string" },
"severity": { "type": "integer" }
}
})
}
}
fn respond_ctx(payload: serde_json::Value) -> WorkflowCtx {
WorkflowCtx::builder(Arc::new(BoxedProvider::new(RespondProvider { payload })))
.build()
.expect("build ctx")
}
#[tokio::test]
async fn typed_schema_round_trips() {
let ctx = respond_ctx(serde_json::json!({ "title": "SQLi", "severity": 9 }));
let found: Option<Finding> = agent(&ctx, "audit")
.schema::<Finding>()
.run()
.await
.expect("run ok");
assert_eq!(
found,
Some(Finding {
title: "SQLi".into(),
severity: 9
})
);
assert_eq!(ctx.budget().spent(), 10);
}
#[tokio::test]
async fn schema_value_returns_raw_json() {
let payload = serde_json::json!({ "anything": [1, 2, 3] });
let ctx = respond_ctx(payload.clone());
let out: Option<serde_json::Value> = agent(&ctx, "task")
.schema_value(serde_json::json!({ "type": "object" }))
.run()
.await
.expect("run ok");
assert_eq!(out, Some(payload));
}
#[tokio::test]
async fn finished_without_respond_is_err_not_none() {
let mock = Arc::new(MockProvider::new(vec![MockProvider::text_response(
"just prose",
5,
5,
)]));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::from_arc(mock)))
.build()
.expect("build ctx");
let result = agent(&ctx, "audit").schema::<Finding>().run().await;
assert!(
result.is_err(),
"missing structured output must be Err, got {result:?}"
);
}
#[derive(serde::Deserialize, Debug)]
#[allow(dead_code)]
struct Choice {
pick: Pick,
}
#[derive(serde::Deserialize, Debug)]
#[serde(rename_all = "lowercase")]
#[allow(dead_code)]
enum Pick {
Yes,
No,
}
impl StructuredSchema for Choice {
fn json_schema() -> serde_json::Value {
serde_json::json!({
"type": "object",
"required": ["pick"],
"properties": { "pick": { "type": "string" } }
})
}
}
#[tokio::test]
async fn serde_catches_what_jsonschema_misses() {
let ctx = respond_ctx(serde_json::json!({ "pick": "maybe" }));
let result = agent(&ctx, "decide").schema::<Choice>().run().await;
match result {
Err(Error::Json(_)) => {}
other => panic!("expected Error::Json from serde, got {other:?}"),
}
}
#[tokio::test]
async fn cancelled_typed_run_is_none() {
let ctx = respond_ctx(serde_json::json!({ "title": "x", "severity": 1 }));
ctx.cancellation_token().cancel();
let found: Option<Finding> = agent(&ctx, "audit")
.schema::<Finding>()
.run()
.await
.expect("run ok");
assert!(found.is_none(), "cancelled typed run must be Ok(None)");
}
#[tokio::test]
async fn raw_json_marker_is_distinct_from_no_schema() {
let ctx = respond_ctx(serde_json::json!({ "k": "v" }));
let call = agent(&ctx, "t").schema_value(serde_json::json!({ "type": "object" }));
let _assert_type: fn(AgentCall<RawJson>) = |_c| {};
_assert_type(call);
}
use super::super::journal::ResumeMode;
use crate::ExecutionContext;
use crate::llm::types::ToolDefinition;
use crate::tool::{Tool, ToolOutput};
struct RecordingTool {
name: String,
invoked: Arc<AtomicUsize>,
}
impl Tool for RecordingTool {
fn definition(&self) -> ToolDefinition {
ToolDefinition {
name: self.name.clone(),
description: "records that it ran".into(),
input_schema: serde_json::json!({ "type": "object" }),
}
}
fn execute(
&self,
_ctx: &ExecutionContext,
_input: serde_json::Value,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = Result<ToolOutput, Error>> + Send + '_>,
> {
let invoked = Arc::clone(&self.invoked);
Box::pin(async move {
invoked.fetch_add(1, Ordering::SeqCst);
Ok(ToolOutput::success("recorded"))
})
}
}
struct CallsToolProvider {
tool_name: String,
turn: AtomicUsize,
}
impl LlmProvider for CallsToolProvider {
async fn complete(&self, _request: CompletionRequest) -> Result<CompletionResponse, Error> {
let turn = self.turn.fetch_add(1, Ordering::SeqCst);
let content = if turn == 0 {
vec![ContentBlock::ToolUse {
id: "call-1".into(),
name: self.tool_name.clone(),
input: serde_json::json!({}),
}]
} else {
vec![ContentBlock::Text {
text: "done".into(),
}]
};
Ok(CompletionResponse {
content,
stop_reason: if turn == 0 {
StopReason::ToolUse
} else {
StopReason::EndTurn
},
usage: TokenUsage {
input_tokens: 1,
output_tokens: 1,
..Default::default()
},
model: None,
reasoning: None,
})
}
fn model_name(&self) -> Option<&str> {
Some("calls-tool-mock")
}
}
fn tool_ctx(tool_name: &str) -> WorkflowCtx {
WorkflowCtx::builder(Arc::new(BoxedProvider::new(CallsToolProvider {
tool_name: tool_name.to_string(),
turn: AtomicUsize::new(0),
})))
.build()
.expect("build ctx")
}
#[tokio::test]
async fn per_call_tools_are_invoked_by_the_agent() {
let invoked = Arc::new(AtomicUsize::new(0));
let tool = Arc::new(RecordingTool {
name: "rec".into(),
invoked: Arc::clone(&invoked),
}) as Arc<dyn Tool>;
let ctx = tool_ctx("rec");
agent(&ctx, "use the tool")
.tools(vec![tool])
.run()
.await
.expect("run ok");
assert_eq!(
invoked.load(Ordering::SeqCst),
1,
"the per-call wired tool must be executed by the agent"
);
}
#[tokio::test]
async fn ctx_base_tools_are_used_when_no_per_call_tools() {
let invoked = Arc::new(AtomicUsize::new(0));
let tool = Arc::new(RecordingTool {
name: "rec".into(),
invoked: Arc::clone(&invoked),
}) as Arc<dyn Tool>;
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::new(CallsToolProvider {
tool_name: "rec".into(),
turn: AtomicUsize::new(0),
})))
.base_tools(vec![tool])
.build()
.expect("build ctx");
agent(&ctx, "use the tool").run().await.expect("run ok");
assert_eq!(
invoked.load(Ordering::SeqCst),
1,
"ctx base_tools must be used when the call sets no tools"
);
}
#[tokio::test]
async fn per_call_tools_override_base_tools() {
let base_invoked = Arc::new(AtomicUsize::new(0));
let per_invoked = Arc::new(AtomicUsize::new(0));
let base = Arc::new(RecordingTool {
name: "base".into(),
invoked: Arc::clone(&base_invoked),
}) as Arc<dyn Tool>;
let per = Arc::new(RecordingTool {
name: "rec".into(),
invoked: Arc::clone(&per_invoked),
}) as Arc<dyn Tool>;
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::new(CallsToolProvider {
tool_name: "rec".into(),
turn: AtomicUsize::new(0),
})))
.base_tools(vec![base])
.build()
.expect("build ctx");
agent(&ctx, "use rec")
.tools(vec![per])
.run()
.await
.expect("run ok");
assert_eq!(per_invoked.load(Ordering::SeqCst), 1, "per-call tool runs");
assert_eq!(
base_invoked.load(Ordering::SeqCst),
0,
"base tool must be shadowed by the per-call override"
);
}
struct CountingProvider {
calls: Arc<AtomicUsize>,
}
impl LlmProvider for CountingProvider {
async fn complete(&self, _request: CompletionRequest) -> Result<CompletionResponse, Error> {
let n = self.calls.fetch_add(1, Ordering::SeqCst);
Ok(CompletionResponse {
content: vec![ContentBlock::Text {
text: format!("live-{n}"),
}],
stop_reason: StopReason::EndTurn,
reasoning: None,
usage: TokenUsage {
input_tokens: 10,
output_tokens: 5,
..Default::default()
},
model: None,
})
}
fn model_name(&self) -> Option<&str> {
Some("counting-mock")
}
}
fn counting_ctx(
calls: &Arc<AtomicUsize>,
path: &std::path::Path,
mode: ResumeMode,
) -> WorkflowCtx {
WorkflowCtx::builder(Arc::new(BoxedProvider::new(CountingProvider {
calls: Arc::clone(calls),
})))
.journal(path, mode)
.expect("open journal")
.build()
.expect("build ctx")
}
#[tokio::test]
async fn resume_replays_without_calling_provider() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("run.jsonl");
let calls1 = Arc::new(AtomicUsize::new(0));
let ctx1 = counting_ctx(&calls1, &path, ResumeMode::Fresh);
let first = agent(&ctx1, "do the task").run().await.expect("run ok");
assert_eq!(first.as_deref(), Some("live-0"));
assert_eq!(calls1.load(Ordering::SeqCst), 1);
let calls2 = Arc::new(AtomicUsize::new(0));
let ctx2 = counting_ctx(&calls2, &path, ResumeMode::Resume);
let replayed = agent(&ctx2, "do the task").run().await.expect("run ok");
assert_eq!(
replayed.as_deref(),
Some("live-0"),
"must replay cached text"
);
assert_eq!(
calls2.load(Ordering::SeqCst),
0,
"replay must NOT call the provider"
);
}
#[tokio::test]
async fn replayed_prefix_costs_zero_budget() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("run.jsonl");
let calls1 = Arc::new(AtomicUsize::new(0));
let ctx1 = counting_ctx(&calls1, &path, ResumeMode::Fresh);
agent(&ctx1, "task").run().await.expect("run ok");
let calls2 = Arc::new(AtomicUsize::new(0));
let ctx2 = WorkflowCtx::builder(Arc::new(BoxedProvider::new(CountingProvider {
calls: Arc::clone(&calls2),
})))
.journal(&path, ResumeMode::Resume)
.expect("open journal")
.budget(1) .build()
.expect("build ctx");
let replayed = agent(&ctx2, "task").run().await.expect("run ok");
assert_eq!(replayed.as_deref(), Some("live-0"));
assert_eq!(
ctx2.budget().spent(),
0,
"replayed call must spend 0 budget"
);
assert!(!ctx2.is_cancelled());
}
#[tokio::test]
async fn changed_prompt_reruns_live() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("run.jsonl");
let calls1 = Arc::new(AtomicUsize::new(0));
let ctx1 = counting_ctx(&calls1, &path, ResumeMode::Fresh);
agent(&ctx1, "original").run().await.expect("run ok");
let calls2 = Arc::new(AtomicUsize::new(0));
let ctx2 = counting_ctx(&calls2, &path, ResumeMode::Resume);
let out = agent(&ctx2, "changed").run().await.expect("run ok");
assert_eq!(out.as_deref(), Some("live-0"));
assert_eq!(
calls2.load(Ordering::SeqCst),
1,
"a changed prompt must re-run live"
);
}
#[tokio::test]
async fn reordered_schema_keys_still_replay() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("run.jsonl");
let calls1 = Arc::new(AtomicUsize::new(0));
let ctx1 = WorkflowCtx::builder(Arc::new(BoxedProvider::new(RespondProvider {
payload: serde_json::json!({ "title": "x", "severity": 1 }),
})))
.journal(&path, ResumeMode::Fresh)
.expect("open journal")
.build()
.expect("build ctx");
let _ = calls1; let f1: Option<Finding> = agent(&ctx1, "audit")
.schema::<Finding>()
.run()
.await
.unwrap();
assert!(f1.is_some());
let calls2 = Arc::new(AtomicUsize::new(0));
let ctx2 = counting_ctx(&calls2, &path, ResumeMode::Resume);
let f2: Option<Finding> = agent(&ctx2, "audit")
.schema::<Finding>()
.run()
.await
.unwrap();
assert_eq!(
f2,
Some(Finding {
title: "x".into(),
severity: 1
})
);
assert_eq!(
calls2.load(Ordering::SeqCst),
0,
"identical schema+prompt must replay, not re-run"
);
}
#[tokio::test]
async fn no_journal_path_is_unaffected() {
let calls = Arc::new(AtomicUsize::new(0));
let ctx = WorkflowCtx::builder(Arc::new(BoxedProvider::new(CountingProvider {
calls: Arc::clone(&calls),
})))
.build()
.expect("build ctx");
agent(&ctx, "a").run().await.expect("run ok");
agent(&ctx, "a").run().await.expect("run ok");
assert_eq!(calls.load(Ordering::SeqCst), 2, "no journal -> always live");
}
}