From 622c9b27eb4e6b99d6e5c7d9c7f622cf295b5b72 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 20:01:12 +0000 Subject: [PATCH 1/2] Add REPL host-embedding: external cancellation + live capability events Fills the three host-embedding gaps a host (openhuman) needs to drive `repl::session::ReplSession` from an async runtime, none of which could be worked around downstream: - External cancellation. New `ReplCancelFlag` (sticky `Arc`), installed via `ReplSession::with_cancel_flag`. Enforced fail-closed at the same two points as the deadline: the engine `on_progress` hook terminates a running script at the next statement (mapped to `Cancelled` via a `CANCELLED_TOKEN` sentinel), and the blocking bridge grows a watcher thread that drops an in-flight (possibly hung) capability future promptly on cancel. `eval_cell` also fails closed before starting a cell if already cancelled, without consuming an iteration. Reuses the existing unit `TinyAgentsError::Cancelled` (the plan's proposed `Cancelled(String)` would have been a breaking change to ~15 existing call sites for no added value). - Live capability-call events. New feature-gated `AgentEvent::ReplCall { session_id, record, phase }` (+ `ReplCallPhase`) streamed on the session `EventSink` at capability-call start and completion, correlated by a `call_id` threaded from the built-in through `record`. A host subscribes an `EventListener` and forwards REPL progress mid-cell, instead of only reading `ReplResult.calls` after the cell returns. - Async-embedding docs. `eval_cell` rustdoc now states it blocks internally (`futures::executor::block_on`) and must be driven via `spawn_blocking`/dedicated thread, and deadlocks on a current-thread runtime. Tests ship with the change (crate convention): cancel-before-eval, cancel-mid-script-loop (`on_progress` path), cancel-during-hanging-capability (bridge path), and start+completion events paired by call_id. `cargo fmt --check`, `cargo clippy --all-targets --features repl -- -D warnings`, `cargo test --features repl`, and plain `cargo test` all pass. --- src/harness/events/types.rs | 42 ++++ src/lib.rs | 4 +- src/repl/session/builtins/authoring.rs | 3 + src/repl/session/builtins/batched.rs | 59 ++++-- src/repl/session/builtins/capabilities.rs | 36 +++- src/repl/session/builtins/mod.rs | 242 ++++++++++++++++++---- src/repl/session/mod.rs | 76 +++++++ src/repl/session/test.rs | 170 +++++++++++++++ src/repl/session/types.rs | 44 ++++ 9 files changed, 609 insertions(+), 67 deletions(-) diff --git a/src/harness/events/types.rs b/src/harness/events/types.rs index be79030..03a1822 100644 --- a/src/harness/events/types.rs +++ b/src/harness/events/types.rs @@ -448,6 +448,46 @@ pub enum AgentEvent { /// Human-readable error description. error: String, }, + + /// A `.ragsh` REPL cell dispatched (or completed) a host capability call + /// (`model_query`, `tool_call`, `agent_query`, or an `emit`). + /// + /// [`ReplResult::calls`](crate::repl::session::ReplResult) is only readable + /// *after* a cell returns, so a long fan-out would otherwise look frozen to + /// a live observer. This event streams each capability call as it starts and + /// again as it completes, letting a host (which subscribes an + /// [`EventListener`] on the session [`EventSink`]) forward REPL progress to + /// its own UI/progress sink mid-cell. Gated behind the `repl` feature so the + /// default build neither pulls in the Rhai engine nor references + /// [`ReplCallRecord`](crate::repl::session::ReplCallRecord). + #[cfg(feature = "repl")] + ReplCall { + /// The session label (its [`SessionId`](crate::harness::ids::SessionId) + /// string) the call was issued from, correlating the event back to the + /// REPL session that produced it. + session_id: String, + /// The call record. On the [`ReplCallPhase::Started`] phase the record's + /// `elapsed` is zero and `detail` is minimal (the completed phase + /// carries the full detail and measured `elapsed`); `call_id` is stable + /// across the two phases so a listener can pair them. + record: crate::repl::session::ReplCallRecord, + /// Whether the call is starting or has completed. + phase: ReplCallPhase, + }, +} + +/// The lifecycle phase of an [`AgentEvent::ReplCall`] — whether a REPL host +/// capability call is just starting or has completed. +/// +/// Gated behind the `repl` feature alongside the event it annotates. +#[cfg(feature = "repl")] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ReplCallPhase { + /// The capability call has been dispatched but has not yet returned. + Started, + /// The capability call returned (successfully or with a recorded error). + Completed, } /// Names the kind of run limit that tripped in an [`AgentEvent::LimitReached`]. @@ -521,6 +561,8 @@ impl AgentEvent { AgentEvent::StreamClosed => "stream.closed", AgentEvent::RunCompleted { .. } => "run.completed", AgentEvent::RunFailed { .. } => "run.failed", + #[cfg(feature = "repl")] + AgentEvent::ReplCall { .. } => "repl.call", } } } diff --git a/src/lib.rs b/src/lib.rs index a0456b1..6efec5e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -242,6 +242,6 @@ pub use graph::testkit::{ // remains available as `repl::ReplSession`. #[cfg(feature = "repl")] pub use repl::session::{ - LanguageCompiler, ReplCallKind, ReplCallRecord, ReplCapabilities, ReplPolicy, ReplResult, - ReplSession, ReplValue, ReplVariables, + LanguageCompiler, ReplCallKind, ReplCallRecord, ReplCancelFlag, ReplCapabilities, ReplPolicy, + ReplResult, ReplSession, ReplValue, ReplVariables, }; diff --git a/src/repl/session/builtins/authoring.rs b/src/repl/session/builtins/authoring.rs index f9ef98f..e49970e 100644 --- a/src/repl/session/builtins/authoring.rs +++ b/src/repl/session/builtins/authoring.rs @@ -92,6 +92,7 @@ pub(super) fn graph_define_impl( .insert(handle.name.clone(), handle.clone()); record( ctx, + new_call_id(), ReplCallKind::Graph, "graph_define", json!({ "name": handle.name }), @@ -166,6 +167,7 @@ pub(super) fn graph_compile_impl( .insert(handle.name.clone(), handle.clone()); record( ctx, + new_call_id(), ReplCallKind::Graph, "graph_compile", json!({ "name": handle.name, "requires_review": handle.requires_review }), @@ -222,6 +224,7 @@ pub(super) fn graph_register_impl( // the REPL never installs generated topology directly. record( ctx, + new_call_id(), ReplCallKind::Graph, "graph_register", json!({ "name": handle.name, "review_id": review_id }), diff --git a/src/repl/session/builtins/batched.rs b/src/repl/session/builtins/batched.rs index f47c176..636d7b7 100644 --- a/src/repl/session/builtins/batched.rs +++ b/src/repl/session/builtins/batched.rs @@ -32,6 +32,7 @@ pub(super) fn model_query_batched_impl( let items = batch_items(ctx, items, "model_query_batched")?; let mut prepared = Vec::with_capacity(items.len()); + let mut call_ids = Vec::with_capacity(items.len()); for params in &items { let model_name = map_str(params, "model") .ok_or_else(|| invalid(ctx, "model_query_batched: missing `model`"))?; @@ -42,12 +43,17 @@ pub(super) fn model_query_batched_impl( .ok_or_else(|| raise(ctx, TinyAgentsError::ModelNotFound(model_name.clone())))?; let request = build_model_request(&model_name, 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. + let call_id = new_call_id(); + emit_call_started(ctx, &call_id, ReplCallKind::Model, &model_name); + call_ids.push(call_id); prepared.push((model_name, model, request, structured)); } let concurrency = ctx.policy.max_concurrency.max(1); let results: Vec> = - bridge_block_on_raw(ctx.buffers.deadline(), async { + bridge_block_on_raw(ctx.buffers.deadline(), &ctx.cancel, async { stream::iter(prepared.iter().map(|(name, model, request, structured)| { let name = name.clone(); let structured = *structured; @@ -65,12 +71,14 @@ pub(super) fn model_query_batched_impl( }) .map_err(|err| raise(ctx, err))?; + // `buffered` preserves input order, so results align 1:1 with `call_ids`. let mut out = Array::with_capacity(results.len()); - for result in results { + for (call_id, result) in call_ids.into_iter().zip(results) { let (name, text, finish_reason, structured, elapsed) = result.map_err(|err| raise(ctx, err))?; record( ctx, + call_id, ReplCallKind::Model, &name, json!({ "chars": text.len() }), @@ -89,6 +97,7 @@ pub(super) fn tool_call_batched_impl( let items = batch_items(ctx, items, "tool_call_batched")?; let mut prepared = Vec::with_capacity(items.len()); + let mut call_ids = Vec::with_capacity(items.len()); for params in &items { let tool_name = map_str(params, "tool") .ok_or_else(|| invalid(ctx, "tool_call_batched: missing `tool`"))?; @@ -98,26 +107,31 @@ pub(super) fn tool_call_batched_impl( .tool(&tool_name) .ok_or_else(|| raise(ctx, TinyAgentsError::ToolNotFound(tool_name.clone())))?; let arguments = map_json(params, "arguments").unwrap_or(Value::Null); + let call_id = new_call_id(); + emit_call_started(ctx, &call_id, ReplCallKind::Tool, &tool_name); + call_ids.push(call_id); prepared.push((tool_name, tool, arguments)); } let concurrency = ctx.policy.max_concurrency.max(1); let results: Vec< Result<(String, crate::harness::tool::ToolResult, Duration), TinyAgentsError>, - > = bridge_block_on_raw(ctx.buffers.deadline(), async { - stream::iter(prepared.iter().map(|(name, tool, arguments)| { - let name = name.clone(); - let call = ToolCall { - id: new_call_id().as_str().to_string(), - name: name.clone(), - arguments: arguments.clone(), - }; - async move { - let start = Instant::now(); - let result = tool.call(&ctx.state, call).await?; - Ok((name, result, start.elapsed())) - } - })) + > = bridge_block_on_raw(ctx.buffers.deadline(), &ctx.cancel, async { + stream::iter(prepared.iter().zip(call_ids.iter()).map( + |((name, tool, arguments), call_id)| { + let name = name.clone(); + let call = ToolCall { + id: call_id.as_str().to_string(), + name: name.clone(), + arguments: arguments.clone(), + }; + async move { + let start = Instant::now(); + let result = tool.call(&ctx.state, call).await?; + Ok((name, result, start.elapsed())) + } + }, + )) .buffered(concurrency) .collect() .await @@ -132,10 +146,11 @@ pub(super) fn tool_call_batched_impl( // above (a harness-level failure, not a tool-reported one) still aborts // the whole batch, since no results exist to preserve in that case. let mut out = Array::with_capacity(results.len()); - for result in results { + for (call_id, result) in call_ids.into_iter().zip(results) { let (name, tool_result, elapsed) = result.map_err(|err| raise(ctx, err))?; record( ctx, + call_id, ReplCallKind::Tool, &name, json!({ "chars": tool_result.content.len() }), @@ -168,6 +183,7 @@ pub(super) fn agent_query_batched_impl( let items = batch_items(ctx, items, "agent_query_batched")?; let mut prepared = Vec::with_capacity(items.len()); + let mut call_ids = Vec::with_capacity(items.len()); for params in &items { let agent_name = map_str(params, "agent") .ok_or_else(|| invalid(ctx, "agent_query_batched: missing `agent`"))?; @@ -186,12 +202,15 @@ pub(super) fn agent_query_batched_impl( if let Some(data) = map_json(params, "input") { input = input.with_data(data); } + let call_id = new_call_id(); + emit_call_started(ctx, &call_id, ReplCallKind::Agent, &agent_name); + call_ids.push(call_id); prepared.push((agent_name, agent, input)); } let concurrency = ctx.policy.max_concurrency.max(1); let results: Vec> = - bridge_block_on_raw(ctx.buffers.deadline(), async { + bridge_block_on_raw(ctx.buffers.deadline(), &ctx.cancel, async { stream::iter(prepared.iter().map(|(name, agent, input)| { let name = name.clone(); async move { @@ -207,9 +226,9 @@ pub(super) fn agent_query_batched_impl( .map_err(|err| raise(ctx, err))?; let mut out = Array::with_capacity(results.len()); - for result in results { + for (call_id, result) in call_ids.into_iter().zip(results) { let (name, text, elapsed) = result.map_err(|err| raise(ctx, err))?; - record(ctx, ReplCallKind::Agent, &name, json!({}), elapsed); + record(ctx, call_id, ReplCallKind::Agent, &name, json!({}), elapsed); out.push(Dynamic::from(text)); } Ok(Dynamic::from_array(out)) diff --git a/src/repl/session/builtins/capabilities.rs b/src/repl/session/builtins/capabilities.rs index 64a9557..48264dc 100644 --- a/src/repl/session/builtins/capabilities.rs +++ b/src/repl/session/builtins/capabilities.rs @@ -20,14 +20,21 @@ pub(super) fn model_query_impl( .model(&model_name) .ok_or_else(|| raise(ctx, TinyAgentsError::ModelNotFound(model_name.clone())))?; let request = build_model_request(&model_name, params); + let call_id = new_call_id(); + emit_call_started(ctx, &call_id, ReplCallKind::Model, &model_name); let start = Instant::now(); - let response = bridge_block_on(ctx.buffers.deadline(), model.invoke(&ctx.state, request)) - .map_err(|err| raise(ctx, err))?; + let response = bridge_block_on( + ctx.buffers.deadline(), + &ctx.cancel, + model.invoke(&ctx.state, request), + ) + .map_err(|err| raise(ctx, err))?; let elapsed = start.elapsed(); let finish_reason = response.finish_reason.clone(); let text = Message::Assistant(response.message).text(); record( ctx, + call_id, ReplCallKind::Model, &model_name, json!({ "chars": text.len() }), @@ -52,17 +59,24 @@ pub(super) fn tool_call_impl( .tool(&tool_name) .ok_or_else(|| raise(ctx, TinyAgentsError::ToolNotFound(tool_name.clone())))?; let arguments = map_json(params, "arguments").unwrap_or(Value::Null); + let call_id = new_call_id(); let call = ToolCall { - id: new_call_id().as_str().to_string(), + id: call_id.as_str().to_string(), name: tool_name.clone(), arguments: arguments.clone(), }; + emit_call_started(ctx, &call_id, ReplCallKind::Tool, &tool_name); let start = Instant::now(); - let result = bridge_block_on(ctx.buffers.deadline(), tool.call(&ctx.state, call)) - .map_err(|err| raise(ctx, err))?; + let result = bridge_block_on( + ctx.buffers.deadline(), + &ctx.cancel, + tool.call(&ctx.state, call), + ) + .map_err(|err| raise(ctx, err))?; let elapsed = start.elapsed(); record( ctx, + call_id, ReplCallKind::Tool, &tool_name, json!({ "arguments": arguments }), @@ -107,11 +121,18 @@ pub(super) fn agent_query_impl( if let Some(data) = map_json(params, "input") { input = input.with_data(data); } + let call_id = new_call_id(); + emit_call_started(ctx, &call_id, ReplCallKind::Agent, &agent_name); let start = Instant::now(); - let output = bridge_block_on(ctx.buffers.deadline(), agent.run(input, ctx.events.clone())) - .map_err(|err| raise(ctx, err))?; + let output = bridge_block_on( + ctx.buffers.deadline(), + &ctx.cancel, + agent.run(input, ctx.events.clone()), + ) + .map_err(|err| raise(ctx, err))?; record( ctx, + call_id, ReplCallKind::Agent, &agent_name, json!({ "model_calls": output.model_calls, "tool_calls": output.tool_calls }), @@ -149,6 +170,7 @@ pub(super) fn graph_run_impl( .clone(); record( ctx, + new_call_id(), ReplCallKind::Graph, &graph_name, json!({ "nodes": blueprint.nodes.len() }), diff --git a/src/repl/session/builtins/mod.rs b/src/repl/session/builtins/mod.rs index 5f624cd..ad33f69 100644 --- a/src/repl/session/builtins/mod.rs +++ b/src/repl/session/builtins/mod.rs @@ -49,13 +49,13 @@ use std::time::{Duration, Instant}; use rhai::{Array, Dynamic, Engine, EvalAltResult, Map, Position}; use serde_json::{Value, json}; -use super::types::{GraphBlueprintHandle, LanguageCompiler, ReplPolicy}; +use super::types::{GraphBlueprintHandle, LanguageCompiler, ReplCancelFlag, ReplPolicy}; use super::{ ReplCallKind, ReplCallRecord, dynamic_to_repl_value, json_to_repl_value, repl_value_to_dynamic, }; use crate::error::TinyAgentsError; -use crate::harness::events::EventSink; -use crate::harness::ids::new_call_id; +use crate::harness::events::{AgentEvent, EventSink, ReplCallPhase}; +use crate::harness::ids::{CallId, new_call_id}; use crate::harness::message::Message; use crate::harness::model::ModelRequest; use crate::harness::tool::ToolCall; @@ -102,6 +102,9 @@ pub(super) struct HostContext { pub run_depth: usize, /// The event sink shared with the run context. pub events: EventSink, + /// External cancellation flag, observed fail-closed by the `on_progress` + /// hook (mid-script) and the blocking capability bridge (mid-call). + pub cancel: ReplCancelFlag, /// Per-cell shared buffers (stdout, calls, answer, host error, vars). pub buffers: super::CellBuffers, /// Session-cumulative call counters. @@ -110,58 +113,117 @@ pub(super) struct HostContext { pub drafts: Arc>>, } -/// Drives an async capability future to completion synchronously, bounded by -/// an optional wall-clock `deadline` — the v1 "blocking bridge" adapter (see -/// the [module docs](self)), now with fail-closed enforcement of -/// [`ReplPolicy::timeout`]. +/// How often the watcher thread in [`bridge_block_on_raw`] wakes to observe an +/// armed [`ReplCancelFlag`] while a capability call is in flight. +/// +/// Small enough that a user cancel releases a hung call promptly, large enough +/// that watching a fast (scripted-test) call costs nothing measurable. +const CANCEL_POLL_INTERVAL: Duration = Duration::from_millis(25); + +/// Why the watcher tripped a bounded [`bridge_block_on_raw`] call. +enum BridgeStop { + /// The per-cell wall-clock deadline elapsed. + Deadline, + /// The external [`ReplCancelFlag`] was tripped. + Cancelled, +} + +/// Drives an async capability future to completion synchronously, bounded by an +/// optional wall-clock `deadline` **and** an external `cancel` flag — the v1 +/// "blocking bridge" adapter (see the [module docs](self)), with fail-closed +/// enforcement of both [`ReplPolicy::timeout`] and host cancellation. /// /// `on_progress` (see [`build_engine`]) only fires between Rhai -/// statements/operations, so it can never interrupt a blocked native call: -/// this is the enforcement point for that case. When `deadline` elapses -/// first, the future is dropped — canceling the underlying request, since -/// providers are built on cancel-safe `reqwest`/`futures` — and a `Timeout` -/// error is returned instead of blocking the session forever. +/// statements/operations, so it can never interrupt a blocked native call: this +/// is the enforcement point for that case. A detached watcher thread races the +/// capability future; when the deadline elapses or `cancel` trips first, the +/// future is dropped — canceling the underlying request, since providers are +/// built on cancel-safe `reqwest`/`futures` — and a `Timeout` or `Cancelled` +/// error is returned instead of blocking the session forever. If the future +/// finishes first, the watcher observes the dropped receiver and exits. fn bridge_block_on_raw( deadline: Option, + cancel: &ReplCancelFlag, future: F, ) -> std::result::Result { - let Some(deadline) = deadline else { - return Ok(futures::executor::block_on(future)); - }; - let now = Instant::now(); - if now >= deadline { + // Fail closed before the call even starts if either bound has already + // tripped, so a cancel/timeout that landed between statements is honored + // without dispatching the call at all. + if cancel.is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + if let Some(deadline) = deadline + && Instant::now() >= deadline + { return Err(TinyAgentsError::Timeout(format!( "{DEADLINE_EXCEEDED_TOKEN} before a host capability call could start" ))); } - let remaining = deadline - now; - let (tx, rx) = futures::channel::oneshot::channel::<()>(); - // A detached timer thread wakes the race below at the deadline. If the - // capability future finishes first this thread simply sleeps out its - // remainder and exits; nothing observes it after that. + + let (tx, rx) = futures::channel::oneshot::channel::(); + let watcher_cancel = cancel.clone(); + // A detached watcher wakes the race below when the deadline elapses or the + // cancel flag trips. If the capability future finishes first, `rx` is + // dropped and `tx.is_canceled()` lets the watcher exit promptly instead of + // polling out a full deadline. std::thread::spawn(move || { - std::thread::sleep(remaining); - let _ = tx.send(()); + loop { + if tx.is_canceled() { + return; + } + if watcher_cancel.is_cancelled() { + let _ = tx.send(BridgeStop::Cancelled); + return; + } + match deadline { + Some(deadline) => { + let now = Instant::now(); + if now >= deadline { + let _ = tx.send(BridgeStop::Deadline); + return; + } + std::thread::sleep((deadline - now).min(CANCEL_POLL_INTERVAL)); + } + None => std::thread::sleep(CANCEL_POLL_INTERVAL), + } + } }); + match futures::executor::block_on(futures::future::select(Box::pin(future), rx)) { - futures::future::Either::Left((output, _timer)) => Ok(output), - futures::future::Either::Right(_) => Err(TinyAgentsError::Timeout(format!( - "{DEADLINE_EXCEEDED_TOKEN} during a host capability call" - ))), + futures::future::Either::Left((output, _watcher)) => Ok(output), + futures::future::Either::Right((stop, _fut)) => match stop { + Ok(BridgeStop::Cancelled) => Err(TinyAgentsError::Cancelled), + Ok(BridgeStop::Deadline) => Err(TinyAgentsError::Timeout(format!( + "{DEADLINE_EXCEEDED_TOKEN} during a host capability call" + ))), + // The watcher dropped its sender without sending — only reachable in + // a race the future has effectively already won; re-check the bounds + // and prefer cancellation, never silently succeeding. + Err(_canceled) => { + if cancel.is_cancelled() { + Err(TinyAgentsError::Cancelled) + } else { + Err(TinyAgentsError::Timeout(format!( + "{DEADLINE_EXCEEDED_TOKEN} during a host capability call" + ))) + } + } + }, } } /// Convenience wrapper over [`bridge_block_on_raw`] for the common case where -/// the capability future itself resolves to a `Result`, flattening the -/// deadline error into the same error channel as the call's own failures. +/// the capability future itself resolves to a `Result`, flattening the deadline +/// / cancellation error into the same error channel as the call's own failures. fn bridge_block_on( deadline: Option, + cancel: &ReplCancelFlag, future: F, ) -> std::result::Result where F: Future>, { - bridge_block_on_raw(deadline, future)? + bridge_block_on_raw(deadline, cancel, future)? } /// One completed `model_query_batched` item: `(model, text, finish_reason, @@ -192,23 +254,67 @@ fn invalid( raise(ctx, TinyAgentsError::Validation(message.into())) } -/// Records a capability call (or emitted event) into the per-cell buffer. +/// Records a completed capability call (or emitted event) into the per-cell +/// buffer **and** streams it live on the session [`EventSink`] as an +/// [`AgentEvent::ReplCall`] with phase [`ReplCallPhase::Completed`]. +/// +/// `call_id` is generated by the caller up front so a preceding +/// [`emit_call_started`] event can carry the same id, letting a host pair the +/// start and completion of one call. fn record( ctx: &HostContext, + call_id: CallId, kind: ReplCallKind, name: &str, detail: Value, elapsed: Duration, ) { - ctx.buffers.push_call(ReplCallRecord { - call_id: new_call_id(), + let record = ReplCallRecord { + call_id, kind, name: name.to_string(), detail, elapsed, + }; + emit_repl_call(ctx, &record, ReplCallPhase::Completed); + ctx.buffers.push_call(record); +} + +/// Emits an [`AgentEvent::ReplCall`] on the session event sink so a live +/// observer sees a capability call as it happens, rather than only in +/// [`ReplResult::calls`](super::ReplResult) after the cell returns. +fn emit_repl_call( + ctx: &HostContext, + record: &ReplCallRecord, + phase: ReplCallPhase, +) { + ctx.events.emit(AgentEvent::ReplCall { + session_id: ctx.session_label.clone(), + record: record.clone(), + phase, }); } +/// Streams a `ReplCall` "started" event for a capability call about to be +/// dispatched, carrying `call_id` (matched by the later [`record`] completion), +/// its kind, and name — but no `detail` (arguments are only in the completed +/// record) and a zero `elapsed`. +fn emit_call_started( + ctx: &HostContext, + call_id: &CallId, + kind: ReplCallKind, + name: &str, +) { + let record = ReplCallRecord { + call_id: call_id.clone(), + kind, + name: name.to_string(), + detail: Value::Null, + elapsed: Duration::default(), + }; + emit_repl_call(ctx, &record, ReplCallPhase::Started); +} + // ── Map argument helpers ──────────────────────────────────────────────────── /// Reads a string field from a Rhai object map argument. @@ -358,6 +464,12 @@ use capabilities::*; /// generic runtime-error path. pub(super) const DEADLINE_EXCEEDED_TOKEN: &str = "ragsh cell exceeded its wall-clock timeout"; +/// Sentinel exception value `on_progress` terminates a script with when an +/// external [`ReplCancelFlag`] is tripped mid-script. `eval_cell`'s +/// `map_rhai_error` recognizes this exact string and maps it to +/// [`TinyAgentsError::Cancelled`] instead of the generic runtime-error path. +pub(super) const CANCELLED_TOKEN: &str = "ragsh cell cancelled by host"; + /// Builds a sandboxed Rhai engine for a session, registering every host-backed /// built-in against the session's live registries and policy. /// @@ -378,6 +490,13 @@ pub(super) fn build_engine(ctx: Arc(ctx: Arc(ctx: Arc(None, async { Ok(7) }).expect("no deadline"); + let out = bridge_block_on::(None, &ReplCancelFlag::new(), async { Ok(7) }) + .expect("no deadline"); assert_eq!(out, 7); } @@ -542,7 +664,8 @@ mod bridge_deadline_test { fn future_finishing_before_the_deadline_succeeds() { let deadline = Instant::now() + Duration::from_secs(5); let out = - bridge_block_on::(Some(deadline), async { Ok(9) }).expect("within deadline"); + bridge_block_on::(Some(deadline), &ReplCancelFlag::new(), async { Ok(9) }) + .expect("within deadline"); assert_eq!(out, 9); } @@ -551,8 +674,9 @@ mod bridge_deadline_test { // Regression test: `ReplPolicy::timeout` used to be parsed but never // enforced anywhere a host capability call could hang forever. let deadline = Instant::now() - Duration::from_millis(1); - let err = bridge_block_on::(Some(deadline), async { Ok(1) }) - .expect_err("deadline already passed"); + let err = + bridge_block_on::(Some(deadline), &ReplCancelFlag::new(), async { Ok(1) }) + .expect_err("deadline already passed"); assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); } @@ -565,6 +689,7 @@ mod bridge_deadline_test { let deadline = start + Duration::from_millis(30); let err = bridge_block_on::( Some(deadline), + &ReplCancelFlag::new(), futures::future::pending::>(), ) .expect_err("hanging call must be cut off at the deadline"); @@ -575,4 +700,45 @@ mod bridge_deadline_test { start.elapsed() ); } + + #[test] + fn cancel_already_set_fails_closed_without_starting_the_call() { + // A flag tripped before the bridge runs must short-circuit to + // `Cancelled` without ever polling the (here, never-resolving) future. + let cancel = ReplCancelFlag::new(); + cancel.cancel(); + let err = bridge_block_on::( + None, + &cancel, + futures::future::pending::>(), + ) + .expect_err("pre-cancelled call must not start"); + assert!(matches!(err, TinyAgentsError::Cancelled), "got {err:?}"); + } + + #[test] + fn a_hanging_call_is_cut_off_promptly_when_the_cancel_flag_trips() { + // With no deadline, a hung capability future must still be released + // once a host trips the cancel flag from another thread — the watcher + // polls the flag and drops the future. + let start = Instant::now(); + let cancel = ReplCancelFlag::new(); + let trigger = cancel.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(40)); + trigger.cancel(); + }); + let err = bridge_block_on::( + None, + &cancel, + futures::future::pending::>(), + ) + .expect_err("hanging call must be cut off on cancel"); + assert!(matches!(err, TinyAgentsError::Cancelled), "got {err:?}"); + assert!( + start.elapsed() < Duration::from_secs(5), + "took {:?}, should return promptly after the ~40ms cancel", + start.elapsed() + ); + } } diff --git a/src/repl/session/mod.rs b/src/repl/session/mod.rs index 6730143..249246a 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -230,6 +230,11 @@ where /// against [`ReplPolicy::max_iterations`]. Each `eval_cell` call is one /// CodeAct-style iteration of a model-driven session. iterations: usize, + /// External cancellation flag. A host holding a clone can abort an in-flight + /// cell (see [`ReplCancelFlag`]); enforced fail-closed in both the engine + /// `on_progress` hook and the blocking capability bridge. Cloned into the + /// engine's [`HostContext`] on every [`rebuild_engine`](Self::rebuild_engine). + cancel: ReplCancelFlag, } impl ReplSession { @@ -284,6 +289,7 @@ impl ReplSession { engine: Engine::new(), buffers, iterations: 0, + cancel: ReplCancelFlag::new(), }; session.rebuild_engine(); session @@ -307,10 +313,31 @@ impl ReplSession { buffers: self.buffers.clone(), counters: self.counters.clone(), drafts: self.drafts.clone(), + cancel: self.cancel.clone(), }); self.engine = builtins::build_engine(ctx); } + /// Installs an external [`ReplCancelFlag`] and rebuilds the engine so the + /// `on_progress` hook and the blocking capability bridge observe it. + /// + /// The host keeps a clone of `flag` and calls [`ReplCancelFlag::cancel`] to + /// abort an in-flight cell; the cell then fails with + /// [`TinyAgentsError::Cancelled`]. Because a cancelled flag is sticky, pass a + /// **fresh** flag when reusing a session whose previous run was cancelled. + pub fn with_cancel_flag(mut self, flag: ReplCancelFlag) -> Self { + self.cancel = flag; + self.rebuild_engine(); + self + } + + /// Returns a clone of this session's cancellation flag, so a host that did + /// not supply one via [`with_cancel_flag`](Self::with_cancel_flag) can still + /// obtain the handle needed to abort an in-flight cell. + pub fn cancel_flag(&self) -> ReplCancelFlag { + self.cancel.clone() + } + /// Replaces the session policy and rebuilds the engine to honor the new /// operation and call limits. pub fn with_policy(mut self, policy: ReplPolicy) -> Self { @@ -364,6 +391,33 @@ impl ReplSession { /// names are restored afterward so the next cell starts from a clean /// capability baseline. /// + /// # Blocking — driving this from an async host + /// + /// **This method blocks the calling thread.** The Rhai engine is + /// synchronous, so each `model_query`/`tool_call`/`agent_query` a cell + /// performs is driven to completion by an internal + /// [`futures::executor::block_on`] (the "blocking bridge"; see + /// [`builtins`](self::builtins)). Calling `eval_cell` directly on an async + /// worker therefore blocks that worker for the whole cell, and on a + /// **current-thread** Tokio runtime it deadlocks — `block_on` parks the only + /// worker the in-flight capability future needs to make progress. + /// + /// An async host **must** run `eval_cell` off the async workers, on a + /// blocking-safe thread: + /// + /// ```ignore + /// let result = tokio::task::spawn_blocking(move || session.eval_cell(&script)).await?; + /// ``` + /// + /// A multi-threaded runtime with a spare worker also works, but + /// `spawn_blocking` (or a dedicated thread) is the contract. Because + /// `eval_cell` takes `&mut self`, only one cell runs per session at a time; + /// the host serializes concurrent calls to the same session. To bound a + /// cell's wall clock from the async side as well, wrap the join handle in a + /// [`tokio::time::timeout`] and install a [`ReplCancelFlag`] via + /// [`with_cancel_flag`](Self::with_cancel_flag) so the blocked worker is + /// released rather than leaked. + /// /// # Errors /// /// * [`TinyAgentsError::LimitExceeded`] — the script exceeds @@ -376,9 +430,22 @@ impl ReplSession { /// model/tool/agent/graph call. /// * [`TinyAgentsError::Validation`] — the script failed to compile or /// raised a runtime error. + /// * [`TinyAgentsError::Cancelled`] — an external [`ReplCancelFlag`] was + /// tripped before or during the cell (mid-script via the `on_progress` + /// hook, or during an in-flight capability call via the blocking bridge). pub fn eval_cell(&mut self, script: &str) -> Result { let start = Instant::now(); + // Fail closed if cancellation was requested before this cell even + // starts: a host that cancels between cells must not have its next cell + // begin any script or capability work. (Mid-cell cancellation is + // enforced separately by the `on_progress` hook and the capability + // bridge.) The iteration counter is left untouched so a cancelled, + // never-run cell does not consume the session's `max_iterations` budget. + if self.cancel.is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + // Each call is one CodeAct-style iteration of a model-driven session; // enforce the cap fail-closed before doing any other work. if self.iterations >= self.policy.max_iterations { @@ -618,6 +685,15 @@ fn map_rhai_error(err: EvalAltResult) -> TinyAgentsError { "ragsh cell exceeded the operation limit (max_operations) at {pos}" )), // The engine's `on_progress` hook (see `builtins::build_engine`) + // terminates the script with this exact sentinel value once an external + // [`ReplCancelFlag`] is tripped mid-script; map it to `Cancelled` so the + // host sees a cancellation rather than a generic validation error. + EvalAltResult::ErrorTerminated(token, _pos) + if token.clone().into_string().ok().as_deref() == Some(builtins::CANCELLED_TOKEN) => + { + TinyAgentsError::Cancelled + } + // The engine's `on_progress` hook (see `builtins::build_engine`) // terminates the script with this exact sentinel value once the // per-cell `ReplPolicy::timeout` deadline elapses. EvalAltResult::ErrorTerminated(token, pos) diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index 2d59cc3..a199d5b 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -482,6 +482,176 @@ fn agent_call_limit_is_independent_of_the_model_call_limit() { } } +// ── External cancellation (Phase 2) ────────────────────────────────────────── + +/// A tool whose call never resolves, modelling a hung provider/tool so a test +/// can prove external cancellation drops the in-flight future via the blocking +/// bridge rather than blocking the session forever. +struct HangingTool; + +#[async_trait::async_trait] +impl crate::harness::tool::Tool<()> for HangingTool { + fn name(&self) -> &str { + "hangs" + } + + fn description(&self) -> &str { + "Never returns; used to test the cancellation bridge." + } + + fn schema(&self) -> crate::harness::tool::ToolSchema { + crate::harness::tool::ToolSchema { + name: self.name().to_string(), + description: self.description().to_string(), + parameters: serde_json::json!({ "type": "object" }), + format: Default::default(), + } + } + + async fn call( + &self, + _state: &(), + _call: crate::harness::tool::ToolCall, + ) -> crate::Result { + futures::future::pending().await + } +} + +fn session_with_hanging_tool(policy: ReplPolicy) -> ReplSession { + let mut registry = crate::registry::CapabilityRegistry::<()>::new(); + registry + .register_tool(std::sync::Arc::new(HangingTool)) + .expect("register hanging tool"); + let capabilities = ReplCapabilities::new(std::sync::Arc::new(registry)); + ReplSession::<()>::new() + .with_policy(policy) + .with_capabilities(capabilities) +} + +#[test] +fn cancel_set_before_eval_fails_closed_without_running_the_cell() { + // A flag tripped before the cell starts short-circuits to `Cancelled`, and + // must not consume the session's iteration budget: a later cell with a + // fresh flag still runs. + let flag = ReplCancelFlag::new(); + let mut s = session().with_cancel_flag(flag.clone()); + flag.cancel(); + + let err = s.eval_cell("let x = 1; x").expect_err("pre-cancelled"); + assert!(matches!(err, TinyAgentsError::Cancelled), "got {err:?}"); + + // A fresh flag re-enables the session — the cancelled cell did not burn an + // iteration or otherwise poison the namespace. + let mut s = s.with_cancel_flag(ReplCancelFlag::new()); + let ok = s.eval_cell("1 + 1").expect("session usable after cancel"); + assert_eq!(ok.value, Some(ReplValue::Int(2))); +} + +#[test] +fn cancel_mid_script_loop_terminates_promptly() { + // A pure `loop {}` with no timeout and an unbounded operation budget can + // only be stopped by the cancel flag, enforced via the `on_progress` hook. + let policy = ReplPolicy { + timeout: None, + max_operations: 0, + ..ReplPolicy::default() + }; + let flag = ReplCancelFlag::new(); + let mut s = session().with_policy(policy).with_cancel_flag(flag.clone()); + + let trigger = flag.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(40)); + trigger.cancel(); + }); + + let start = Instant::now(); + let err = s + .eval_cell("let total = 0; loop { total += 1; }") + .expect_err("cancel must terminate the loop"); + assert!(matches!(err, TinyAgentsError::Cancelled), "got {err:?}"); + assert!( + start.elapsed() < Duration::from_secs(5), + "eval_cell took {:?}, should return near the ~40ms cancel", + start.elapsed() + ); +} + +#[test] +fn cancel_during_a_hanging_capability_call_terminates_promptly() { + // A tool call that never resolves can only be released by the cancel flag, + // enforced via the blocking bridge (`on_progress` never fires inside a + // blocked native call). No timeout is configured so only cancel can stop it. + let policy = ReplPolicy { + timeout: None, + ..ReplPolicy::default() + }; + let flag = ReplCancelFlag::new(); + let mut s = session_with_hanging_tool(policy).with_cancel_flag(flag.clone()); + + let trigger = flag.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(40)); + trigger.cancel(); + }); + + let start = Instant::now(); + let err = s + .eval_cell(r#"tool_call(#{ tool: "hangs" })"#) + .expect_err("cancel must drop the hung tool future"); + assert!(matches!(err, TinyAgentsError::Cancelled), "got {err:?}"); + assert!( + start.elapsed() < Duration::from_secs(5), + "eval_cell took {:?}, should return near the ~40ms cancel", + start.elapsed() + ); +} + +// ── Live capability-call events (Phase 2) ──────────────────────────────────── + +#[test] +fn capability_calls_stream_started_and_completed_events() { + use crate::harness::events::{AgentEvent, ReplCallPhase}; + + // `fail_id` "none" never matches, so the "1" call succeeds. + let mut s = session_with_sometimes_failing_tool("none"); + let recorder = std::sync::Arc::new(crate::harness::events::RecordingListener::new()); + s.events.subscribe(recorder.clone()); + + let session_label = s.session_id.as_str().to_string(); + s.eval_cell(r#"tool_call(#{ tool: "sometimes_fails", arguments: #{ id: "1" } })"#) + .expect("tool call"); + + // Exactly one Started and one Completed ReplCall event, paired by call_id, + // both naming the tool and correlated back to this session. + let repl_calls: Vec<(ReplCallPhase, String, String)> = recorder + .events() + .into_iter() + .filter_map(|rec| match rec.event { + AgentEvent::ReplCall { + session_id, + record, + phase, + } => Some((phase, session_id, record.call_id.as_str().to_string())), + _ => None, + }) + .collect(); + + assert_eq!( + repl_calls.len(), + 2, + "expected start + completion: {repl_calls:?}" + ); + assert_eq!(repl_calls[0].0, ReplCallPhase::Started); + assert_eq!(repl_calls[1].0, ReplCallPhase::Completed); + assert_eq!(repl_calls[0].1, session_label, "session_id correlates"); + assert_eq!(repl_calls[1].1, session_label); + assert_eq!( + repl_calls[0].2, repl_calls[1].2, + "start and completion share one call_id" + ); +} + #[test] fn map_and_array_values_round_trip_to_json() { let mut s = session(); diff --git a/src/repl/session/types.rs b/src/repl/session/types.rs index 2058bb1..1c984ec 100644 --- a/src/repl/session/types.rs +++ b/src/repl/session/types.rs @@ -10,6 +10,7 @@ use std::collections::BTreeMap; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use serde::{Deserialize, Serialize}; @@ -17,6 +18,49 @@ use serde::{Deserialize, Serialize}; use crate::language::types::{Blueprint, Origin}; use crate::registry::CapabilityRegistry; +// ── External cancellation ──────────────────────────────────────────────────── + +/// A shared, cheaply-cloneable cancellation flag for a [`super::ReplSession`]. +/// +/// A cell can otherwise only be stopped by its wall-clock +/// [`ReplPolicy::timeout`]. A host that drives sessions from an async runtime +/// (for example when a user cancels the surrounding run) needs to abort an +/// **in-flight** cell on demand, without waiting out the timeout. It holds a +/// clone of this flag and calls [`ReplCancelFlag::cancel`]; the session observes +/// it fail-closed at the same two enforcement points as the deadline: +/// +/// - the engine `on_progress` hook, which terminates a running *script* at the +/// next statement/operation (a pure `while true {}` loop with no host calls); +/// - the blocking bridge around every `model_query`/`tool_call`/`agent_query` +/// call, which drops an in-flight (possibly hung) capability future promptly. +/// +/// In both cases [`super::ReplSession::eval_cell`] returns +/// [`TinyAgentsError::Cancelled`](crate::error::TinyAgentsError::Cancelled). The +/// flag is *sticky*: once cancelled it stays cancelled, so a session observed as +/// cancelled will refuse to start further cells until a fresh flag is installed. +/// Construct one with [`ReplCancelFlag::new`] and install it with +/// [`super::ReplSession::with_cancel_flag`]. +#[derive(Clone, Debug, Default)] +pub struct ReplCancelFlag(Arc); + +impl ReplCancelFlag { + /// Creates a fresh, un-cancelled flag. + pub fn new() -> Self { + Self(Arc::new(AtomicBool::new(false))) + } + + /// Requests cancellation. Idempotent and safe to call from any thread; every + /// clone of this flag observes the change. + pub fn cancel(&self) { + self.0.store(true, Ordering::SeqCst); + } + + /// Returns whether cancellation has been requested. + pub fn is_cancelled(&self) -> bool { + self.0.load(Ordering::SeqCst) + } +} + // ── Reserved names ────────────────────────────────────────────────────────── /// Reserved built-in *variable* names seeded into every session scope. From df391c46891cb69ffd39f69023f497a5d95bec0e Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 4 Jul 2026 21:01:54 +0000 Subject: [PATCH 2/2] Add ReplSession::set_cancel_flag for per-cell cancel reset `with_cancel_flag` consumes `self`, which suits construction but not a long-lived session held behind a lock by a host session manager. Because a `ReplCancelFlag` is sticky (once cancelled, stays cancelled), a persistent session reused after a cancelled cell would otherwise refuse every subsequent cell. `set_cancel_flag(&mut self, flag)` installs a fresh flag in place and rebuilds the engine (the persistent variable namespace is untouched, so `let` bindings survive the swap), letting a host arm a new per-cell flag before each `eval_cell` and stay resumable after a cancel. Test covers the swap + namespace preservation. --- src/repl/session/mod.rs | 15 +++++++++++++++ src/repl/session/test.rs | 22 ++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/src/repl/session/mod.rs b/src/repl/session/mod.rs index 249246a..ec0715a 100644 --- a/src/repl/session/mod.rs +++ b/src/repl/session/mod.rs @@ -338,6 +338,21 @@ impl ReplSession { self.cancel.clone() } + /// Installs a fresh cancellation flag in place (`&mut self`), rebuilding the + /// engine so the `on_progress` hook and the blocking bridge observe it. + /// + /// Unlike the consuming [`with_cancel_flag`](Self::with_cancel_flag), this + /// swaps the flag on a session already owned behind a lock — the shape a + /// long-lived session manager needs. Because a cancelled flag is sticky, + /// installing a **fresh** flag before each cell lets a persistent session + /// stay resumable after a prior cell was cancelled. The persistent variable + /// namespace is untouched (only the engine is rebuilt), so `let` bindings + /// survive the swap. + pub fn set_cancel_flag(&mut self, flag: ReplCancelFlag) { + self.cancel = flag; + self.rebuild_engine(); + } + /// Replaces the session policy and rebuilds the engine to honor the new /// operation and call limits. pub fn with_policy(mut self, policy: ReplPolicy) -> Self { diff --git a/src/repl/session/test.rs b/src/repl/session/test.rs index a199d5b..a047603 100644 --- a/src/repl/session/test.rs +++ b/src/repl/session/test.rs @@ -547,6 +547,28 @@ fn cancel_set_before_eval_fails_closed_without_running_the_cell() { assert_eq!(ok.value, Some(ReplValue::Int(2))); } +#[test] +fn set_cancel_flag_swaps_in_place_and_preserves_the_namespace() { + // A long-lived session behind a lock swaps its flag with `set_cancel_flag` + // (not the consuming builder). A fresh flag re-enables a session whose + // prior cell was cancelled, and the persistent namespace survives the swap. + let mut s = session(); + s.eval_cell("let kept = 42;").expect("seed binding"); + + let tripped = ReplCancelFlag::new(); + tripped.cancel(); + s.set_cancel_flag(tripped); + let err = s + .eval_cell("kept") + .expect_err("cancelled flag blocks the cell"); + assert!(matches!(err, TinyAgentsError::Cancelled), "got {err:?}"); + + // A fresh flag re-enables the session, and `kept` is still in scope. + s.set_cancel_flag(ReplCancelFlag::new()); + let ok = s.eval_cell("kept").expect("session usable again"); + assert_eq!(ok.value, Some(ReplValue::Int(42))); +} + #[test] fn cancel_mid_script_loop_terminates_promptly() { // A pure `loop {}` with no timeout and an unbounded operation budget can