diff --git a/docs/modules/graph/checkpointing.md b/docs/modules/graph/checkpointing.md index 51b3b1a..07e8e26 100644 --- a/docs/modules/graph/checkpointing.md +++ b/docs/modules/graph/checkpointing.md @@ -177,6 +177,27 @@ The checkpoint core lives in `src/graph/checkpoint/`: `Async` degrades to `Sync`. `Exit` persists only the terminal checkpoint and any interrupt boundary (interrupts must persist so the run can resume). Set it with `CompiledGraph::with_durability(mode)`. + +Two guarantees make `Async` durability safe to reason about. **Background writes +are ordered**: each spawned write awaits the previously spawned one before its +own `put`, because every bundled backend defines a thread's "latest" checkpoint +by insertion order — an unserialized writer could let boundary N land after +boundary N+1 and make a concurrent reader (or a `retry` racing a straggler +append) see a stale record as latest. A write whose predecessor failed is +skipped entirely rather than appending a record whose `parent_checkpoint_id` +points at something that never persisted; the predecessor's error is what the +run reports. **Every exit drains**: not only the terminal, interrupt, and +failure boundaries but also the mid-run aborts — recursion limit, run deadline, +node-visit limit, reducer and routing errors — settle in-flight writes before +returning `Err`. Dropping the tracker would detach those tasks, discarding their +outcome and racing a caller that immediately calls `retry(thread_id)`. + +`prune` is namespace-aware: the recency window applies per namespace, so the +lineages of embedded subgraphs (which share the parent's thread id under their +own namespace, and are never referenced by a parent record's +`parent_checkpoint_id`) are retained rather than deleted along with old parent +records. + - `CheckpointConfig { thread_id, checkpoint_id, namespace }` — checkpoint coordinates. `CheckpointConfig::latest(thread_id)` addresses the newest checkpoint at the root namespace. diff --git a/src/graph/checkpoint/mod.rs b/src/graph/checkpoint/mod.rs index 12eddb3..822beb1 100644 --- a/src/graph/checkpoint/mod.rs +++ b/src/graph/checkpoint/mod.rs @@ -266,7 +266,13 @@ where /// /// Strategy (lineage- and delta-safe): /// - /// 1. Protect the most recent `keep_last` checkpoints (listing order). + /// 1. Protect the most recent `keep_last` checkpoints (listing order) *of + /// every namespace present in the thread*. An embedded subgraph writes + /// its checkpoints under the parent's thread id but its own namespace, + /// and its lineage is disjoint from the parent's (no parent-namespace + /// record ever references a child-namespace id), so a thread-wide + /// recency window would delete the child lineage outright and leave the + /// thread unresumable. /// 2. Walk the `parent_checkpoint_id` chain of every protected checkpoint /// and protect every ancestor reached. This is what honors the /// delta-channel warning: a kept checkpoint that only stores a delta (or @@ -283,18 +289,26 @@ where if metas.is_empty() { return Ok(0); } - let keep_last = keep_last.max(1).min(metas.len()); + let keep_last = keep_last.max(1); // Index by id so ancestor walks are O(depth). let mut parent_of: HashMap<&str, Option<&str>> = HashMap::new(); for m in &metas { parent_of.insert(m.checkpoint_id.as_str(), m.parent_checkpoint_id.as_deref()); } + // Group by namespace so each lineage (the root run and every embedded + // subgraph run sharing this thread) gets its own recency window. + let mut by_namespace: HashMap<&Vec, Vec<&CheckpointMetadata>> = HashMap::new(); + for m in &metas { + by_namespace.entry(&m.namespace).or_default().push(m); + } let mut protected: HashSet = HashSet::new(); - // Step 1: the recency window. - for m in metas.iter().rev().take(keep_last) { - protected.insert(m.checkpoint_id.clone()); + // Step 1: the recency window, per namespace. + for group in by_namespace.values() { + for m in group.iter().rev().take(keep_last) { + protected.insert(m.checkpoint_id.clone()); + } } // Step 2: expand to every ancestor of a protected checkpoint. let window: Vec = protected.iter().cloned().collect(); diff --git a/src/graph/checkpoint/test.rs b/src/graph/checkpoint/test.rs index 0dc0969..945c006 100644 --- a/src/graph/checkpoint/test.rs +++ b/src/graph/checkpoint/test.rs @@ -345,6 +345,36 @@ async fn prune_zero_keeps_latest_and_its_chain() { assert_eq!(cp.count("t"), 2); } +#[tokio::test] +async fn prune_keeps_a_window_per_namespace() { + let cp = InMemoryCheckpointer::::new(); + // An embedded subgraph writes under the parent's thread but its own + // namespace, interleaved with the parent's records. The two lineages are + // disjoint — no parent record ever references a child id — so a + // thread-wide recency window would delete the child lineage outright and + // leave the thread unresumable. + let sub = |id: &str, parent: Option<&str>, step: usize| { + let mut c = checkpoint("t", id, parent, step); + c.namespace = vec!["sub".to_string()]; + c + }; + cp.put(checkpoint("t", "p1", None, 1)).await.unwrap(); + cp.put(sub("s1", None, 1)).await.unwrap(); + cp.put(sub("s2", Some("s1"), 2)).await.unwrap(); + cp.put(checkpoint("t", "p2", Some("p1"), 2)).await.unwrap(); + + // Keep the last 1 of each namespace plus its ancestors: {p2, p1} and + // {s2, s1} — nothing is deleted, and the subgraph stays resolvable. + let removed = cp.prune("t", 1).await.unwrap(); + assert_eq!(removed, 0); + let child = cp + .get_scoped("t", None, &["sub".to_string()]) + .await + .unwrap() + .expect("the subgraph namespace must stay resumable after prune"); + assert_eq!(child.checkpoint_id, "s2"); +} + // ---- File-backed checkpointer --------------------------------------------- mod file_backend { diff --git a/src/graph/compiled/executor.rs b/src/graph/compiled/executor.rs index 7b3330e..d12bbbb 100644 --- a/src/graph/compiled/executor.rs +++ b/src/graph/compiled/executor.rs @@ -182,10 +182,24 @@ where )); } + // The resume value belongs to the node(s) that actually interrupted. The + // pending set is deliberately wider than that at an interrupt boundary + // (it also carries the successors of branches that completed before the + // interrupt), so fanning the value across it would hand `ctx.resume` to + // nodes that have never run. A boundary that recorded no interrupt (a + // failure boundary, resumed via `retry` with no value) keeps the old + // fan-across-pending behaviour. let mut resume_map = HashMap::new(); if let Some(value) = command.resume { - for activation in &active { - resume_map.insert(activation.node.clone(), value.clone()); + let interrupted = interrupted_nodes(&checkpoint, &active); + if interrupted.is_empty() { + for activation in &active { + resume_map.insert(activation.node.clone(), value.clone()); + } + } else { + for node in interrupted { + resume_map.insert(node, value.clone()); + } } } @@ -414,9 +428,16 @@ where .min(self.recursion_policy.max_total_steps); if steps >= step_limit { let err = TinyAgentsError::RecursionLimit(step_limit); - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + return self + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; - return Err(err); } // Whole-run wall-clock deadline: stop *between* super-steps once the // elapsed run time reaches it, leaving the last committed boundary @@ -430,17 +451,31 @@ where "graph run exceeded its {deadline:?} deadline after {steps} super-step(s) \ ({elapsed:?} elapsed)" )); - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + return self + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; - return Err(err); } } // Node-loop recursion: enforce `max_visits_per_node` per activation. for activation in &active { if let Err(err) = recursion.record_node_visit(&mut node_visits, &activation.node) { - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + return self + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; - return Err(err); } } steps += 1; @@ -486,9 +521,16 @@ where } = match run_result { Ok(step_run) => step_run, Err(err) => { - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + return self + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; - return Err(err); } }; @@ -500,7 +542,14 @@ where Ok(state) => state, Err(err) => { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; } }; @@ -540,7 +589,14 @@ where Ok(successors) => successors, Err(route_err) => { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, route_err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + route_err, + ) .await; } }; @@ -593,9 +649,16 @@ where // too. Then return control to the caller. if let Some((index, emitted)) = interrupt { if let Err(err) = self.require_interrupt_durability(&thread_id) { - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) + return self + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; - return Err(err); } let successors = match self.route_completed( &active[..index], @@ -606,7 +669,14 @@ where Ok(successors) => successors, Err(route_err) => { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, route_err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + route_err, + ) .await; } }; @@ -620,7 +690,14 @@ where // (a broken lineage cannot be safely resumed from). if let Err(err) = async_writes.drain().await { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; } let checkpoint_id = match self @@ -631,6 +708,7 @@ where &pending, &activation_nodes(&active[..index]), vec![emitted.clone()], + std::slice::from_ref(&active[index].node), &barrier_arrivals, parent_checkpoint.clone(), steps, @@ -643,7 +721,14 @@ where Ok(id) => id, Err(persist_err) => { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, persist_err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + persist_err, + ) .await; } }; @@ -680,7 +765,14 @@ where Ok(next) => next, Err(route_err) => { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, route_err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + route_err, + ) .await; } }; @@ -700,7 +792,14 @@ where // continuing with a hole in its lineage. if let Some(err) = async_writes.take_failure().await { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; } let terminal = next.is_empty(); @@ -728,7 +827,14 @@ where // synchronously in every mode. if terminal && let Err(err) = async_writes.drain().await { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + err, + ) .await; } self.persist_checkpoint( @@ -738,6 +844,7 @@ where &next, &completed_nodes, Vec::new(), + &[], &barrier_arrivals, parent_checkpoint.clone(), steps, @@ -751,7 +858,14 @@ where Ok(id) => id, Err(persist_err) => { return self - .fail_and_return(&run_id, &thread_id, started_at, steps, persist_err) + .fail_and_return( + &run_id, + &thread_id, + started_at, + steps, + &mut async_writes, + persist_err, + ) .await; } } @@ -829,14 +943,23 @@ where /// a reducer merge, a routing resolution, or a checkpoint persist — still /// transitions the run to `Failed` (rather than leaving observers to see it /// stuck in `Running` forever) before the error unwinds out of the run. + /// + /// Any in-flight `Async` background write is drained first: dropping the + /// tracker would detach those tasks, discarding their outcome (contrary to + /// [`AsyncCheckpointWrites`]' contract) and racing a caller that + /// immediately `retry`s the thread. A background write error must not + /// replace the error that aborted the run, so it is dropped here — exactly + /// as at the failure boundary. async fn fail_and_return( &self, run_id: &RunId, thread_id: &Option, started_at: SystemTime, steps: usize, + writes: &mut AsyncCheckpointWrites, err: TinyAgentsError, ) -> Result { + let _ = writes.drain().await; self.fail_run(run_id, thread_id, started_at, steps, &err, None) .await; Err(err) @@ -1337,6 +1460,7 @@ where pending: &[Activation], completed_tasks: &[NodeId], interrupts: Vec, + interrupted: &[NodeId], barrier_arrivals: &HashMap>, parent: Option, step: usize, @@ -1354,6 +1478,7 @@ where pending, completed_tasks, interrupts, + interrupted, barrier_arrivals, parent, step, @@ -1413,6 +1538,7 @@ where pending, completed_tasks, Vec::new(), + &[], barrier_arrivals, parent, step, @@ -1426,7 +1552,7 @@ where Ok(handle) => { let checkpointer = Arc::clone(checkpointer); let sink = self.event_sink.clone(); - writes.push(handle.spawn(async move { + writes.spawn_ordered(&handle, async move { let id = checkpointer.put(checkpoint).await?; if let Some(sink) = sink { sink.emit(GraphEvent::CheckpointSaved { @@ -1434,7 +1560,7 @@ where }); } Ok(id) - })); + }); Ok(Some(id)) } Err(_) => { @@ -1458,6 +1584,7 @@ where pending: &[Activation], completed_tasks: &[NodeId], interrupts: Vec, + interrupted: &[NodeId], barrier_arrivals: &HashMap>, parent: Option, step: usize, @@ -1465,6 +1592,23 @@ where recursion: &serde_json::Value, child_runs: &serde_json::Value, ) -> Checkpoint { + let mut metadata = serde_json::json!({ + "source": source, + "step": step, + "recursion": recursion, + "child_runs": child_runs, + }); + // Which node of *this* graph paused, as opposed to the (possibly + // re-emitted, child-owned) `Interrupt::node`. Resume keys the resume + // value on it; omitted entirely when nothing interrupted. + if !interrupted.is_empty() { + metadata["interrupted_nodes"] = serde_json::json!( + interrupted + .iter() + .map(|n| n.to_string()) + .collect::>() + ); + } Checkpoint { thread_id: thread.to_string(), checkpoint_id: next_checkpoint_id(), @@ -1478,12 +1622,7 @@ where pending_activations: Some(pending.iter().map(PendingActivation::from).collect()), barrier_arrivals: barriers_to_persisted(barrier_arrivals), interrupts, - metadata: serde_json::json!({ - "source": source, - "step": step, - "recursion": recursion, - "child_runs": child_runs, - }), + metadata, } } diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index 28c5f04..c7c5b31 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -235,6 +235,44 @@ fn activation_nodes(active: &[Activation]) -> Vec { active.iter().map(|a| a.node.clone()).collect() } +/// The nodes of *this* graph that paused at an interrupt boundary — the ones a +/// resume value belongs to. +/// +/// Read from the `interrupted_nodes` metadata the interrupt boundary stamps, +/// which is the local activation rather than +/// [`Interrupt::node`](crate::graph::command::Interrupt::node): a subgraph node +/// re-emits its *child's* interrupt, so the recorded interrupt can name a node +/// that does not exist in this graph. Falls back to the recorded interrupts for +/// checkpoints written before that metadata existed — intersected with +/// `pending`, because a legacy checkpoint of a re-emitted child interrupt names +/// a node this graph never schedules. An empty result tells the caller to fan +/// the resume value across the pending set, which is what those checkpoints got +/// before the metadata existed. +fn interrupted_nodes(checkpoint: &Checkpoint, pending: &[Activation]) -> Vec { + let stamped = checkpoint + .metadata + .get("interrupted_nodes") + .and_then(serde_json::Value::as_array) + .map(|nodes| { + nodes + .iter() + .filter_map(serde_json::Value::as_str) + .map(NodeId::from) + .collect::>() + }) + .unwrap_or_default(); + if stamped.is_empty() { + let scheduled: HashSet<&NodeId> = pending.iter().map(|a| &a.node).collect(); + return checkpoint + .interrupts + .iter() + .map(|i| i.node.clone()) + .filter(|node| scheduled.contains(node)) + .collect(); + } + stamped +} + impl CompiledGraph { /// Internal constructor used by the builder. #[allow(clippy::too_many_arguments)] diff --git a/src/graph/compiled/state_api.rs b/src/graph/compiled/state_api.rs index de40080..52d4ab5 100644 --- a/src/graph/compiled/state_api.rs +++ b/src/graph/compiled/state_api.rs @@ -71,7 +71,11 @@ where /// attributed node). A command node cannot be used as `as_node` (it routes /// dynamically and has no static successors); doing so returns /// [`TinyAgentsError::Graph`] rather than silently producing a non-resumable - /// checkpoint. With `as_node == None` the latest pending node set is + /// checkpoint. A successor reached by a waiting edge is barrier-gated + /// exactly as it would be during a run: it is scheduled only once every + /// required predecessor has arrived (the write records the attributed + /// node's arrival), so a manual write can never fire a join ahead of a + /// still-pending branch. With `as_node == None` the latest pending node set is /// preserved. Requires a configured checkpointer and an existing checkpoint /// for the thread. pub async fn update_state( @@ -111,26 +115,73 @@ where let parent_id = base.checkpoint_id.clone(); let new_state = self.reducer.apply(base.state, update)?; + // Manual writes preserve any accumulated barrier arrivals, and an + // attributed write records its own arrival into them. + let mut arrivals = barriers_from_persisted(&base.barrier_arrivals); // Pending nodes: the attributed node's successors, or the inherited set. + let mut withheld_by_barrier = false; + // Set only when the base checkpoint's pending set was reinstated below, + // which is the one case where the pending *activations* must be + // reinstated with it. Keying that off `withheld_by_barrier` would let a + // routing that withholds one target while scheduling another persist + // `next_nodes` and `pending_activations` that disagree — and resume + // prefers the activations, silently dropping the scheduled successor. + let mut used_base_fallback = false; let next_nodes: Vec = match &as_node { - Some(node) => self - .route(node, None, &new_state)? - .into_iter() - .map(|t| t.node().clone()) - .filter(|n| n.as_str() != END) - .collect(), + Some(node) => { + let mut next = Vec::new(); + for target in self.route(node, None, &new_state)? { + let tnode = target.node().clone(); + if tnode.as_str() == END { + continue; + } + // Apply the same barrier gate the executor applies in + // `route_completed`: a waiting node stays unscheduled until + // every required predecessor has arrived. Without this an + // attributed write would fire a join ahead of a predecessor + // that is still pending — the data loss the waiting edge + // exists to prevent. + if let Some(required) = self.waiting.get(&tnode) { + let arrived = arrivals.entry(tnode.clone()).or_default(); + arrived.insert(node.clone()); + if !required.is_subset(arrived) { + withheld_by_barrier = true; + continue; + } + arrivals.remove(&tnode); + } + next.push(tnode); + } + // An unsatisfied barrier leaves nothing to schedule, which would + // make the checkpoint non-resumable. Keep the base checkpoint's + // still-pending nodes (minus the attributed one) so the barrier's + // remaining predecessors still run and clear the join. + if next.is_empty() && withheld_by_barrier { + used_base_fallback = true; + next.extend(base.next_nodes.iter().filter(|n| *n != node).cloned()); + } + next + } None => base.next_nodes.clone(), }; let completed_tasks: Vec = as_node.iter().cloned().collect(); // With `as_node`, pending becomes that node's (plain) successors, so no // send args carry over; without it, inherit the base checkpoint's - // pending activations verbatim so any pending `Send` args survive. - let pending_activations = match &as_node { - Some(_) => None, - None => base.pending_activations.clone(), + // pending activations verbatim so any pending `Send` args survive. The + // barrier-withheld fallback above re-schedules base pending nodes, so it + // keeps their activations (and `Send` args) too. + let pending_activations = match (&as_node, used_base_fallback) { + (Some(node), true) => base.pending_activations.as_ref().map(|pending| { + pending + .iter() + .filter(|activation| activation.node != *node) + .cloned() + .collect() + }), + (Some(_), false) => None, + (None, _) => base.pending_activations.clone(), }; - // Manual writes preserve any accumulated barrier arrivals. - let barrier_arrivals = base.barrier_arrivals.clone(); + let barrier_arrivals = barriers_to_persisted(&arrivals); let checkpoint_id = next_checkpoint_id(); let config = self.config_for(thread_id, Some(&checkpoint_id)); diff --git a/src/graph/compiled/test.rs b/src/graph/compiled/test.rs index 8eb370c..a8e308e 100644 --- a/src/graph/compiled/test.rs +++ b/src/graph/compiled/test.rs @@ -2435,3 +2435,476 @@ async fn async_durability_surfaces_background_write_failure_in_run_result() { "unexpected error: {err}" ); } + +// --------------------------------------------------------------------------- +// Regression: resume value targeting, barrier-gated manual writes, async drains +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn resume_value_reaches_only_the_interrupted_node() { + // Parallel [a, b]: a routes to successor `x` and completes; b interrupts. + // The interrupt boundary schedules both `x` and `b`, but only `b` actually + // interrupted — `x` has never run, so it must observe `ctx.resume == None` + // on its first activation. Fanning the resume value across the whole + // pending set used to drive `x` down its "already approved" arm. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let graph = GraphBuilder::::new() + .with_parallel(true) + .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { + s.value += u; + s.log.push(format!("+{u}")); + Ok(s) + })) + .add_node("super", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Command( + Command::default().with_goto(["a", "b"]), + )) + }) + .add_node("a", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Update(1)) + }) + .add_node("b", |_s: Counter, c: NodeContext| async move { + match c.resume { + Some(_) => Ok(NodeResult::Update(100)), + None => Ok(NodeResult::Interrupt(Interrupt::new("b", json!({})))), + } + }) + // 10 on a first (unresumed) activation, 1000 if it wrongly sees a + // resume value it never asked for. + .add_node("x", |_s: Counter, c: NodeContext| async move { + Ok(NodeResult::Update(if c.resume.is_some() { + 1000 + } else { + 10 + })) + }) + .set_entry("super") + .mark_command_routing("super") + .add_edge("a", "x") + .set_finish("b") + .set_finish("x") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let paused = graph + .run_with_thread( + "t-resume-scope", + Counter { + value: 0, + log: vec![], + }, + ) + .await + .unwrap(); + assert!(paused.is_interrupted()); + + let done = graph + .resume("t-resume-scope", Command::resume(json!({"approved": true}))) + .await + .unwrap(); + // 1 (a) + 100 (b, resumed) + 10 (x, first activation, no resume value). + assert_eq!( + done.state.value, 111, + "the resume value must reach only the node that interrupted" + ); +} + +#[tokio::test] +async fn attributed_update_does_not_fire_an_unsatisfied_barrier() { + // Diamond `super -> {b, c} -> merge` joined by waiting edges. `c` interrupts + // on its first activation, so the pause leaves `c` pending with only `b` + // arrived at the barrier. A manual write attributed to `b` must not + // schedule `merge` (that would run the join without `c`'s contribution), + // and the recorded arrival must let `merge` fire once `c` completes. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let interrupted = Arc::new(AtomicBool::new(false)); + let once = interrupted.clone(); + let graph = GraphBuilder::::new() + .with_parallel(true) + .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { + s.value += u; + s.log.push(format!("+{u}")); + Ok(s) + })) + .add_node("super", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Command( + Command::default().with_goto(["b", "c"]), + )) + }) + .add_node("b", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Update(1)) + }) + .add_node("c", move |_s: Counter, _c: NodeContext| { + let once = once.clone(); + async move { + if once.swap(true, AtomicOrdering::SeqCst) { + Ok(NodeResult::Update(2)) + } else { + Ok(NodeResult::Interrupt(Interrupt::new("c", json!({})))) + } + } + }) + .add_node("merge", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Update(100)) + }) + .set_entry("super") + .mark_command_routing("super") + .add_waiting_edge("b", "merge") + .add_waiting_edge("c", "merge") + .set_finish("merge") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let paused = graph + .run_with_thread( + "t-barrier-update", + Counter { + value: 0, + log: vec![], + }, + ) + .await + .unwrap(); + assert!(paused.is_interrupted()); + + // Operator edits state and attributes the write to `b`, the barrier + // predecessor that already ran. + graph + .update_state("t-barrier-update", 10, Some(NodeId::from("b"))) + .await + .unwrap(); + let written = cp.get("t-barrier-update", None).await.unwrap().unwrap(); + assert!( + !written.next_nodes.iter().any(|n| n.as_str() == "merge"), + "an unsatisfied barrier must not be scheduled by an attributed write" + ); + assert!( + written.next_nodes.iter().any(|n| n.as_str() == "c"), + "the still-pending barrier predecessor must stay scheduled" + ); + // Resume prefers `pending_activations` over `next_nodes`, so the two must + // never disagree: a node named by only one of them would be silently + // dropped (or scheduled without its `Send` arg). + if let Some(pending) = &written.pending_activations { + assert_eq!( + pending.iter().map(|a| a.node.clone()).collect::>(), + written.next_nodes, + "pending activations and next nodes must describe the same schedule" + ); + } + + let done = graph.retry("t-barrier-update").await.unwrap(); + assert!( + done.visited.iter().any(|n| n.as_str() == "merge"), + "merge must fire once the remaining predecessor arrives" + ); + // 1 (b) + 10 (manual write) + 2 (c) + 100 (merge). + assert_eq!(done.state.value, 113); +} + +#[tokio::test] +async fn async_durability_drains_background_writes_on_abort() { + // The recursion-limit abort returns `Err` mid-run. Any in-flight background + // checkpoint write must be settled first: dropping the tracker would detach + // the tasks, discarding their outcome and racing a caller that immediately + // retries the thread. + use crate::graph::checkpoint::DurabilityMode; + + let completed_puts = Arc::new(AtomicUsize::new(0)); + let cp = Arc::new(SlowCheckpointer { + inner: Arc::new(InMemoryCheckpointer::new()), + delay: Duration::from_millis(50), + completed_puts: completed_puts.clone(), + }); + let graph = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("a") + .add_edge("a", "a") + .with_recursion_limit(3) + .compile() + .unwrap() + .with_checkpointer(cp) + .with_durability(DurabilityMode::Async); + + let err = graph + .run_with_thread("t-async-abort", 0) + .await + .expect_err("the run must abort at the recursion limit"); + assert!(matches!(err, TinyAgentsError::RecursionLimit(3))); + assert_eq!( + completed_puts.load(AtomicOrdering::SeqCst), + 3, + "every background boundary write must be settled before the abort returns" + ); +} + +/// Delegating checkpointer that makes the *first* `put` slow and every later +/// one instant, so an unserialized background writer would append boundary 2 +/// before boundary 1. +struct FirstPutSlowCheckpointer { + inner: Arc>, + started: AtomicUsize, +} + +#[async_trait::async_trait] +impl Checkpointer for FirstPutSlowCheckpointer { + async fn put( + &self, + checkpoint: crate::graph::checkpoint::Checkpoint, + ) -> crate::error::Result { + if self.started.fetch_add(1, AtomicOrdering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(100)).await; + } + self.inner.put(checkpoint).await + } + + async fn get( + &self, + thread_id: &str, + checkpoint_id: Option<&str>, + ) -> crate::error::Result>> { + self.inner.get(thread_id, checkpoint_id).await + } + + async fn list( + &self, + thread_id: &str, + ) -> crate::error::Result> { + self.inner.list(thread_id).await + } + + async fn list_threads(&self) -> crate::error::Result> { + self.inner.list_threads().await + } + + async fn delete_thread(&self, thread_id: &str) -> crate::error::Result<()> { + self.inner.delete_thread(thread_id).await + } + + async fn delete_checkpoints( + &self, + thread_id: &str, + ids: &[String], + ) -> crate::error::Result { + self.inner.delete_checkpoints(thread_id, ids).await + } +} + +#[tokio::test] +async fn async_durability_appends_boundaries_in_order() { + // Every bundled backend defines a thread's "latest" checkpoint by insertion + // order, so background writes must land in boundary order even when an + // earlier `put` is slower than a later one. + use crate::graph::checkpoint::DurabilityMode; + + let cp = Arc::new(FirstPutSlowCheckpointer { + inner: Arc::new(InMemoryCheckpointer::new()), + started: AtomicUsize::new(0), + }); + let graph = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("c", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("a") + .add_edge("a", "b") + .add_edge("b", "c") + .set_finish("c") + .compile() + .unwrap() + .with_checkpointer(cp.clone()) + .with_durability(DurabilityMode::Async); + + let run = graph.run_with_thread("t-async-order", 0).await.unwrap(); + assert_eq!(run.state, 3); + // Insertion order is listing order: each record's parent must be the one + // appended just before it. + let list = cp.list("t-async-order").await.unwrap(); + assert_eq!(list.len(), 3); + for pair in list.windows(2) { + assert_eq!( + pair[1].parent_checkpoint_id.as_deref(), + Some(pair[0].checkpoint_id.as_str()), + "boundary writes landed out of order" + ); + } +} + +#[tokio::test] +async fn legacy_interrupt_checkpoint_without_stamped_nodes_still_resumes() { + // Checkpoints written before the `interrupted_nodes` metadata existed carry + // only `Interrupt::node`, which for a re-emitted child interrupt (what a + // subgraph node does) names a node this graph never schedules. The resume + // value must still reach the paused node — falling back to the pending set — + // rather than being addressed to a node that does not exist here, which + // would re-run the paused node unresumed and interrupt forever. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let graph = GraphBuilder::::new() + .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { + s.value += u; + s.log.push(format!("+{u}")); + Ok(s) + })) + .add_node("gate", |_s: Counter, c: NodeContext| async move { + match c.resume { + // The interrupt names a foreign node, as a subgraph node's + // re-emitted child interrupt does. + None => Ok(NodeResult::Interrupt(Interrupt::new( + "child-gate", + json!({}), + ))), + Some(_) => Ok(NodeResult::Update(5)), + } + }) + .set_entry("gate") + .set_finish("gate") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let paused = graph + .run_with_thread( + "t-legacy-resume", + Counter { + value: 0, + log: vec![], + }, + ) + .await + .unwrap(); + assert!(paused.is_interrupted()); + + // Age the boundary checkpoint into its pre-upgrade shape. + let mut legacy = cp.get("t-legacy-resume", None).await.unwrap().unwrap(); + legacy + .metadata + .as_object_mut() + .unwrap() + .remove("interrupted_nodes"); + cp.put(legacy).await.unwrap(); + + let done = graph + .resume("t-legacy-resume", Command::resume(json!("go"))) + .await + .unwrap(); + assert!( + !done.is_interrupted(), + "a legacy interrupt checkpoint must still deliver the resume value" + ); + assert_eq!(done.state.value, 5); +} + +/// Delegating checkpointer whose first `put` fails *slowly* (so the next +/// boundary's write is already chained behind it) and records every attempt. +struct FirstPutFailsSlowlyCheckpointer { + inner: Arc>, + attempts: Arc, +} + +#[async_trait::async_trait] +impl Checkpointer for FirstPutFailsSlowlyCheckpointer { + async fn put( + &self, + checkpoint: crate::graph::checkpoint::Checkpoint, + ) -> crate::error::Result { + if self.attempts.fetch_add(1, AtomicOrdering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(100)).await; + return Err(crate::error::TinyAgentsError::Checkpoint( + "injected background write failure".to_string(), + )); + } + self.inner.put(checkpoint).await + } + + async fn get( + &self, + thread_id: &str, + checkpoint_id: Option<&str>, + ) -> crate::error::Result>> { + self.inner.get(thread_id, checkpoint_id).await + } + + async fn list( + &self, + thread_id: &str, + ) -> crate::error::Result> { + self.inner.list(thread_id).await + } + + async fn list_threads(&self) -> crate::error::Result> { + self.inner.list_threads().await + } + + async fn delete_thread(&self, thread_id: &str) -> crate::error::Result<()> { + self.inner.delete_thread(thread_id).await + } + + async fn delete_checkpoints( + &self, + thread_id: &str, + ids: &[String], + ) -> crate::error::Result { + self.inner.delete_checkpoints(thread_id, ids).await + } +} + +#[tokio::test] +async fn async_durability_skips_a_write_whose_predecessor_failed() { + // Writing boundary N+1 after boundary N failed would durably append a + // record whose `parent_checkpoint_id` points at something that never + // persisted. The chained write must skip its own `put` and report the + // failure that broke the lineage. + use crate::graph::checkpoint::DurabilityMode; + + let attempts = Arc::new(AtomicUsize::new(0)); + let cp = Arc::new(FirstPutFailsSlowlyCheckpointer { + inner: Arc::new(InMemoryCheckpointer::new()), + attempts: attempts.clone(), + }); + let graph = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("c", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("a") + .add_edge("a", "b") + .add_edge("b", "c") + .set_finish("c") + .compile() + .unwrap() + .with_checkpointer(cp.clone()) + .with_durability(DurabilityMode::Async); + + let err = graph + .run_with_thread("t-async-orphan", 0) + .await + .expect_err("a lost background checkpoint must fail the run"); + assert!( + err.to_string() + .contains("injected background write failure"), + "unexpected error: {err}" + ); + assert_eq!( + attempts.load(AtomicOrdering::SeqCst), + 1, + "the write chained behind a failed one must not be attempted" + ); + assert!( + cp.list("t-async-orphan").await.unwrap().is_empty(), + "no orphaned checkpoint may be appended after a broken lineage" + ); +} diff --git a/src/graph/compiled/types.rs b/src/graph/compiled/types.rs index 2419bdb..994df5d 100644 --- a/src/graph/compiled/types.rs +++ b/src/graph/compiled/types.rs @@ -143,20 +143,48 @@ impl Clone for CompiledGraph { /// in-flight writes at its terminal/interrupt boundary ([`Self::drain`]). A /// run therefore never reports success while one of its checkpoints silently /// failed to persist. +/// +/// # Ordering +/// +/// Writes are *chained*, not merely spawned: each background task awaits the +/// previously spawned one before performing its own `put`. Every bundled +/// backend defines a thread's "latest" checkpoint by insertion order, so +/// letting boundary N+1 land before boundary N would make a concurrent reader +/// (or a `retry` racing a straggler append) observe a stale record as latest. #[derive(Default)] pub(crate) struct AsyncCheckpointWrites { /// In-flight (or finished-but-unharvested) background writes, in the order - /// they were spawned. + /// they were spawned. Because each spawn chains onto its predecessor, this + /// holds at most the tail of the chain. handles: Vec>>, } impl AsyncCheckpointWrites { - /// Tracks one spawned background write. - pub(crate) fn push( - &mut self, - handle: tokio::task::JoinHandle>, - ) { - self.handles.push(handle); + /// Spawns `write` so it runs only after every previously spawned write has + /// settled, keeping appends in boundary order. + /// + /// The spawned task reports the first error along the chain (in spawn + /// order), so a predecessor's failure is never lost by being folded behind + /// a later success. Awaiting a panicked predecessor yields a join error + /// rather than deadlocking the chain. + pub(crate) fn spawn_ordered(&mut self, runtime: &tokio::runtime::Handle, write: F) + where + F: std::future::Future> + Send + 'static, + { + let previous = self.handles.pop(); + self.handles.push(runtime.spawn(async move { + let prior = match previous { + Some(handle) => Self::harvest(handle.await), + None => None, + }; + // A predecessor that never persisted would leave this record's + // `parent_checkpoint_id` dangling, so skip the write entirely and + // report the failure that broke the lineage. + match prior { + Some(err) => Err(err), + None => write.await, + } + })); } /// Harvests writes that have already finished, without blocking on the diff --git a/src/graph/goals/test.rs b/src/graph/goals/test.rs index aa4c17b..839ec8b 100644 --- a/src/graph/goals/test.rs +++ b/src/graph/goals/test.rs @@ -274,6 +274,7 @@ mod tool_tests { depth: 0, max_turn_output_tokens: None, events: EventSink::new(), + cancellation: crate::harness::cancel::CancellationToken::new(), streaming: false, workspace: None, } diff --git a/src/graph/todos/test.rs b/src/graph/todos/test.rs index 03a3586..077ce78 100644 --- a/src/graph/todos/test.rs +++ b/src/graph/todos/test.rs @@ -498,6 +498,7 @@ mod tool_tests { depth: 0, max_turn_output_tokens: None, events: EventSink::new(), + cancellation: crate::harness::cancel::CancellationToken::new(), streaming: false, workspace: None, } diff --git a/src/harness/agent_loop/entry.rs b/src/harness/agent_loop/entry.rs index a281b1e..e927f7d 100644 --- a/src/harness/agent_loop/entry.rs +++ b/src/harness/agent_loop/entry.rs @@ -205,8 +205,13 @@ impl AgentHarness { status.set_last_event(record.id); status.mark_failed(error.to_string()); // Surface the failure to every middleware. Inner errors are - // ignored so the originating error is never masked. - let _ = self.middleware.run_on_error(&mut ctx, &error).await; + // ignored so the originating error is never masked. A failure + // raised *inside* a lifecycle hook was already fanned out by the + // stack before it propagated here, so skip it rather than + // delivering the same failure twice. + if !ctx.take_on_error_dispatched() { + let _ = self.middleware.run_on_error(&mut ctx, &error).await; + } Err(error) } } diff --git a/src/harness/agent_loop/run_loop.rs b/src/harness/agent_loop/run_loop.rs index 1dfe343..9f71ceb 100644 --- a/src/harness/agent_loop/run_loop.rs +++ b/src/harness/agent_loop/run_loop.rs @@ -333,6 +333,17 @@ impl AgentHarness { continue; } + // This turn resolved without scheduling a truncated-empty + // retry, so the recovery state must not leak into later turns: + // a stale `boosted_max_tokens` would override the caller's + // per-turn cap on every subsequent call, and a spent retry + // counter would deny recovery to a later turn that needs it. + reset_truncated_empty_recovery( + &mut truncated_empty_retries_used, + &mut boosted_max_tokens, + &mut truncation_base, + ); + // The model says it is not finished (`ModelResponse::continue_turn`). // Hand the floor back and ask for another reply instead of taking // this response as the turn's answer. Checked after truncated-empty @@ -375,6 +386,15 @@ impl AgentHarness { break; } + // A tool-calling response is a resolved turn too: clear the + // recovery state before the tools run so the next turn starts from + // the caller's configured cap and a full retry budget. + reset_truncated_empty_recovery( + &mut truncated_empty_retries_used, + &mut boosted_max_tokens, + &mut truncation_base, + ); + // Execute requested tools: serial admission -> serial or // concurrent execution -> ordered fold. Multi-call turns run // concurrently when no tool-wrap middleware is registered; see @@ -435,3 +455,18 @@ impl AgentHarness { Some((Arc::clone(cache), cache_key(request))) } } + +/// Clears the per-turn truncated-empty recovery state (see +/// [`crate::harness::runtime::RunPolicy::truncated_empty_retries`]). +/// +/// The state is scoped to a single logical turn: the boosted token cap and the +/// retry counter must not carry over into the turns that follow a recovered one. +fn reset_truncated_empty_recovery( + retries_used: &mut u32, + boosted_max_tokens: &mut Option, + truncation_base: &mut Option, +) { + *retries_used = 0; + *boosted_max_tokens = None; + *truncation_base = None; +} diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index 498debb..9663f56 100644 --- a/src/harness/agent_loop/test.rs +++ b/src/harness/agent_loop/test.rs @@ -775,6 +775,65 @@ async fn truncated_empty_response_retries_then_succeeds() { ); } +#[tokio::test] +async fn truncated_empty_boost_does_not_leak_into_later_turns() { + // Regression test: the boosted token cap and the retry counter are per-turn + // recovery state, but they used to live for the whole run — so every turn + // after a recovered one was dispatched at the boosted cap (overriding the + // caller's `max_turn_output_tokens`) and a later truncation got no retry. + let model = Arc::new(crate::harness::testkit::ScriptedModel::new(vec![ + truncated_empty_response(2048), + tool_call_response("c1", "fake", json!({})), + text_response("done", 4, 3), + ])); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness + .register_model("mock", Arc::clone(&model) as _) + .register_tool(Arc::new(FakeTool::new("fake", "tool output"))); + + let ctx = RunContext::new( + RunConfig::new("truncated-leak").with_max_turn_output_tokens(2048), + (), + ); + let run = harness + .invoke_in_context(&(), ctx, vec![Message::user("hi")]) + .await + .expect("the recovered run finishes"); + + assert_eq!(run.text(), Some("done".to_string())); + let sent: Vec> = model.requests().iter().map(|r| r.max_tokens).collect(); + assert_eq!( + sent, + vec![Some(2048), Some(4096), Some(2048)], + "only the retry of the truncated turn carries the boost; the next turn \ + is back at the configured per-turn cap" + ); +} + +#[tokio::test] +async fn truncated_empty_retry_budget_is_restored_for_a_later_turn() { + // The retry budget is per turn too: a second truncated-empty completion in a + // later turn must still be recoverable with the default single retry. + let model = Arc::new(crate::harness::testkit::ScriptedModel::new(vec![ + truncated_empty_response(2048), + tool_call_response("c1", "fake", json!({})), + truncated_empty_response(2048), + text_response("done", 4, 3), + ])); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness + .register_model("mock", Arc::clone(&model) as _) + .register_tool(Arc::new(FakeTool::new("fake", "tool output"))); + + let run = harness + .invoke_default(&(), vec![Message::user("hi")]) + .await + .expect("both truncated turns recover"); + + assert_eq!(run.text(), Some("done".to_string())); + assert_eq!(run.model_calls, 4); +} + #[tokio::test] async fn truncated_empty_retry_budget_stays_unset_when_request_had_none() { // With no per-turn token cap the budget cannot be doubled, but the retry is @@ -2721,6 +2780,52 @@ async fn request_cache_policy_overrides_run_policy_to_disable_caching() { ); } +/// Middleware whose `before_model` hook always fails, counting how many times +/// the failure is delivered back to it through `on_error`. +struct FailingHookMiddleware { + on_error_calls: Arc>, +} + +#[async_trait] +impl Middleware<(), ()> for FailingHookMiddleware { + fn name(&self) -> &str { + "failing_hook" + } + async fn before_model( + &self, + _ctx: &mut RunContext<()>, + _state: &(), + _request: &mut ModelRequest, + ) -> Result<()> { + Err(TinyAgentsError::Model("hook failed".to_string())) + } + async fn on_error(&self, _ctx: &mut RunContext<()>, _error: &TinyAgentsError) -> Result<()> { + *self.on_error_calls.lock().unwrap() += 1; + Ok(()) + } +} + +#[tokio::test] +async fn hook_failure_delivers_on_error_exactly_once() { + // Regression test: the stack fans `on_error` out itself before propagating a + // failed lifecycle hook, and the driver used to run the same hook again for + // the propagated error — so a guardrail counting failures, alerting, or + // compensating did it twice for one failure. + let calls = Arc::new(Mutex::new(0usize)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::new(MockModel::constant("hi"))); + harness.push_middleware(Arc::new(FailingHookMiddleware { + on_error_calls: calls.clone(), + })); + + let err = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect_err("the failing hook aborts the run"); + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + assert_eq!(*calls.lock().unwrap(), 1, "one failure, one on_error"); +} + /// Middleware that requests an early stop-with-final control outcome after the /// first model response, exercising the harness control channel (gap #13). struct EarlyStopMiddleware; diff --git a/src/harness/context/mod.rs b/src/harness/context/mod.rs index c6b0efe..29a70ba 100644 --- a/src/harness/context/mod.rs +++ b/src/harness/context/mod.rs @@ -44,6 +44,12 @@ use crate::harness::ids::{RunId, ThreadId}; use crate::harness::limits::{LimitTracker, RunLimits}; use crate::harness::store::StoreRegistry; +/// Mints the next process-unique [`RunContext::instance_id`]. +fn next_context_instance_id() -> u64 { + static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); + NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed) +} + // ── RunConfig ───────────────────────────────────────────────────────────────── impl RunConfig { @@ -184,6 +190,7 @@ impl RunContext { // (a durable journal replayed after a restart re-mints the same ids). let events = EventSink::with_stream_id(config.run_id.as_str()); Self { + instance_id: next_context_instance_id(), config, data, stores: StoreRegistry::new(), @@ -193,10 +200,34 @@ impl RunContext { cancellation: CancellationToken::new(), control: std::sync::Arc::new(std::sync::Mutex::new(None)), workspace: None, + on_error_dispatched: false, streaming: false, } } + /// Returns this context's process-unique instance id. + /// + /// Two concurrent runs can carry the same [`RunConfig::run_id`] (it is a + /// caller-supplied label), so keep per-run bookkeeping keyed on this + /// instead when it is shared across runs. + pub fn instance_id(&self) -> u64 { + self.instance_id + } + + /// Records that the middleware stack already dispatched `on_error` for the + /// error currently unwinding this run. + pub(crate) fn mark_on_error_dispatched(&mut self) { + self.on_error_dispatched = true; + } + + /// Takes (and clears) the flag set by [`Self::mark_on_error_dispatched`]. + /// + /// `true` means the stack already delivered `on_error` for this failure, so + /// the driver must not dispatch it again. + pub(crate) fn take_on_error_dispatched(&mut self) -> bool { + std::mem::take(&mut self.on_error_dispatched) + } + /// Attaches an isolated workspace descriptor that is threaded into every /// [`ToolExecutionContext`][crate::harness::tool::ToolExecutionContext] this /// run creates, so tools read their allowed root from context. To prepare diff --git a/src/harness/context/types.rs b/src/harness/context/types.rs index 54118a4..c7e0a4f 100644 --- a/src/harness/context/types.rs +++ b/src/harness/context/types.rs @@ -152,6 +152,12 @@ impl MiddlewareControl { /// Unlike [`RunConfig`], `RunContext` is **not** serializable: it owns live /// counters, listener lists, and user handles. pub struct RunContext { + /// Process-unique identity of *this context instance*, minted on + /// construction. Unlike [`RunConfig::run_id`] — a caller-supplied label two + /// concurrent runs may well share — it distinguishes concurrent runs, so + /// shared per-run bookkeeping (a middleware's in-flight reservation, say) + /// can be keyed on it. Read it with [`RunContext::instance_id`]. + pub(crate) instance_id: u64, /// The declarative configuration this context was built from. pub config: RunConfig, /// Arbitrary user-supplied run data. @@ -193,6 +199,12 @@ pub struct RunContext { /// [`WorkspaceIsolation`][crate::harness::workspace::WorkspaceIsolation] /// provider; `None` means no workspace policy is in effect. pub workspace: Option, + /// Whether the middleware stack already fanned `on_error` out to every + /// middleware for the error currently unwinding this run. The stack sets it + /// when a lifecycle hook fails (it dispatches `on_error` itself before + /// propagating), and the agent-loop driver reads it so the same failure is + /// not delivered to every middleware a second time. + pub(crate) on_error_dispatched: bool, /// Whether this run is being driven through the streaming loop path /// (`ChatModel::stream`), set by the agent-loop driver. Threaded into each /// [`ToolExecutionContext`][crate::harness::tool::ToolExecutionContext] so a diff --git a/src/harness/middleware/README.md b/src/harness/middleware/README.md index 72f0d31..695d889 100644 --- a/src/harness/middleware/README.md +++ b/src/harness/middleware/README.md @@ -54,6 +54,11 @@ to log/redact/react), then the *original* error is returned to the caller. Errors raised from `on_error` itself are ignored — they cannot mask the root cause or replace it with a different error. +One failure delivers exactly one `on_error` per middleware. The stack marks the +`RunContext` when it fans a hook failure out, and a driver that also handles the +propagated error (the agent loop does, for failures raised by models, tools, or +the loop itself) skips its own dispatch for that error. + ## Public surface - `Middleware` — the lifecycle trait described above. diff --git a/src/harness/middleware/library/budget.rs b/src/harness/middleware/library/budget.rs index d7b1dc5..c3e0dec 100644 --- a/src/harness/middleware/library/budget.rs +++ b/src/harness/middleware/library/budget.rs @@ -146,7 +146,7 @@ impl BudgetMiddleware { limits, tracker: BudgetTracker::new(), pricing: std::collections::HashMap::new(), - pending_reservation: std::sync::Mutex::new(0), + pending_reservations: std::sync::Mutex::new(std::collections::HashMap::new()), } } @@ -172,6 +172,29 @@ impl BudgetMiddleware { self.tracker.clone() } + /// Records `estimated` as `run`'s outstanding preflight reservation. + /// + /// A run only ever has one model call in flight, so an existing entry means + /// a prior reservation was abandoned; add to it rather than dropping it, so + /// the shared tracker is never left holding an unreleasable amount. + fn record_reservation(&self, run: u64, estimated: u64) { + let mut guard = self + .pending_reservations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *guard.entry(run).or_insert(0) += estimated; + } + + /// Takes `run`'s outstanding reservation, leaving nothing behind. Returns 0 + /// when the run has none (a second release for one reservation is a no-op). + fn take_reservation(&self, run: u64) -> u64 { + self.pending_reservations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&run) + .unwrap_or(0) + } + fn price(&self, response: &ModelResponse) -> crate::harness::cost::CostTotals { let Some(usage) = response.usage else { return crate::harness::cost::CostTotals::default(); @@ -243,12 +266,9 @@ impl Middleware for BudgetMidd } // Remember this run's own outstanding reservation for reconciliation - // in `after_model` (local to this middleware instance, so concurrent - // runs sharing the tracker never clobber each other's amount). - *self - .pending_reservation - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = estimated; + // in `after_model`, keyed by the run so concurrent runs on this + // middleware never release each other's amount. + self.record_reservation(ctx.instance_id(), estimated); ctx.emit(AgentEvent::BudgetReserved { estimated_input_tokens: estimated, @@ -266,12 +286,7 @@ impl Middleware for BudgetMidd // usage came back, so a call that fails to report usage (or errors // out before this hook) never leaks a permanent reservation that // starves later calls on a shared tracker. - let reserved = std::mem::take( - &mut *self - .pending_reservation - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner), - ); + let reserved = self.take_reservation(ctx.instance_id()); { let mut guard = self.tracker.lock_recovering(); guard.reserved_input_total = guard.reserved_input_total.saturating_sub(reserved); @@ -323,7 +338,7 @@ impl Middleware for BudgetMidd Ok(()) } - async fn on_error(&self, _ctx: &mut RunContext, _error: &TinyAgentsError) -> Result<()> { + async fn on_error(&self, ctx: &mut RunContext, _error: &TinyAgentsError) -> Result<()> { // A model call that fails (retries/fallback exhausted, hard provider // error, middleware timeout, ...) short-circuits with `?` before // `after_model` ever runs, so the reservation `before_model` added to @@ -331,12 +346,7 @@ impl Middleware for BudgetMidd // here so a run of failures cannot permanently inflate // `reserved_input_total` and starve every future call on a // process-lifetime-shared `BudgetTracker`. - let reserved = std::mem::take( - &mut *self - .pending_reservation - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner), - ); + let reserved = self.take_reservation(ctx.instance_id()); if reserved > 0 { let mut guard = self.tracker.lock_recovering(); guard.reserved_input_total = guard.reserved_input_total.saturating_sub(reserved); diff --git a/src/harness/middleware/library/test.rs b/src/harness/middleware/library/test.rs index 3daee75..2a68fc4 100644 --- a/src/harness/middleware/library/test.rs +++ b/src/harness/middleware/library/test.rs @@ -804,6 +804,48 @@ async fn reservation_is_released_when_model_call_errors() { .expect("reservation was released so a new call fits the budget again"); } +#[tokio::test] +async fn concurrent_runs_on_one_middleware_release_their_own_reservations() { + // Regression test: one `BudgetMiddleware` serves every run on the harness it + // is registered with (`invoke` takes `&self`), so a single scalar slot for + // the outstanding reservation was clobbered whenever two runs interleaved — + // one release subtracted the sibling's amount and the other subtracted + // nothing, permanently inflating `reserved_input_total` until the preflight + // rejected every future call. Both runs deliberately share a `run_id`, as + // `invoke_default` mints the same one for every run. + use crate::harness::middleware::Middleware; + + let mw = BudgetMiddleware::new(BudgetLimits::default()); + let tracker = mw.tracker(); + let mut small: RunContext = RunContext::new(RunConfig::new("run"), ()); + let mut large: RunContext = RunContext::new(RunConfig::new("run"), ()); + + let mut small_req = ModelRequest::new(vec![Message::user("x".repeat(40))]); + let mut large_req = ModelRequest::new(vec![Message::user("y".repeat(400))]); + mw.before_model(&mut small, &(), &mut small_req) + .await + .unwrap(); + mw.before_model(&mut large, &(), &mut large_req) + .await + .unwrap(); + let reserved = tracker.snapshot().reserved_input_total; + assert_eq!(reserved, 110, "both reservations are on the shared tracker"); + + let mut response = ModelResponse::assistant("ok"); + mw.after_model(&mut small, &(), &mut response) + .await + .unwrap(); + mw.after_model(&mut large, &(), &mut response) + .await + .unwrap(); + + assert_eq!( + tracker.snapshot().reserved_input_total, + 0, + "each run must release exactly what it reserved" + ); +} + #[test] fn poisoned_tracker_stays_fail_closed() { // A poisoned mutex still holds a valid last-written spend value; a diff --git a/src/harness/middleware/library/types.rs b/src/harness/middleware/library/types.rs index 65205a4..9f2e3f3 100644 --- a/src/harness/middleware/library/types.rs +++ b/src/harness/middleware/library/types.rs @@ -213,9 +213,9 @@ pub struct BudgetSpend { /// reconciled in `after_model`. Shared trackers (handed to concurrent /// sub-agent runs) can have more than one outstanding reservation at /// once, so this is a running total, not a single call's estimate; - /// each in-flight call's own reservation is tracked separately by its - /// [`BudgetMiddleware`] instance and released from this total when it - /// reconciles (or is abandoned). + /// each in-flight call's own reservation is tracked separately, per run, by + /// its [`BudgetMiddleware`] and released from this total when it reconciles + /// (or is abandoned). pub reserved_input_total: u64, } @@ -243,11 +243,19 @@ pub struct BudgetMiddleware { pub(crate) limits: BudgetLimits, pub(crate) tracker: BudgetTracker, pub(crate) pricing: std::collections::HashMap, - /// This run's own outstanding preflight reservation (input tokens), - /// awaiting reconciliation in `after_model`. Local to this middleware - /// instance (one per run) so concurrent runs sharing the same - /// [`BudgetTracker`] never clobber each other's reservation. - pub(crate) pending_reservation: std::sync::Mutex, + /// Outstanding preflight reservations (input tokens) awaiting + /// reconciliation in `after_model`, keyed by + /// [`RunContext::instance_id`][crate::harness::context::RunContext::instance_id]. + /// + /// One middleware instance serves every run on the harness it is registered + /// with (`invoke` takes `&self`), so a single scalar here would be + /// overwritten whenever two runs interleave their model calls — each would + /// then release the other's amount and permanently skew + /// [`BudgetTracker::reserved_input_total`][crate::harness::middleware::BudgetSpend::reserved_input_total]. + /// Keying by the run context's process-unique instance id (not its + /// caller-supplied `run_id`, which concurrent runs may share) keeps each + /// run releasing exactly what it reserved. + pub(crate) pending_reservations: std::sync::Mutex>, } // ── ToolPolicyMiddleware ────────────────────────────────────────────────────── diff --git a/src/harness/middleware/mod.rs b/src/harness/middleware/mod.rs index 294c704..ba4df6e 100644 --- a/src/harness/middleware/mod.rs +++ b/src/harness/middleware/mod.rs @@ -147,7 +147,12 @@ impl MiddlewareStack { /// Fans `on_error` out to every middleware, ignoring their results so the /// original error is never masked. No start/completed events are emitted on /// this internal recovery path. + /// + /// Marks the context so a driver that also handles the propagated error + /// (the agent loop does) skips its own dispatch: one failure must deliver + /// exactly one `on_error` per middleware. async fn fan_out_on_error(&self, ctx: &mut RunContext, error: &TinyAgentsError) { + ctx.mark_on_error_dispatched(); for mw in self.middlewares.iter() { let _ = mw.on_error(ctx, error).await; } diff --git a/src/harness/model/mod.rs b/src/harness/model/mod.rs index 54730d6..15789c3 100644 --- a/src/harness/model/mod.rs +++ b/src/harness/model/mod.rs @@ -844,11 +844,20 @@ impl StreamAccumulator { let tool_calls = self .tool_chunks .into_iter() - .map(|(id, args, name)| ToolCall { - name: name.unwrap_or_default(), - arguments: serde_json::from_str(&args).unwrap_or(Value::Null), - id, - invalid: None, + .map(|(id, args, name)| { + let name = name.unwrap_or_default(); + // Mirror the non-streaming adapters' argument semantics: an + // empty fragment is a well-formed zero-argument call, while an + // unparseable one is marked invalid with the raw text preserved + // so the agent loop can feed the parse error back to the model + // instead of silently executing a `null`/`{}` argument set. + if args.trim().is_empty() { + return ToolCall::new(id, name, Value::Object(serde_json::Map::new())); + } + match serde_json::from_str(&args) { + Ok(arguments) => ToolCall::new(id, name, arguments), + Err(err) => ToolCall::invalid(id, name, args, err.to_string()), + } }) .collect(); let message = AssistantMessage { diff --git a/src/harness/model/test.rs b/src/harness/model/test.rs index bc163d4..667102e 100644 --- a/src/harness/model/test.rs +++ b/src/harness/model/test.rs @@ -775,6 +775,50 @@ fn finish_names_reconstructed_tool_call_from_the_call_opening_delta_name() { assert_eq!(call.arguments, serde_json::json!({ "q": "rust" })); } +#[test] +fn finish_marks_unparseable_reconstructed_tool_arguments_invalid() { + // A stream cut mid-arguments (or a model emitting malformed JSON) used to + // reconstruct as `arguments: null` with `invalid: None`, hiding the parse + // failure from the agent loop's tool-error recovery path — which then either + // aborted the run on schema validation or executed the tool with `{}`. + use crate::harness::tool::ToolDelta; + + let mut acc = StreamAccumulator::new(); + acc.push(&ModelStreamItem::ToolCallDelta(ToolDelta { + call_id: "call-1".into(), + content: r#"{"path": "/tm"#.into(), + tool_name: Some("read_file".into()), + })); + + let finished = acc.finish().unwrap(); + let call = &finished.message.tool_calls[0]; + assert!(call.is_invalid(), "unparseable arguments must be flagged"); + assert_eq!( + call.arguments, + serde_json::Value::String(r#"{"path": "/tm"#.to_string()), + "the raw fragment is preserved for the model to correct" + ); +} + +#[test] +fn finish_reconstructs_empty_tool_arguments_as_an_empty_object() { + // A zero-argument call whose deltas carried no argument text is well-formed, + // not malformed: it must reconstruct as `{}` rather than an invalid call. + use crate::harness::tool::ToolDelta; + + let mut acc = StreamAccumulator::new(); + acc.push(&ModelStreamItem::ToolCallDelta(ToolDelta { + call_id: "call-1".into(), + content: String::new(), + tool_name: Some("ping".into()), + })); + + let finished = acc.finish().unwrap(); + let call = &finished.message.tool_calls[0]; + assert!(!call.is_invalid()); + assert_eq!(call.arguments, serde_json::json!({})); +} + /// Round-trips a [`ModelStreamItem`] through JSON and asserts the re-serialized /// form is byte-for-byte stable, proving every variant survives serde. fn roundtrip_stream_item(item: ModelStreamItem) { diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 1328d4b..b46a2a8 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -100,6 +100,7 @@ use sse::*; use transport::{ Degrade, auth_headers, degrade_for_400, effective_temperature, glob_match, is_stream_required_error, merge_provider_options, merge_system_into_user, request_timeout, + unary_fold_timeout_ms, }; #[cfg(test)] diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index c06b02e..487bfef 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -1744,6 +1744,66 @@ fn request_timeout_defaults_by_call_kind() { assert_eq!(request_timeout(None, true), None); } +#[test] +fn unary_fold_timeout_ms_caps_a_degraded_unary_call_at_the_default() { + // Before the fix, `invoke_with_streaming` passed the request's `timeout_ms` + // through unchanged, and `stream()` resolves an unset one to `None` overall + // deadline. A unary `invoke` call degraded onto the streaming path must + // still be capped, so a provider that returns `200 text/event-stream` and + // then stops sending bytes cannot hang the caller forever. + assert_eq!( + unary_fold_timeout_ms(None, false), + Some(DEFAULT_REQUEST_TIMEOUT_SECS * 1_000) + ); + // A genuine `stream()` caller is unaffected: this function is only + // consulted on the degraded-unary path. +} + +#[test] +fn unary_fold_timeout_ms_prefers_an_explicit_override() { + assert_eq!(unary_fold_timeout_ms(Some(1_500), false), Some(1_500)); + assert_eq!(unary_fold_timeout_ms(Some(1_500), true), Some(1_500)); +} + +#[test] +fn unary_fold_timeout_ms_leaves_caller_owned_clients_uncapped() { + // `with_client` callers manage their own client-level timeout policy + // (see `effective_request_timeout`'s caller-owned-client opt-out); the + // degraded-unary injection must not override that. + assert_eq!(unary_fold_timeout_ms(None, true), None); +} + +#[test] +fn list_models_request_carries_the_default_timeout() { + // Regression: `list_models` used to build its GET request with no + // `.timeout(...)` at all, so a reachable-but-wedged endpoint (TCP connect + // succeeds, then the server never responds) hung the call forever. Every + // other outbound path resolves a deadline via `effective_request_timeout`; + // this inspects the actual `reqwest::Request` `list_models` now builds. + let m = model(); + let built = m + .list_models_request(&format!("{}/models", m.base_url())) + .build() + .expect("request builds"); + assert_eq!( + built.timeout(), + Some(&Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS)), + "list_models must apply the default unary request timeout" + ); +} + +#[test] +fn list_models_request_respects_caller_owned_client_opt_out() { + // A caller-owned client (`with_client`) manages its own timeout policy; + // `list_models` must not override it with a timeout of its own. + let m = model().with_client(reqwest::Client::new()); + let built = m + .list_models_request(&format!("{}/models", m.base_url())) + .build() + .expect("request builds"); + assert_eq!(built.timeout(), None); +} + #[test] fn from_env_errors_when_api_key_missing() { // Snapshot and clear the key so the missing-key path is exercised @@ -2539,6 +2599,74 @@ fn degrade_for_400_unions_with_existing_baseline_degrade() { ); } +// The named-tool-choice / json_object 400-driven degradation used to be +// discovered by `degrade_for_400` and applied only to the one retry, then +// thrown away — `named_tool_choice_supported`/`json_object_format_supported` +// were plain `bool`s only the builder ever wrote. `post_chat_with_degrade` +// takes `&self`, so every later call to the same rejecting endpoint replayed +// the un-degraded baseline body and paid a guaranteed second 400. This mirrors +// `stream_required_constraint_latches_after_discovery` for the two request-shape +// knobs. +#[test] +fn shape_degrade_latches_after_discovery_so_baseline_is_already_degraded() { + let m = model(); + assert_eq!( + m.baseline_degrade(), + Degrade::default(), + "must start un-latched" + ); + + m.latch_degrade(Degrade { + named_tool_choice: true, + json_object: false, + }); + assert_eq!( + m.baseline_degrade(), + Degrade { + named_tool_choice: true, + json_object: false, + }, + "a discovered named_tool_choice rejection must be remembered so the next \ + call's baseline body is already degraded, instead of re-paying the 400" + ); + + // Idempotent, and unions rather than clobbers: a later json_object discovery + // keeps the earlier named_tool_choice latch. + m.latch_degrade(Degrade { + named_tool_choice: true, + json_object: true, + }); + assert_eq!( + m.baseline_degrade(), + Degrade { + named_tool_choice: true, + json_object: true, + } + ); +} + +#[test] +fn shape_degrade_latch_survives_through_a_shared_handle() { + // Same reasoning as `stream_required_latch_survives_through_a_shared_handle`: + // production holds models as `Arc`. + let shared: std::sync::Arc = std::sync::Arc::new(model()); + let clone = std::sync::Arc::clone(&shared); + assert_eq!(clone.baseline_degrade(), Degrade::default()); + + shared.latch_degrade(Degrade { + named_tool_choice: false, + json_object: true, + }); + assert_eq!( + clone.baseline_degrade(), + Degrade { + named_tool_choice: false, + json_object: true, + }, + "the latch must be visible to every holder of the shared model" + ); +} + #[test] fn stream_cleanup_scrubs_leaked_markup_from_live_deltas() { use crate::harness::message::{AssistantMessage, ContentBlock, MessageDelta}; diff --git a/src/harness/providers/openai/transport.rs b/src/harness/providers/openai/transport.rs index 3204e27..b35e29f 100644 --- a/src/harness/providers/openai/transport.rs +++ b/src/harness/providers/openai/transport.rs @@ -66,12 +66,21 @@ pub struct OpenAiModel { /// `tool_choice:"required"` with the `tools` array filtered to the named tool /// — some local runtimes (LM Studio, llama.cpp server) 400 on the object form. /// See [`Self::with_named_tool_choice`]. - named_tool_choice_supported: bool, + /// + /// An [`AtomicBool`] (like [`Self::stream_required`]) because a 400 that + /// implicates this shape latches it `false` in [`Self::post_chat_with_degrade`] + /// so later calls on the same shared `&self` skip the doomed baseline + /// request instead of re-discovering the rejection every time. + named_tool_choice_supported: AtomicBool, /// Whether the endpoint accepts `response_format:{"type":"json_object"}`. /// `true` by default. When `false`, a [`ResponseFormat::JsonObject`] request is /// degraded to a permissive `json_schema` wire form — some local runtimes 400 /// on `json_object`. See [`Self::with_json_object_format`]. - json_object_format_supported: bool, + /// + /// An [`AtomicBool`] for the same reason as + /// [`Self::named_tool_choice_supported`]: the 400-driven discovery is latched + /// on the instance instead of thrown away after the retry. + json_object_format_supported: AtomicBool, /// Default model id used when a request does not override it. model: String, /// Provider family identifier used in profiles and normalized errors. @@ -339,8 +348,8 @@ impl OpenAiModel { temperature_unsupported: Vec::new(), temperature_override: None, merge_system_into_user: false, - named_tool_choice_supported: true, - json_object_format_supported: true, + named_tool_choice_supported: AtomicBool::new(true), + json_object_format_supported: AtomicBool::new(true), model: DEFAULT_MODEL.to_string(), provider: "openai".to_string(), base_url: DEFAULT_BASE_URL.to_string(), @@ -399,6 +408,37 @@ impl OpenAiModel { } } + /// Remembers a newly-discovered request-shape rejection so later calls skip + /// straight to the degraded body instead of re-paying a guaranteed 400. + /// + /// Mirrors [`Self::latch_stream_required`]: `degrade` is the union + /// [`degrade_for_400`] just computed, so this stores `false` into exactly + /// the knob(s) that flipped on this call. Logs on each transition only. + pub(super) fn latch_degrade(&self, degrade: Degrade) { + if degrade.named_tool_choice + && self + .named_tool_choice_supported + .swap(false, Ordering::Relaxed) + { + tracing::info!( + provider = %self.provider, + model = %self.model, + "[openai] provider rejects named tool_choice; latching degraded shape for subsequent calls" + ); + } + if degrade.json_object + && self + .json_object_format_supported + .swap(false, Ordering::Relaxed) + { + tracing::info!( + provider = %self.provider, + model = %self.model, + "[openai] provider rejects response_format:json_object; latching degraded shape for subsequent calls" + ); + } + } + /// Routes calls to the OpenAI **Responses API** (`/v1/responses`) instead of /// Chat Completions. Required for the OpenAI Codex OAuth backend; pair with /// [`with_extra_query_param`](Self::with_extra_query_param) + @@ -494,9 +534,12 @@ impl OpenAiModel { /// then degraded to `tool_choice:"required"` with the wire `tools` array /// filtered down to just the named tool, preserving the "must call *this* /// tool" semantics. Independent of this flag, a 400 whose body implicates - /// `tool_choice` triggers the same degraded retry automatically (once). - pub fn with_named_tool_choice(mut self, supported: bool) -> Self { - self.named_tool_choice_supported = supported; + /// `tool_choice` triggers the same degraded retry automatically — and, once + /// discovered, is latched on the instance so later calls skip the doomed + /// baseline request instead of re-paying the 400 every time. + pub fn with_named_tool_choice(self, supported: bool) -> Self { + self.named_tool_choice_supported + .store(supported, Ordering::Relaxed); self } @@ -508,9 +551,12 @@ impl OpenAiModel { /// [`ResponseFormat::JsonObject`] request is then degraded to a permissive /// `json_schema` wire form (an empty object schema with `strict:false`). /// Independent of this flag, a 400 whose body implicates `response_format` - /// triggers the same degraded retry automatically (once). - pub fn with_json_object_format(mut self, supported: bool) -> Self { - self.json_object_format_supported = supported; + /// triggers the same degraded retry automatically — and, once discovered, + /// is latched on the instance so later calls skip the doomed baseline + /// request instead of re-paying the 400 every time. + pub fn with_json_object_format(self, supported: bool) -> Self { + self.json_object_format_supported + .store(supported, Ordering::Relaxed); self } @@ -749,9 +795,8 @@ impl OpenAiModel { /// decoded. pub async fn list_models(&self) -> Result> { let url = format!("{}/models", self.base_url); - let response = self - .send_checked(self.authorized(self.client.get(&url)), "request", &url) + .send_checked(self.list_models_request(&url), "request", &url) .await?; let text = response.text().await.map_err(|e| { @@ -954,8 +999,8 @@ impl OpenAiModel { /// shape on the wire". pub(super) fn baseline_degrade(&self) -> Degrade { Degrade { - named_tool_choice: !self.named_tool_choice_supported, - json_object: !self.json_object_format_supported, + named_tool_choice: !self.named_tool_choice_supported.load(Ordering::Relaxed), + json_object: !self.json_object_format_supported.load(Ordering::Relaxed), } } @@ -1137,6 +1182,25 @@ impl OpenAiModel { builder } + /// Builds the authorized GET request for [`Self::list_models`], applying + /// the same deadline policy as every other outbound call. + /// + /// Every other outbound path (`post_json`, `send_responses`) resolves a + /// deadline through [`Self::effective_request_timeout`]; `list_models` used + /// not to, so a reachable-but-wedged endpoint (an Ollama/LM Studio server + /// that accepts the TCP connect and then never responds — exactly the + /// runtime model discovery this method exists for) hung the call forever. + /// Factored out (rather than inlined in `list_models`) so the timeout + /// policy is unit-testable via [`reqwest::RequestBuilder::build`] without a + /// network round trip. + pub(super) fn list_models_request(&self, url: &str) -> reqwest::RequestBuilder { + let mut builder = self.authorized(self.client.get(url)); + if let Some(timeout) = self.effective_request_timeout(None, false) { + builder = builder.timeout(timeout); + } + builder + } + /// The `/responses` endpoint URL — a sibling of `/chat/completions` under the /// same base URL, tolerating a base that already ends in `/responses` or a /// `…/v1` chat base. @@ -1320,6 +1384,7 @@ impl OpenAiModel { if err.status == Some(400) && let Some(degrade) = degrade_for_400(&err.message, request, baseline) => { + self.latch_degrade(degrade); let retry = self.build_chat_body(request, degrade, streaming)?; self.post_json(&retry, request.timeout_ms, streaming, what) .await @@ -1558,6 +1623,28 @@ pub(super) fn request_timeout(timeout_ms: Option, streaming: bool) -> Optio } } +/// Resolves the `timeout_ms` [`invoke_with_streaming`] should carry when a +/// unary [`ChatModel::invoke`] call degrades onto the streaming wire path. +/// +/// An explicit `timeout_ms` always wins (unchanged). Otherwise the caller's +/// mode is unary, not the wire mode — so unlike [`request_timeout`] with +/// `streaming: true`, an unset `timeout_ms` here must **not** resolve to `None`: +/// `stream()` folds `None` into no overall deadline, which would let a provider +/// that returns `200 text/event-stream` and then stops sending bytes hang the +/// unary caller forever. Skipped for caller-owned clients (`with_client`), +/// which manage their own client-level timeout policy and opt out of +/// [`OpenAiModel::effective_request_timeout`] the same way. +pub(super) fn unary_fold_timeout_ms( + timeout_ms: Option, + caller_owned_client: bool, +) -> Option { + match timeout_ms { + Some(ms) => Some(ms), + None if caller_owned_client => None, + None => Some(DEFAULT_REQUEST_TIMEOUT_SECS * 1_000), + } +} + /// Merges baked `defaults` under a request's own `overrides` provider options. /// /// Keys present in `overrides` win over `defaults`. A `Null` on either side @@ -1898,8 +1985,14 @@ pub(super) fn clean_stream_item( /// `ChatModel` impl (the stream method never reads it). async fn invoke_with_streaming( model: &OpenAiModel, - request: ModelRequest, + mut request: ModelRequest, ) -> Result { + // `invoke` is a unary call to its caller and is documented (and, below, + // tested) to be capped at `DEFAULT_REQUEST_TIMEOUT_SECS` when the caller did + // not set an explicit `timeout_ms`. Folding onto the streaming wire path + // must not silently drop that cap — see `unary_fold_timeout_ms`. + request.timeout_ms = unary_fold_timeout_ms(request.timeout_ms, model.caller_owned_client); + let mut stream = model.stream(&(), request).await?; let mut acc = StreamAccumulator::new(); diff --git a/src/harness/subagent/mod.rs b/src/harness/subagent/mod.rs index 48d30e0..362f6ce 100644 --- a/src/harness/subagent/mod.rs +++ b/src/harness/subagent/mod.rs @@ -226,7 +226,11 @@ impl SubAgent { parent.thread_id(), parent.config.max_turn_output_tokens, )?; - let ctx = RunContext::new(config, ctx_data).with_events(parent.events.clone()); + // Share the parent's cancellation token so one `cancel()` unwinds the + // whole nested-run tree instead of stopping at this boundary. + let ctx = RunContext::new(config, ctx_data) + .with_events(parent.events.clone()) + .with_cancellation(parent.cancellation.clone()); self.run_child(state, ctx, input.into(), parent.streaming) .await } @@ -593,7 +597,12 @@ where return Err(error); } }; - let ctx = RunContext::new(config, Ctx::default()).with_events(context.events); + // Inherit the caller's cancellation token: a cancel requested on the + // parent run must also stop the child loop this tool drives, otherwise + // it runs to completion while the parent waits on this `await`. + let ctx = RunContext::new(config, Ctx::default()) + .with_events(context.events) + .with_cancellation(context.cancellation); // Match the parent's drive mode: when the parent run streams, the child // streams too, so its deltas flow onto the shared sink and reach the // parent's `invoke_stream` consumer with the child's own lineage. diff --git a/src/harness/subagent/test.rs b/src/harness/subagent/test.rs index 4ff5e65..4d5d7b5 100644 --- a/src/harness/subagent/test.rs +++ b/src/harness/subagent/test.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use serde_json::json; use crate::error::TinyAgentsError; +use crate::harness::cancel::CancellationToken; use crate::harness::context::{RunConfig, RunContext}; use crate::harness::events::{AgentEvent, EventSink, RecordingListener}; use crate::harness::limits::RunLimits; @@ -99,6 +100,35 @@ impl Tool<()> for SpinTool { } } +/// A child-side tool that counts its invocations and cancels the shared token +/// on the first one, so the child's next loop checkpoint must observe it. +struct CancelOnFirstCallTool { + calls: Arc, + token: CancellationToken, +} + +#[async_trait::async_trait] +impl Tool<()> for CancelOnFirstCallTool { + fn name(&self) -> &str { + "spin" + } + + fn description(&self) -> &str { + "cancels the run on its first call" + } + + fn schema(&self) -> ToolSchema { + ToolSchema::new("spin", "cancels the run on its first call", json!({})) + } + + async fn call(&self, _state: &(), call: ToolCall) -> crate::Result { + if self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 { + self.token.cancel(); + } + Ok(ToolResult::text(call.id, "spin", "again")) + } +} + fn looping_child_harness_with_max_model_calls(max_model_calls: usize) -> AgentHarness<()> { let mut harness: AgentHarness<()> = AgentHarness::new(); harness @@ -328,6 +358,82 @@ async fn parent_can_continue_after_subagent_tool_hits_child_limit() { ); } +#[tokio::test] +async fn call_with_context_propagates_parent_cancellation_into_the_child() { + // Regression test: the child run used to be built with a fresh, never + // cancelled `CancellationToken`, so cancelling the parent's token was a + // no-op for the whole sub-agent delegation — the child ran until it hit a + // limit while the parent was blocked awaiting this tool. + let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let token = CancellationToken::new(); + + let mut child: AgentHarness<()> = AgentHarness::new(); + child + .register_model( + "child-model", + Arc::new(MockModel::with_tool_call("spin", json!({}))), + ) + .register_tool(Arc::new(CancelOnFirstCallTool { + calls: calls.clone(), + token: token.clone(), + })) + // Bound the child so a non-propagating token fails the assertion + // instead of looping forever. + .with_policy(RunPolicy { + limits: RunLimits::default().with_max_model_calls(5), + ..RunPolicy::default() + }); + + let tool = SubAgentTool::new(Arc::new(SubAgent::new( + "worker", + "spins until cancelled", + Arc::new(child), + ))); + + let parent_ctx: RunContext<()> = + RunContext::new(RunConfig::new("parent"), ()).with_cancellation(token); + let context = ToolExecutionContext::from_run_context(&parent_ctx); + + let error = tool + .call_with_context( + &(), + ToolCall::new("c1", "worker", json!({ "input": "spin" })), + context, + ) + .await + .expect_err("the cancelled child unwinds instead of running to its limit"); + + assert!( + matches!(error, TinyAgentsError::Cancelled), + "unexpected child error: {error}" + ); + assert_eq!( + calls.load(std::sync::atomic::Ordering::SeqCst), + 1, + "the child should stop at its next checkpoint after the cancel" + ); +} + +#[tokio::test] +async fn invoke_in_parent_propagates_parent_cancellation_into_the_child() { + let subagent = SubAgent::new("worker", "does work", Arc::new(child_harness("done"))); + + let token = CancellationToken::new(); + token.cancel(); + let parent_ctx: RunContext<()> = + RunContext::new(RunConfig::new("parent"), ()).with_cancellation(token); + + let error = subagent + .invoke_in_parent(&(), (), &parent_ctx, "go") + .await + .expect_err("a child started under a cancelled parent token stops immediately"); + + assert!( + matches!(error, TinyAgentsError::Cancelled), + "unexpected child error: {error}" + ); +} + #[tokio::test] async fn invoke_with_events_emits_lifecycle_on_shared_sink() { let subagent = SubAgent::new( diff --git a/src/harness/tool/prompt.rs b/src/harness/tool/prompt.rs index 55cbc47..3a5142e 100644 --- a/src/harness/tool/prompt.rs +++ b/src/harness/tool/prompt.rs @@ -485,14 +485,37 @@ pub fn apply_prompt_tool_calls(mut response: ModelResponse) -> ModelResponse { return response; } response.message.tool_calls.extend(calls); - response.message.content = if cleaned.is_empty() { - Vec::new() - } else { - vec![ContentBlock::Text(cleaned)] - }; + response.message.content = replace_text_blocks(response.message.content, cleaned); response } +/// Rebuild a content vector, keeping every non-[`ContentBlock::Text`] block (e.g. +/// `Thinking`) in place and substituting the single cleaned text at the position +/// of the first original `Text` block. If the original content had no `Text` +/// block, the cleaned text (when non-empty) is appended; if `cleaned` is empty, +/// no text block is emitted at all. +fn replace_text_blocks(content: Vec, cleaned: String) -> Vec { + let mut out = Vec::with_capacity(content.len()); + let mut inserted = false; + for block in content { + match block { + ContentBlock::Text(_) => { + if !inserted { + if !cleaned.is_empty() { + out.push(ContentBlock::Text(cleaned.clone())); + } + inserted = true; + } + } + other => out.push(other), + } + } + if !inserted && !cleaned.is_empty() { + out.push(ContentBlock::Text(cleaned)); + } + out +} + /// Parse a single tool-call body into a [`ToolCall`] with a synthetic 1-based id. fn parse_one(inner: &str, index: usize) -> Option { let value: Value = serde_json::from_str(inner).ok()?; diff --git a/src/harness/tool/prompt_test.rs b/src/harness/tool/prompt_test.rs index e6fa9fb..ef5d5c2 100644 --- a/src/harness/tool/prompt_test.rs +++ b/src/harness/tool/prompt_test.rs @@ -2,6 +2,7 @@ use super::*; use crate::harness::message::{ContentBlock, ImageRef, Message}; +use crate::harness::model::ModelResponse; fn schema(name: &str) -> ToolSchema { ToolSchema { @@ -476,3 +477,36 @@ fn scrubber_matches_batch_parser_on_the_visible_text() { let refs: Vec<&str> = frags.iter().map(String::as_str).collect(); assert_eq!(scrub_all(&refs).trim(), batch); } + +#[test] +fn apply_prompt_tool_calls_preserves_a_leading_thinking_block() { + // A prompt-guided reasoning model emits a `Thinking` block followed by the + // `` text. Recovering the call must not discard the reasoning. + let mut response = ModelResponse::assistant( + r#"reply {"name":"search","arguments":{"q":"x"}}"#, + ); + response.message.content.insert( + 0, + ContentBlock::Thinking { + text: "chain of thought".to_string(), + signature: None, + }, + ); + + let out = apply_prompt_tool_calls(response); + + assert_eq!(out.message.tool_calls.len(), 1); + assert_eq!(out.message.tool_calls[0].name, "search"); + assert_eq!( + out.message.content[0], + ContentBlock::Thinking { + text: "chain of thought".to_string(), + signature: None, + }, + "the thinking block must survive the content rebuild" + ); + assert_eq!( + out.message.content[1], + ContentBlock::Text("reply".to_string()) + ); +} diff --git a/src/harness/tool/types.rs b/src/harness/tool/types.rs index 042f2f8..a27a32b 100644 --- a/src/harness/tool/types.rs +++ b/src/harness/tool/types.rs @@ -16,6 +16,7 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::Result; +use crate::harness::cancel::CancellationToken; use crate::harness::context::RunContext; use crate::harness::events::EventSink; use crate::harness::ids::{RunId, ThreadId}; @@ -186,6 +187,11 @@ pub struct ToolExecutionContext { pub max_turn_output_tokens: Option, /// Shared event sink for nested run observability. pub events: EventSink, + /// The caller run's cancellation token. A recursive tool such as a + /// sub-agent installs this on its child run so one `cancel()` unwinds the + /// whole nested-run tree, as + /// [`crate::harness::cancel`] documents. + pub cancellation: CancellationToken, /// Whether the caller run is being driven through the streaming loop path. /// A sub-agent tool uses this to run its child in the matching mode so the /// child's deltas propagate onto the shared [`EventSink`]. @@ -207,6 +213,7 @@ impl ToolExecutionContext { depth: ctx.config.depth, max_turn_output_tokens: ctx.config.max_turn_output_tokens, events: ctx.events.clone(), + cancellation: ctx.cancellation.clone(), streaming: ctx.streaming, workspace: ctx.workspace.clone(), } diff --git a/src/language/ast.rs b/src/language/ast.rs index 36eecb2..eb31257 100644 --- a/src/language/ast.rs +++ b/src/language/ast.rs @@ -156,10 +156,31 @@ pub struct NodeDecl { pub retry: Vec<(String, Literal)>, /// Node-level `metadata { key value … }` entries. pub metadata: Vec<(String, Literal)>, + /// A `steering { … }` policy declaration for a `subagent` node. + pub steering: Option, /// Source position of the `node` keyword. pub span: Span, } +/// A `steering { parent allow [...] human allow [...] delivery "…" }` +/// declaration: the narrowed set of steering commands a parent orchestrator +/// or a human may send to this `subagent` node, plus a delivery policy. +/// +/// Parsing accepts the shape documented in +/// `docs/modules/expressive-language/reference.md`; lowering it into a +/// runtime `harness::steering` policy is not yet implemented (tracked +/// separately), so [`crate::language::compiler::compile`] currently parses +/// and retains this declaration without acting on it. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SteeringDecl { + /// Steering commands a parent orchestrator run may send (`parent allow [...]`). + pub parent_allow: Vec, + /// Steering commands a human may send (`human allow [...]`). + pub human_allow: Vec, + /// The delivery policy name (`delivery "safe_boundary"`), if declared. + pub delivery: Option, +} + impl NodeDecl { /// Creates an empty node declaration with the given name and span and all /// optional fields unset. Keeps the parser's construction site concise as @@ -185,6 +206,7 @@ impl NodeDecl { timeout: None, retry: Vec::new(), metadata: Vec::new(), + steering: None, span, } } diff --git a/src/language/capability_resolver.rs b/src/language/capability_resolver.rs index e28c22d..f94cb13 100644 --- a/src/language/capability_resolver.rs +++ b/src/language/capability_resolver.rs @@ -295,6 +295,36 @@ impl CapabilityResolver { Some(PrimaryReference { class, target }) } + /// The secondary `model` reference a node may carry alongside its primary, + /// kind-specific reference. + /// + /// `subagent` and `repl_agent` nodes both accept a documented `model` + /// field in addition to their primary `agent`/`script` reference (for + /// `repl_agent` it is the model-driven CodeAct form), so that value must + /// be validated against the model allowlist too — it is never the + /// primary target for those kinds. `subgraph`/`graph` nodes only carry a + /// secondary model check when a *dedicated* subgraph/graph field was + /// used (`dedicated_subgraph_field`); otherwise `model` is itself the + /// fallback primary target already covered by + /// [`classify_reference`](Self::classify_reference). Every other kind + /// resolves `model` as its *primary* reference, so it has no secondary + /// check here. + /// + /// Shares the same "edit the policy once" shape as `classify_reference` + /// so [`bind_blueprint`](Self::bind_blueprint) and both + /// [`crate::language::resolver::Resolver`] paths stay in lockstep. + pub fn secondary_model_reference<'a>( + kind: &str, + model: Option<&'a str>, + dedicated_subgraph_field: bool, + ) -> Option<&'a str> { + match kind { + "subagent" | "repl_agent" => model, + "subgraph" | "graph" if dedicated_subgraph_field => model, + _ => None, + } + } + /// Returns true when `target` is allowed for the given reference `class`. pub fn reference_allowed(&self, class: ReferenceClass, target: &str) -> bool { match class { @@ -324,6 +354,10 @@ impl CapabilityResolver { /// references to a registered agent, `repl_agent` node references to a /// registered script, and all other nodes' `model` references to a /// registered model (via the shared [`classify_reference`](Self::classify_reference) policy); + /// - a `subagent`/`repl_agent` node's `model` field (and a `subgraph`/`graph` + /// node's `model` field when a dedicated subgraph field was also used) + /// also resolves to a registered model (via + /// [`secondary_model_reference`](Self::secondary_model_reference)); /// - every `channel` reducer reference is registered. /// /// # Errors @@ -359,6 +393,18 @@ impl CapabilityResolver { ))); } + if let Some(model) = Self::secondary_model_reference( + &node.kind, + node.model.as_deref(), + node.subgraph.is_some(), + ) && !self.model_allowed(model) + { + return Err(TinyAgentsError::Capability(format!( + "node `{}` references unknown model `{}`", + node.name, model + ))); + } + for tool in &node.tools { if !self.tool_allowed(tool) { return Err(TinyAgentsError::Capability(format!( diff --git a/src/language/compiler.rs b/src/language/compiler.rs index 0da34be..ea22609 100644 --- a/src/language/compiler.rs +++ b/src/language/compiler.rs @@ -55,6 +55,17 @@ use crate::registry::CapabilityRegistry; /// /// All failures are reported as [`TinyAgentsError::Compile`]. pub fn compile(program: &Program) -> Result> { + reject_duplicate_graphs(program)?; + program.graphs.iter().map(compile_graph).collect() +} + +/// Rejects a [`Program`] that declares the same graph name more than once. +/// +/// Shared by [`compile`] and [`compile_with_provenance`] so both entry points +/// enforce the same semantic guard — two blueprints sharing one `graph_id` +/// would otherwise collide when registered (`DuplicateComponent` mid-way +/// through, or a silent overwrite via `replace_graph_blueprint`). +fn reject_duplicate_graphs(program: &Program) -> Result<()> { let mut graph_ids: HashSet<&str> = HashSet::new(); for graph in &program.graphs { if !graph_ids.insert(graph.name.as_str()) { @@ -64,7 +75,7 @@ pub fn compile(program: &Program) -> Result> { ))); } } - program.graphs.iter().map(compile_graph).collect() + Ok(()) } fn compile_graph(graph: &crate::language::types::GraphDecl) -> Result { @@ -411,6 +422,7 @@ fn compile_graph(graph: &crate::language::types::GraphDecl) -> Result /// Returns [`TinyAgentsError::Compile`] for the same semantic failures as /// [`compile`]. pub fn compile_with_provenance(program: &Program, origin: Origin) -> Result> { + reject_duplicate_graphs(program)?; program .graphs .iter() diff --git a/src/language/parser.rs b/src/language/parser.rs index 4a0edbd..de84f85 100644 --- a/src/language/parser.rs +++ b/src/language/parser.rs @@ -17,7 +17,7 @@ use crate::language::diagnostic::Diagnostic; use crate::language::source::SourceFile; use crate::language::types::{ ChannelDecl, CommandDecl, EdgeDecl, GraphDecl, IoFieldDecl, JoinDecl, Literal, NodeDecl, - Program, RouteDecl, SendDecl, Span, SpannedToken, Token, + Program, RouteDecl, SendDecl, Span, SpannedToken, SteeringDecl, Token, }; /// Tokenises and parses `source` in one step. @@ -59,6 +59,20 @@ pub fn parse(tokens: &[SpannedToken]) -> Result { .with_primary_label("here") .into_parse_error(None)); } + // The cursor helpers (`current`/`advance`) assume the stream ends with an + // `Eof` sentinel so that `pos` can pin at the last index once exhausted. + // A slice missing that sentinel (e.g. produced by a caller filtering or + // truncating a token stream) would otherwise let the pinned cursor make + // zero progress forever in productions that can succeed on a no-op + // `advance()` — an unkillable hang rather than a parse error. + let last = tokens.last().expect("checked non-empty above"); + if !matches!(last.token, Token::Eof) { + return Err( + Diagnostic::error("token stream missing end-of-input sentinel", last.span) + .with_primary_label("here") + .into_parse_error(None), + ); + } Parser { tokens, pos: 0, @@ -482,6 +496,10 @@ impl Parser<'_> { self.advance(); node.metadata = self.parse_defaults_block()?; } + "steering" => { + self.advance(); + node.steering = Some(self.parse_steering_block()?); + } other => { return Err(self.error(format!("unknown node item `{other}`"), tok.span)); } @@ -533,6 +551,37 @@ impl Parser<'_> { Ok(CommandDecl { goto, update, span }) } + /// Parses a `steering { parent allow [...] human allow [...] delivery "…" }` + /// block. The `steering` keyword has already been consumed. + fn parse_steering_block(&mut self) -> Result { + self.expect(&Token::LBrace)?; + let mut steering = SteeringDecl::default(); + while !matches!(self.current().token, Token::RBrace) { + if self.at_eof() { + return Err(self.error("unexpected end of input inside `steering`", self.span())); + } + if self.is_keyword("parent") { + self.advance(); + self.expect_keyword("allow")?; + steering.parent_allow = self.parse_string_list()?; + } else if self.is_keyword("human") { + self.advance(); + self.expect_keyword("allow")?; + steering.human_allow = self.parse_string_list()?; + } else if self.is_keyword("delivery") { + self.advance(); + steering.delivery = Some(self.expect_string()?); + } else { + return Err(self.error( + "expected `parent`, `human`, or `delivery` inside `steering`", + self.span(), + )); + } + } + self.expect(&Token::RBrace)?; + Ok(steering) + } + /// Parses a `sends [ send ["input"] … ]` block. The `sends` keyword /// has already been consumed. fn parse_sends_block(&mut self) -> Result> { diff --git a/src/language/resolver.rs b/src/language/resolver.rs index 9e57b9b..a78f989 100644 --- a/src/language/resolver.rs +++ b/src/language/resolver.rs @@ -171,7 +171,26 @@ impl Resolver { ); } - // 3. Every referenced tool must be registered. + // 3. `subagent`/`repl_agent` nodes (and `subgraph`/`graph` nodes that + // used a dedicated subgraph field) may also carry a secondary + // `model` reference, checked against the model allowlist. + if let Some(model) = CapabilityResolver::secondary_model_reference( + kind, + node.model.as_deref(), + node.graph.is_some(), + ) { + self.check_ref( + self.caps.model_allowed(model), + &node.name, + "model", + model, + node.span, + CODE_UNKNOWN_MODEL, + out, + ); + } + + // 4. Every referenced tool must be registered. for tool in &node.tools { self.check_ref( self.caps.tool_allowed(tool), @@ -288,6 +307,14 @@ impl Resolver { reference.target, )); } + if let Some(model) = CapabilityResolver::secondary_model_reference( + &node.kind, + node.model.as_deref(), + node.subgraph.is_some(), + ) && !self.caps.model_allowed(model) + { + return Err(unregistered("model", &node.name, model)); + } for tool in &node.tools { if !self.caps.tool_allowed(tool) { return Err(unregistered("tool", &node.name, tool)); diff --git a/src/language/test/compiler.rs b/src/language/test/compiler.rs index f0d2f27..5a00fe8 100644 --- a/src/language/test/compiler.rs +++ b/src/language/test/compiler.rs @@ -125,6 +125,22 @@ fn duplicate_graph_id_is_a_compile_error() { assert!(err.to_string().contains("duplicate graph"), "{err}"); } +#[test] +fn compile_with_provenance_also_rejects_duplicate_graph_id() { + // `compile_with_provenance` documents "the same semantic validation and + // lowering as `compile`", so it must reject a duplicate graph name too — + // otherwise two distinct blueprints share one `graph_id` and collide when + // registered. + use crate::language::compiler::compile_with_provenance; + use crate::language::types::Origin; + + let src = "graph g { start a node a { } } graph g { start b node b { } }"; + let program = parse_str(src).unwrap(); + let err = compile_with_provenance(&program, Origin::generated()).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Compile(_))); + assert!(err.to_string().contains("duplicate graph"), "{err}"); +} + #[test] fn next_and_command_goto_conflict_is_a_compile_error() { let src = "graph g { start a node a { next b command { goto c } } node b { } node c { } }"; diff --git a/src/language/test/extended_grammar.rs b/src/language/test/extended_grammar.rs index ad0c0dd..773f29c 100644 --- a/src/language/test/extended_grammar.rs +++ b/src/language/test/extended_grammar.rs @@ -275,3 +275,130 @@ fn bind_blueprint_rejects_unregistered_subagent_and_script() { assert!(err.to_string().contains("unknown script"), "{err}"); assert!(err.to_string().contains("triage_script"), "{err}"); } + +#[test] +fn bind_blueprint_rejects_unregistered_secondary_model_on_subagent_and_repl_agent() { + // `subagent` and `repl_agent` nodes may carry a documented `model` field + // alongside their primary `agent`/`script` reference. The strict gate + // must validate it too, not silently let it through because only the + // primary reference was classified. + let node_kinds = || { + crate::language::capability_resolver::DEFAULT_NODE_KINDS + .iter() + .map(|k| k.to_string()) + }; + + let subagent_src = r#" + graph g { + start r + node r { + kind subagent + agent "researcher" + model "totally-unregistered" + next END + } + } + "#; + let bp = compile(&parse_str(subagent_src).unwrap()) + .unwrap() + .remove(0); + let resolver = CapabilityResolver::new() + .allow_agent("researcher") + .with_node_kinds(node_kinds()); + let err = resolver.bind_blueprint(&bp).unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown model"), "{err}"); + assert!(err.to_string().contains("totally-unregistered"), "{err}"); + + let repl_agent_src = r#" + graph g { + start r + node r { + kind repl_agent + script "triage" + model "totally-unregistered" + next END + } + } + "#; + let bp = compile(&parse_str(repl_agent_src).unwrap()) + .unwrap() + .remove(0); + let resolver = CapabilityResolver::new() + .allow_script("triage") + .with_node_kinds(node_kinds()); + let err = resolver.bind_blueprint(&bp).unwrap_err(); + assert!(err.to_string().contains("unknown model"), "{err}"); + assert!(err.to_string().contains("totally-unregistered"), "{err}"); + + // A registered model passes, alongside the registered primary reference. + let ok_resolver = CapabilityResolver::new() + .allow_agent("researcher") + .allow_model("totally-unregistered") + .with_node_kinds(node_kinds()); + let bp = compile(&parse_str(subagent_src).unwrap()) + .unwrap() + .remove(0); + ok_resolver.bind_blueprint(&bp).unwrap(); +} + +#[test] +fn steering_block_on_subagent_node_parses() { + // The reference doc's own worked example + // (docs/modules/expressive-language/reference.md, `subagent` section) must + // parse: it previously failed with "unknown node item `steering`" even + // though `steering` is documented as a supported `subagent` field and + // README.md's grammar declares `steering_decl = "steering" object`. + let src = r#" + graph g { + start research + node research { + kind subagent + agent "researcher" + steering { + parent allow ["add_instruction", "request_status", "cancel"] + human allow ["add_instruction", "pause", "resume", "cancel"] + delivery "safe_boundary" + } + next synthesize + } + node synthesize { + kind model + next END + } + } + "#; + let program = parse_str(src).unwrap(); + let node = &program.graphs[0].nodes[0]; + assert_eq!(node.name, "research"); + let steering = node.steering.as_ref().expect("steering block parsed"); + assert_eq!( + steering.parent_allow, + vec!["add_instruction", "request_status", "cancel"] + ); + assert_eq!( + steering.human_allow, + vec!["add_instruction", "pause", "resume", "cancel"] + ); + assert_eq!(steering.delivery.as_deref(), Some("safe_boundary")); + + // The declaration survives compilation (lowering into the blueprint stays + // a no-op; the point is that documented source is no longer rejected). + compile(&program).unwrap(); +} + +#[test] +fn steering_block_rejects_unknown_item() { + let src = + "graph g { start a node a { kind subagent agent \"x\" steering { bogus 1 } next END } }"; + let err = parse_str(src).unwrap_err(); + match err { + crate::error::TinyAgentsError::Parse { message, .. } => { + assert!( + message.contains("`parent`, `human`, or `delivery`"), + "{message}" + ); + } + other => panic!("expected parse error, got {other:?}"), + } +} diff --git a/src/language/test/parser.rs b/src/language/test/parser.rs index c3cd084..b270b25 100644 --- a/src/language/test/parser.rs +++ b/src/language/test/parser.rs @@ -78,3 +78,23 @@ fn parse_rejects_unknown_node_item() { other => panic!("expected parse error, got {other:?}"), } } + +#[test] +fn parse_rejects_token_stream_missing_eof_sentinel_instead_of_hanging() { + // The lexer always terminates a stream with `Eof`; the cursor helpers in + // `Parser` rely on that sentinel to guarantee forward progress. A caller + // that slices off the trailing `Eof` (e.g. after filtering/truncating a + // token stream) must get a parse error, not an infinite loop. + let tokens = tokenize("graph g { start").unwrap(); + assert!(matches!(tokens.last().unwrap().token, Token::Eof)); + let truncated = &tokens[..tokens.len() - 1]; + assert!(!matches!(truncated.last().unwrap().token, Token::Eof)); + + let err = parse(truncated).unwrap_err(); + match err { + crate::error::TinyAgentsError::Parse { message, .. } => { + assert!(message.contains("end-of-input sentinel"), "{message}"); + } + other => panic!("expected parse error, got {other:?}"), + } +} diff --git a/src/language/test/resolver.rs b/src/language/test/resolver.rs index a2aaf7a..a2d201f 100644 --- a/src/language/test/resolver.rs +++ b/src/language/test/resolver.rs @@ -115,6 +115,36 @@ fn resolver_collects_multiple_diagnostics() { assert!(codes.contains(&"E-rag-unknown-reducer"), "{codes:?}"); } +#[test] +fn resolver_rejects_unregistered_secondary_model_on_subagent() { + // `subagent` nodes may carry a `model` field alongside `agent`; the + // registry-backed `Resolver` must validate it too, mirroring the strict + // `bind_blueprint` gate. + let reg = full_registry(); + let caps = reg.capability_resolver().allow_agent("researcher"); + + let src = r#"graph g { start r node r { kind subagent agent "researcher" model "totally-unregistered" next END } }"#; + let program = parse_str(src).unwrap(); + let diagnostics = Resolver::from_capabilities(caps.clone()).resolve_program(&program); + assert_eq!(diagnostics.len(), 1, "{diagnostics:?}"); + assert_eq!(diagnostics[0].code.as_deref(), Some("E-rag-unknown-model")); + assert!( + diagnostics[0] + .message + .contains("unknown model `totally-unregistered`"), + "{:?}", + diagnostics[0] + ); + + // The span-less blueprint path shares the same gap-closing gate. + let bp = compile(&parse_str(src).unwrap()).unwrap().remove(0); + let err = Resolver::from_capabilities(caps) + .resolve_blueprint(&bp) + .unwrap_err(); + assert!(matches!(err, crate::error::TinyAgentsError::Capability(_))); + assert!(err.to_string().contains("unknown model"), "{err}"); +} + #[test] fn resolver_blueprint_path_matches_registry_binding() { // The span-less blueprint path mirrors the legacy gate's variants/messages. diff --git a/src/repl/session/builtins/batched.rs b/src/repl/session/builtins/batched.rs index 813a23e..4566993 100644 --- a/src/repl/session/builtins/batched.rs +++ b/src/repl/session/builtins/batched.rs @@ -41,7 +41,7 @@ pub(super) fn model_query_batched_impl( .registry .model(&model_name) .ok_or_else(|| raise(ctx, TinyAgentsError::ModelNotFound(model_name.clone())))?; - let request = build_model_request(&model_name, params); + let request = build_model_request(params); let structured = map_bool(params, "structured").unwrap_or(false); // Stream a "started" event for every fan-out leg up front, so a live // observer sees the whole batch dispatch before any leg completes. diff --git a/src/repl/session/builtins/capabilities.rs b/src/repl/session/builtins/capabilities.rs index 01f0f47..3a84d0d 100644 --- a/src/repl/session/builtins/capabilities.rs +++ b/src/repl/session/builtins/capabilities.rs @@ -19,7 +19,7 @@ pub(super) fn model_query_impl( .registry .model(&model_name) .ok_or_else(|| raise(ctx, TinyAgentsError::ModelNotFound(model_name.clone())))?; - let request = build_model_request(&model_name, params); + let request = build_model_request(params); let call_id = new_call_id(); emit_call_started(ctx, &call_id, ReplCallKind::Model, &model_name); let start = Instant::now(); diff --git a/src/repl/session/builtins/mod.rs b/src/repl/session/builtins/mod.rs index ad33f69..fe44a28 100644 --- a/src/repl/session/builtins/mod.rs +++ b/src/repl/session/builtins/mod.rs @@ -235,11 +235,41 @@ type AgentBatchItem = (String, String, Duration); // ── Error / recording helpers ─────────────────────────────────────────────── +/// Returns whether a capability error must abort the cell (a policy bound +/// tripped) instead of surfacing inside the script as an ordinary, catchable +/// runtime error. Mirrors [`crate::rlm::host::is_fatal`] — kept as a separate +/// copy here since the `repl` and `rlm` cargo features are independent, so +/// this module cannot assume the `rlm` module is compiled in. +fn is_fatal(err: &TinyAgentsError) -> bool { + matches!( + err, + TinyAgentsError::LimitExceeded(_) + | TinyAgentsError::Timeout(_) + | TinyAgentsError::Cancelled + | TinyAgentsError::SubAgentDepth(_) + ) +} + /// Stashes the precise crate error so `eval_cell` can surface it verbatim, and /// returns the stringly-typed Rhai runtime error the engine propagates. +/// +/// Only *fatal* errors (see [`is_fatal`] — a policy bound such as a call +/// limit, timeout, cancellation, or recursion depth) are stashed as the +/// cell-aborting `host_error`: `on_progress` polls that flag and terminates +/// the script at the next statement, and `eval_cell`'s success path prefers +/// it even when the script otherwise completed normally. A *recoverable* +/// capability failure (unknown tool/model/agent, a tool-reported error, …) +/// must remain an ordinary catchable Rhai runtime error so `try`/`catch` in +/// the script actually works — it is stashed only into the non-aborting +/// `last_capability_error` slot, which `eval_cell`'s error path consults to +/// recover the typed error for an error the script left uncaught. fn raise(ctx: &HostContext, err: TinyAgentsError) -> Box { let message = err.to_string(); - ctx.buffers.set_host_error(err); + if is_fatal(&err) { + ctx.buffers.set_host_error(err); + } else { + ctx.buffers.set_last_capability_error(err); + } Box::new(EvalAltResult::ErrorRuntime( Dynamic::from(message), Position::NONE, @@ -418,7 +448,17 @@ fn check_depth(ctx: &HostContext) -> Result<(), Box ModelRequest { +/// +/// `model` here is the *registry* alias the script named (`map_str(params, +/// "model")`), not a provider model id — `CapabilityRegistry::register_model` +/// allows them to differ. Leave `ModelRequest::model` unset: the resolved +/// `ChatModel` already carries its own provider configuration, and a provider +/// transport that reads `request.model` (falling back to its own model only +/// when unset) would otherwise send the registry alias itself as the model id +/// on the wire. Mirrors `RlmHost::handle_llm` / `RlmRunner::run` in +/// `src/rlm/`, which build `ModelRequest { messages, ..Default::default() }` +/// for exactly this reason. +fn build_model_request(params: &Map) -> ModelRequest { let mut messages = Vec::new(); if let Some(system) = map_str(params, "system") { messages.push(Message::system(system)); @@ -428,7 +468,6 @@ fn build_model_request(model: &str, params: &Map) -> ModelRequest { } ModelRequest { messages, - model: Some(model.to_string()), ..Default::default() } } diff --git a/src/repl/session/mod.rs b/src/repl/session/mod.rs index ec0715a..e7ec3ff 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -71,6 +71,15 @@ pub(super) struct CellBuffers { calls: Arc>>, answer: Arc>>, host_error: Arc>>, + /// The most recent *recoverable* capability error raised this cell (see + /// `builtins::raise`/`builtins::is_fatal`), regardless of whether the + /// script caught it. Unlike `host_error`, this is never consulted by + /// `on_progress` or the success path, so a `try`/`catch`ed error has no + /// further effect once the script continues normally — only + /// [`ReplSession::eval_cell`]'s error path reads it, to recover the typed + /// error for a capability failure the script left uncaught instead of + /// falling back to a stringly-wrapped [`TinyAgentsError::Validation`]. + last_capability_error: Arc>>, vars_snapshot: Arc>>, /// The wall-clock instant the current cell's [`ReplPolicy::timeout`] /// expires at, if the policy configures one. Set at the start of @@ -530,7 +539,20 @@ impl ReplSession { if let Some(host_err) = self.buffers.take_host_error() { return Err(host_err); } - return Err(map_rhai_error(*err)); + // The script left a *recoverable* capability error uncaught. + // `raise` stashed its typed form (without aborting the cell); + // recover it here rather than reporting the generic, + // stringly-wrapped Rhai runtime error — but only when the + // propagated error is actually that same failure (an + // unrelated later error must not be misreported as the + // earlier, already-handled one). + let mapped = map_rhai_error(*err); + if let Some(last) = self.buffers.take_last_capability_error() + && matches!(&mapped, TinyAgentsError::Validation(msg) if msg.contains(&last.to_string())) + { + return Err(last); + } + return Err(mapped); } }; @@ -583,6 +605,10 @@ impl CellBuffers { self.calls.lock().expect("calls poisoned").clear(); *self.answer.lock().expect("answer poisoned") = None; *self.host_error.lock().expect("host_error poisoned") = None; + *self + .last_capability_error + .lock() + .expect("last_capability_error poisoned") = None; *self.deadline.lock().expect("deadline poisoned") = None; *self .max_output_bytes @@ -627,6 +653,14 @@ impl CellBuffers { self.host_error.lock().expect("host_error poisoned").take() } + /// Takes the most recently stashed recoverable capability error, if any. + fn take_last_capability_error(&self) -> Option { + self.last_capability_error + .lock() + .expect("last_capability_error poisoned") + .take() + } + // ── Accessors used by the capability built-ins (in `builtins.rs`). ── /// Pushes a recorded capability call/event. @@ -683,6 +717,17 @@ impl CellBuffers { *self.host_error.lock().expect("host_error poisoned") = Some(err); } + /// Stashes a *recoverable* capability error (see `builtins::is_fatal`) + /// without aborting the cell, so an uncaught occurrence can still be + /// reported with its precise type by [`ReplSession::eval_cell`]'s error + /// path. + pub(super) fn set_last_capability_error(&self, err: TinyAgentsError) { + *self + .last_capability_error + .lock() + .expect("last_capability_error poisoned") = Some(err); + } + /// Returns the pre-cell namespace snapshot for `show_vars()`. pub(super) fn vars_snapshot(&self) -> BTreeMap { self.vars_snapshot diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index a047603..2baa46b 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -415,6 +415,41 @@ fn tool_call_batched_keeps_successes_when_one_item_tool_errors() { assert_eq!(items[2]["content"], serde_json::json!("ok:3")); } +#[test] +fn a_recoverable_capability_error_is_catchable_by_try_catch() { + // Regression test: `raise()` used to stash *every* capability error + // (recoverable or not) into `host_error`, which `eval_cell`'s success + // path and `on_progress` both treat as fatal — so a script that caught + // the error and recovered still failed the whole cell. An unknown tool + // name is a recoverable failure (`TinyAgentsError::ToolNotFound`), not a + // policy bound, so `try`/`catch` around it must actually work. + let mut s = session(); + + let result = s + .eval_cell(r#"let ok = 0; try { tool_call(#{ tool: "nope" }); } catch(e) { ok = 1; } ok"#) + .expect("a caught recoverable capability error must not fail the cell"); + + assert_eq!(result.value, Some(ReplValue::Int(1))); +} + +#[test] +fn an_uncaught_recoverable_capability_error_still_reports_its_typed_form() { + // The typed-error contract for an *uncaught* recoverable failure must + // survive the fix above: `eval_cell` should still report + // `TinyAgentsError::ToolNotFound`, not a generic stringly-wrapped + // `Validation` error. + let mut s = session(); + + let err = s + .eval_cell(r#"tool_call(#{ tool: "nope" })"#) + .expect_err("an unregistered tool must fail the cell when uncaught"); + + assert!( + matches!(err, TinyAgentsError::ToolNotFound(ref t) if t == "nope"), + "expected ToolNotFound(nope), got {err:?}" + ); +} + /// A trivial [`HarnessAgent`] that returns a fixed response, for exercising /// `agent_query` without a real model/harness run. struct StubAgent; diff --git a/src/rlm/interpreter/rhai_cell.rs b/src/rlm/interpreter/rhai_cell.rs index 6dd2eeb..5db4478 100644 --- a/src/rlm/interpreter/rhai_cell.rs +++ b/src/rlm/interpreter/rhai_cell.rs @@ -34,7 +34,23 @@ const CANCELLED_TOKEN: &str = "rlm cell cancelled by host"; /// The embedded Rhai backend. See the [module docs](self). pub struct RhaiInterpreter { max_operations: u64, - scope: Scope<'static>, + /// The persistent notebook scope, behind a shared handle rather than + /// owned outright by `Self`. + /// + /// `eval_cell` hands the scope to a `spawn_blocking` task. If that were a + /// bare `Scope` moved in and out via `mem::take`, dropping the + /// `eval_cell` future (a caller-side `timeout`/`select!`/abort — the + /// documented cancellation shape) would detach the blocking task and + /// silently discard the scope it was about to write back, leaving `self` + /// with the empty scope `mem::take` left behind and no indication + /// anything was lost. Keeping the scope behind `Arc>` instead + /// means a dropped future no longer owns the only copy: the orphaned + /// task still writes into the shared scope, and the mutex serializes any + /// next cell behind it rather than starting from empty. A poisoned lock + /// (the blocking closure panicked mid-eval) is treated as an + /// unrecoverable session error instead of silently falling back to an + /// empty namespace. + scope: Arc>>, } impl RhaiInterpreter { @@ -43,7 +59,7 @@ impl RhaiInterpreter { pub fn new(max_operations: u64) -> Self { Self { max_operations, - scope: Scope::new(), + scope: Arc::new(Mutex::new(Scope::new())), } } } @@ -331,6 +347,12 @@ Rhai syntax notes (Rhai is NOT JavaScript or Rust): async fn set_variable(&mut self, name: &str, value: Value) -> Result<()> { self.scope + .lock() + .map_err(|_| { + TinyAgentsError::Model( + "rlm rhai interpreter scope poisoned by a previous panic".to_string(), + ) + })? .set_value(name.to_string(), json_to_dynamic(&value)); Ok(()) } @@ -338,19 +360,38 @@ Rhai syntax notes (Rhai is NOT JavaScript or Rust): async fn eval_cell(&mut self, code: &str, host: Arc) -> Result { let cell: SharedCellState = Arc::new(Mutex::new(CellState::default())); let engine = build_engine(host, cell.clone(), self.max_operations); - let mut scope = std::mem::take(&mut self.scope); + let scope = self.scope.clone(); let code = code.to_string(); // Rhai is synchronous and the capability closures block through the // bridge, so evaluate on the blocking pool to keep the async runtime // (which drives the actual provider I/O) responsive. - let (scope_back, eval) = tokio::task::spawn_blocking(move || { - let result = engine.eval_with_scope::(&mut scope, &code); - (scope, result) + // + // The scope is locked *inside* the blocking closure (not moved out of + // `self` via `mem::take`) so a caller that drops this `eval_cell` + // future — a `timeout`/`select!`/abort around `RlmSession::eval` or + // `RlmRunner::run` — never leaves `self.scope` empty: the orphaned + // blocking task still holds the only route back to the scope and + // still writes its updates into it before the lock releases. + let eval = tokio::task::spawn_blocking(move || { + // A poisoned lock means an earlier cell's blocking task panicked + // while holding the scope: surface that as a distinct outcome + // rather than silently continuing on a scope whose consistency + // is no longer guaranteed. + match scope.lock() { + Ok(mut guard) => Ok(engine.eval_with_scope::(&mut guard, &code)), + Err(_poisoned) => Err(()), + } }) .await - .map_err(|err| TinyAgentsError::Model(format!("rlm rhai eval task failed: {err}")))?; - self.scope = scope_back; + .map_err(|err| TinyAgentsError::Model(format!("rlm rhai eval task failed: {err}")))? + .map_err(|()| { + TinyAgentsError::Model( + "rlm rhai interpreter state lost: a previous cell panicked while holding the \ + notebook scope" + .to_string(), + ) + })?; let mut state = cell.lock().expect("cell state poisoned"); if let Some(fatal) = state.fatal.take() { diff --git a/src/rlm/runner.rs b/src/rlm/runner.rs index e862da3..a8616cd 100644 --- a/src/rlm/runner.rs +++ b/src/rlm/runner.rs @@ -154,7 +154,18 @@ impl RlmRunner { let mut nudged = false; let outcome = loop { - if steps.len() >= self.config.policy.max_cells { + // Gate on the *session-cumulative* cell count (matching + // `RlmSession::eval`'s own enforcement, and `docs/modules/rlm`'s + // documented "counters are session-cumulative" contract) rather + // than `steps.len()`, which resets to zero on every `run()` call. + // With `steps.len()` the two checks only agreed on the very + // first run: a second `run()` on the same (long-lived, + // `&mut self`) runner would see an empty `steps`, pay for a full + // driver-model call, and only then have `self.session.eval` + // return a hard `LimitExceeded` error instead of the graceful + // `CellBudgetExhausted` outcome the same condition produces on + // the first run. + if self.session.cells_run() >= self.config.policy.max_cells { break RlmOutcome { answer: None, stop_reason: RlmStopReason::CellBudgetExhausted, diff --git a/src/rlm/session.rs b/src/rlm/session.rs index ac3929c..36608b2 100644 --- a/src/rlm/session.rs +++ b/src/rlm/session.rs @@ -19,6 +19,19 @@ use crate::error::{Result, TinyAgentsError}; /// [`RlmPolicy::max_output_bytes`] and is truncated. const TRUNCATION_MARKER: &str = "\n… [output truncated by rlm policy]"; +/// Truncates `s` to at most `max` bytes, walking back to the nearest UTF-8 +/// char boundary at or below `max` first. `String::truncate` panics on a +/// non-boundary byte index, and captured cell output is arbitrary text (a +/// multi-byte character can straddle the raw budget), so the cut point must +/// be found before truncating. +fn truncate_at_char_boundary(s: &mut String, max: usize) { + let mut end = max; + while end > 0 && !s.is_char_boundary(end) { + end -= 1; + } + s.truncate(end); +} + /// One sandboxed script workspace: a persistent interpreter plus the /// capability host its cells call back into. pub struct RlmSession { @@ -108,16 +121,21 @@ impl RlmSession { let mut eval = eval?; // Bound what flows back into the driver conversation. Truncation is - // explicit (marked) so the model knows it saw a prefix. + // explicit (marked) so the model knows it saw a prefix. `truncate` + // cuts at a raw byte offset, so the budget is first walked back to + // the nearest UTF-8 char boundary — output is arbitrary + // model/script-authored text (CJK, emoji, accents included), and a + // multi-byte character straddling `max_output_bytes` would otherwise + // panic `String::truncate`. if eval.stdout.len() > policy.max_output_bytes { - eval.stdout.truncate(policy.max_output_bytes); + truncate_at_char_boundary(&mut eval.stdout, policy.max_output_bytes); eval.stdout.push_str(TRUNCATION_MARKER); } if let Some(value) = &eval.value { let rendered = value.to_string(); if rendered.len() > policy.max_output_bytes { let mut clipped = rendered; - clipped.truncate(policy.max_output_bytes); + truncate_at_char_boundary(&mut clipped, policy.max_output_bytes); clipped.push_str(TRUNCATION_MARKER); eval.value = Some(Value::String(clipped)); } diff --git a/src/rlm/test.rs b/src/rlm/test.rs index e84b3d9..27f2494 100644 --- a/src/rlm/test.rs +++ b/src/rlm/test.rs @@ -3,11 +3,12 @@ //! against deterministic capability doubles. use std::sync::Arc; +use std::time::Duration; use serde_json::json; use super::*; -use crate::harness::testkit::{FakeTool, ScriptedModel}; +use crate::harness::testkit::{FakeTool, ScriptedModel, SlowModel}; use crate::registry::CapabilityRegistry; fn registry_with_mock(replies: Vec<&str>) -> Arc> { @@ -238,6 +239,46 @@ async fn oversized_stdout_is_truncated_with_a_marker() { assert!(outcome.stdout.contains("truncated")); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn oversized_multibyte_stdout_is_truncated_without_panicking() { + // Regression test: `String::truncate` panics unless the cut index is a + // UTF-8 char boundary. `max_output_bytes: 64` is not a multiple of the + // 3-byte-wide "日" character printed below, so the naive raw-byte cut + // used to panic partway through evaluating the cell. + let policy = RlmPolicy { + max_output_bytes: 64, + ..RlmPolicy::default() + }; + let mut session = rhai_session(registry_with_mock(vec![]), policy); + let outcome = session + .eval(r#"for i in 0..100 { print("日日日日日日日日日日"); }"#) + .await + .expect("cell must not panic on a multi-byte truncation boundary"); + assert!(outcome.stdout.contains("truncated")); + // The truncated prefix must itself still be valid UTF-8 (no half-cut + // multi-byte character), which `String::truncate` guarantees once the + // cut lands on a char boundary. + assert!(std::str::from_utf8(outcome.stdout.as_bytes()).is_ok()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn oversized_multibyte_value_is_truncated_without_panicking() { + // Same boundary hazard as stdout, but for the rendered cell value. + let policy = RlmPolicy { + max_output_bytes: 64, + ..RlmPolicy::default() + }; + let mut session = rhai_session(registry_with_mock(vec![]), policy); + let outcome = session + .eval(r#"let s = ""; for i in 0..100 { s += "日"; } s"#) + .await + .expect("cell must not panic on a multi-byte truncation boundary"); + let value = outcome.value.expect("truncated value"); + let text = value.as_str().expect("string value"); + assert!(text.contains("truncated")); + assert!(std::str::from_utf8(text.as_bytes()).is_ok()); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn context_variable_is_visible_to_scripts() { let mut session = rhai_session(registry_with_mock(vec![]), RlmPolicy::default()); @@ -249,6 +290,88 @@ async fn context_variable_is_visible_to_scripts() { assert_eq!(outcome.value, Some(json!("beta"))); } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn dropping_a_cell_future_does_not_lose_the_persistent_scope() { + // Regression test: `eval_cell` used to `mem::take` the scope out of + // `self` and only restore it after its `spawn_blocking` task joined. A + // caller that drops the `eval_cell` future mid-flight — the documented + // `tokio::time::timeout`/`select!`/task-abort cancellation shape — left + // `self.scope` permanently empty, since the detached blocking task's + // `(scope, result)` was discarded along with the dropped future. Keeping + // the scope behind `Arc>` means the orphaned task still writes + // its updates into the *shared* scope, so a later cell still sees them. + let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); + registry + .register_model( + "mock", + Arc::new(SlowModel::new(Duration::from_millis(150), "done")), + ) + .expect("register model"); + let mut session = rhai_session(Arc::new(registry), RlmPolicy::default()); + + // `x` is assigned before the script blocks on the slow `llm(...)` call, + // so the assignment has already happened by the time this future is + // dropped. + let cancelled = tokio::time::timeout( + Duration::from_millis(20), + session.eval(r#"let x = 42; llm("wait"); x"#), + ) + .await; + assert!( + cancelled.is_err(), + "the timeout must fire before the slow llm() call resolves" + ); + + // Give the orphaned blocking task time to actually finish the call and + // write the scope back. + tokio::time::sleep(Duration::from_millis(300)).await; + + let outcome = session.eval("x").await.expect("second cell"); + assert_eq!(outcome.value, Some(json!(42))); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_panicked_cell_poisons_the_scope_instead_of_silently_emptying_it() { + // The other half of the scope-loss defect: a blocking closure that + // panics mid-eval (holding the scope lock) must not leave the + // interpreter usable-but-amnesiac. The mutex poisons, and every + // subsequent cell must fail loudly with a clear diagnostic rather than + // silently running against an empty namespace. + struct PanicModel; + + #[async_trait::async_trait] + impl crate::harness::model::ChatModel<()> for PanicModel { + async fn invoke( + &self, + _state: &(), + _request: crate::harness::model::ModelRequest, + ) -> crate::error::Result { + panic!("simulated provider panic while holding the rlm scope"); + } + } + + let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); + registry + .register_model("mock", Arc::new(PanicModel)) + .expect("register model"); + let mut session = rhai_session(Arc::new(registry), RlmPolicy::default()); + + let first = session.eval(r#"let x = 1; llm("boom")"#).await; + assert!( + first.is_err(), + "the panicking cell must fail, not silently succeed" + ); + + let second = session + .eval("x") + .await + .expect_err("the scope is poisoned, so the next cell must fail loudly"); + assert!( + matches!(&second, crate::error::TinyAgentsError::Model(msg) if msg.contains("interpreter state lost")), + "got {second:?}" + ); +} + // ── The model-driven runner ───────────────────────────────────────────────── #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -329,6 +452,45 @@ async fn runner_stops_at_the_cell_budget() { assert_eq!(outcome.steps.len(), 2); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_second_run_after_the_cell_budget_stops_gracefully_instead_of_erroring() { + // Regression test: `RlmRunner::run` used to gate its loop on + // `steps.len()`, a per-call counter that resets to zero on every `run()` + // call, while `RlmSession::eval` enforces `max_cells` against its own + // session-cumulative `cells_run` counter that nothing ever reset. The two + // checks agreed only on the first call: a second `run()` on the same + // (long-lived, `&mut self`) runner — legal, and a natural thing to do — + // saw an empty `steps`, paid for a driver-model call it didn't need, and + // then hit `self.session.eval`'s hard `LimitExceeded` error instead of + // the graceful `CellBudgetExhausted` outcome the identical condition + // produces on the first run. + let cells: Vec<&str> = vec!["```rhai\n1\n```"; 4]; + let registry = registry_with_mock(cells); + let config = RlmConfig { + driver_model: Some("mock".to_string()), + policy: RlmPolicy { + max_cells: 2, + ..RlmPolicy::default() + }, + ..RlmConfig::default() + }; + let mut runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("build runner"); + + let first = runner.run("loop forever").await.expect("first run"); + assert_eq!(first.stop_reason, RlmStopReason::CellBudgetExhausted); + assert_eq!(first.steps.len(), 2); + + let second = runner + .run("try again") + .await + .expect("a second run must stop gracefully, not return a hard error"); + assert_eq!(second.stop_reason, RlmStopReason::CellBudgetExhausted); + // No cells executed (the budget was already spent) and no driver call + // wasted producing one that would only be rejected. + assert_eq!(second.steps.len(), 0); + assert_eq!(second.driver_calls, 0); +} + // ── Cancellation ──────────────────────────────────────────────────────────── #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/src/rlm/types.rs b/src/rlm/types.rs index 7fc2303..7fcada7 100644 --- a/src/rlm/types.rs +++ b/src/rlm/types.rs @@ -95,7 +95,9 @@ impl InterpreterSpec { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(default)] pub struct RlmPolicy { - /// Maximum code cells one [`super::RlmRunner::run`] loop may execute. + /// Maximum code cells a session may execute, cumulative across every + /// [`super::RlmRunner::run`] call on the same runner/session (matching + /// [`super::RlmSession::eval`]'s own enforcement) — not reset per call. pub max_cells: usize, /// Maximum source size, in bytes, of a single cell. pub max_script_bytes: usize, diff --git a/tests/feature_repl_session.rs b/tests/feature_repl_session.rs index 29066a7..8106983 100644 --- a/tests/feature_repl_session.rs +++ b/tests/feature_repl_session.rs @@ -58,6 +58,34 @@ fn model_query_calls_a_registered_model_and_records_the_call() { assert_eq!(result.calls[0].name, "assistant"); } +#[test] +fn model_query_does_not_leak_the_registry_alias_as_the_provider_model_id() { + // Regression test: `build_model_request` used to set `ModelRequest.model` + // to the *registry* name the script passed, not a provider model id. + // `CapabilityRegistry::register_model` allows those to differ (a host + // may register a `"fast"` alias for `gpt-4o-mini`), and a real provider + // transport sends `request.model` verbatim on the wire when set — so the + // request must leave `model` unset and let the resolved model supply its + // own provider id, exactly like `RlmHost::handle_llm` / `RlmRunner::run`. + let scripted = Arc::new(ScriptedModel::replies(vec!["hi"])); + let mut registry = CapabilityRegistry::<()>::new(); + registry + .register_model("fast", scripted.clone()) + .expect("register model"); + let mut s = + ReplSession::<()>::new().with_capabilities(ReplCapabilities::new(Arc::new(registry))); + + s.eval_cell(r#"model_query(#{ model: "fast", prompt: "hi" })"#) + .expect("model_query"); + + let requests = scripted.requests(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].model, None, + "the registry alias `fast` must not be sent as the provider model id" + ); +} + #[test] fn model_query_structured_returns_content_and_finish_reason() { let mut s = session_with_model(vec!["structured-reply"]);