diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f3e693e3..e655f2aa 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -39,28 +39,28 @@ jobs: run: cargo fmt --all -- --check - name: Clippy - run: cargo clippy --all-targets -- -D warnings + run: cargo clippy --workspace --all-targets -- -D warnings - name: Clippy all features - run: cargo clippy --all-targets --all-features -- -D warnings + run: cargo clippy --workspace --all-targets --all-features -- -D warnings - name: Build - run: cargo build --all-targets + run: cargo build --workspace --all-targets - name: Build all features - run: cargo build --all-targets --all-features + run: cargo build --workspace --all-targets --all-features - name: Test - run: cargo test + run: cargo test --workspace - name: Test all features - run: cargo test --all-features + run: cargo test --workspace --all-features - name: Test optional features independently run: | - cargo test --no-default-features --features sqlite - cargo test --no-default-features --features tools - cargo test --no-default-features --features multimodal + cargo test --workspace --no-default-features --features sqlite + cargo test --workspace --no-default-features --features tools + cargo test --workspace --no-default-features --features multimodal - name: Coverage gate uses: taiki-e/install-action@cargo-llvm-cov @@ -70,3 +70,19 @@ jobs: cargo llvm-cov --all-features --workspace --ignore-filename-regex '(^|/)(tests?|examples)/|/test(_.*)?\.rs$' --fail-under-lines 80 + + # TODO: drop `continue-on-error` once the ~160 existing broken + # intra-doc link warnings across the workspace are fixed and this can + # block merges instead of only reporting. + - name: Doc lints + env: + RUSTDOCFLAGS: -D rustdoc::broken_intra_doc_links + run: cargo doc --workspace --no-deps + continue-on-error: true + + # TODO: drop `continue-on-error` once the existing unused-dependency + # findings across the workspace are cleaned up and this can block + # merges instead of only reporting. + - name: Unused dependencies + uses: bnjbvr/cargo-machete@main + continue-on-error: true diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 7f8c615f..bec23a7d 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -52,10 +52,26 @@ jobs: run: cargo fmt --all -- --check - name: Run clippy - run: cargo clippy --all-targets -- -D warnings + run: cargo clippy --workspace --all-targets -- -D warnings - name: Run tests - run: cargo test + run: cargo test --workspace + + # TODO: drop `continue-on-error` once the existing ~160 broken + # intra-doc link warnings across the workspace are fixed and this can + # block merges instead of only reporting. + - name: Doc lints + env: + RUSTDOCFLAGS: -D rustdoc::broken_intra_doc_links + run: cargo doc --workspace --no-deps + continue-on-error: true + + # TODO: drop `continue-on-error` once the existing unused-dependency + # findings across the workspace are cleaned up and this can block + # merges instead of only reporting. + - name: Unused dependencies + uses: bnjbvr/cargo-machete@main + continue-on-error: true - name: Compute next version id: version diff --git a/Cargo.lock b/Cargo.lock index 0bcfff81..2917f2f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -777,15 +777,6 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" -[[package]] -name = "lock_api" -version = "0.4.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" -dependencies = [ - "scopeguard", -] - [[package]] name = "log" version = "0.4.33" @@ -846,29 +837,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" -[[package]] -name = "parking_lot" -version = "0.12.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" -dependencies = [ - "lock_api", - "parking_lot_core", -] - -[[package]] -name = "parking_lot_core" -version = "0.9.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" -dependencies = [ - "cfg-if", - "libc", - "redox_syscall", - "smallvec", - "windows-link", -] - [[package]] name = "percent-encoding" version = "2.3.2" @@ -1037,15 +1005,6 @@ dependencies = [ "getrandom 0.3.4", ] -[[package]] -name = "redox_syscall" -version = "0.5.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" -dependencies = [ - "bitflags", -] - [[package]] name = "redox_users" version = "0.5.3" @@ -1232,12 +1191,6 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" -[[package]] -name = "scopeguard" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" - [[package]] name = "serde" version = "1.0.229" @@ -1459,6 +1412,7 @@ version = "2.1.2" dependencies = [ "async-trait", "serde", + "serde_json", "tokio", ] @@ -1474,10 +1428,10 @@ dependencies = [ "serde_json", "tempfile", "tinyagents-harness", - "tinyagents-tracing", "tinyinference-llm", "tinytools", "tokio", + "tracing", ] [[package]] @@ -1492,8 +1446,6 @@ dependencies = [ "dirs", "flate2", "futures", - "libc", - "log", "regex", "reqwest", "rusqlite", @@ -1503,12 +1455,12 @@ dependencies = [ "tempfile", "thiserror", "tinyagents-definition", - "tinyagents-tracing", "tinyinference-embeddings", "tinyinference-llm", "tinytools", "tinytools-agent", "tokio", + "tracing", "uuid", "wait-timeout", ] @@ -1545,9 +1497,9 @@ version = "2.1.2" dependencies = [ "anyhow", "async-trait", + "chrono", "dotenvy", "futures", - "parking_lot", "serde", "serde_json", "thiserror", @@ -1567,6 +1519,7 @@ version = "2.1.2" dependencies = [ "anyhow", "async-trait", + "chrono", "serde", "serde_json", "tinyagents-definition", @@ -1598,20 +1551,14 @@ name = "tinyagents-session" version = "2.1.2" dependencies = [ "anyhow", + "async-trait", "chrono", - "log", "rusqlite", "serde", "serde_json", "tempfile", "tinyagents-harness", - "tinyagents-tracing", -] - -[[package]] -name = "tinyagents-tracing" -version = "2.1.2" -dependencies = [ + "tokio", "tracing", ] diff --git a/Cargo.toml b/Cargo.toml index 0817835f..8db469d7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,9 +16,24 @@ version = "2.1.2" edition = "2024" license = "GPL-3.0-only" repository = "https://github.com/tinyhumansai/tinyagents" +rust-version = "1.88" + +[workspace.dependencies] +anyhow = "1" +async-trait = "0.1" +chrono = { version = "0.4", features = ["serde"] } +futures = "0.3" +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } +rusqlite = { version = "0.40", features = ["bundled"] } +serde = { version = "1", features = ["derive", "rc"] } +serde_json = "1" +tempfile = "3" +tokio = { version = "1", default-features = false, features = ["macros", "rt-multi-thread"] } +tracing = "0.1" +uuid = { version = "1", features = ["v4"] } [workspace.lints.rust] -unsafe_code = "allow" +unsafe_code = "deny" [workspace.lints.clippy] all = { level = "warn", priority = -1 } diff --git a/README.md b/README.md index c3c87b29..9fb17b0d 100644 --- a/README.md +++ b/README.md @@ -38,11 +38,16 @@ TinyAgents is a Cargo workspace, not one crate. Depend on the pieces you need: agents, graphs, and routers), plus an offline model price/capability catalog. - **`tinyagents-session`** — a SQLite-backed store for session history, messages, tool calls, cost, and run lineage. +- **`tinyagents-definition`** — the host-owned agent definition vocabulary: + identity, description, declared model/tools/delegates, and a read-only + catalogue seam. Authorization, prompt construction, and execution stay with + the host and harness. - **`tinyagents-runtime`** — host-neutral stateful turns over the harness and append-only transcript seam; hosts retain policy, prompt composition, authorization, and durable-dialect conversion. -- **`tinyagents-tracing`** — the `tracing` macros the other crates gate behind - their `tracing` feature. Compiled out by default. +- **`tinyagents-orchestration`** — host-neutral composition of durable + multi-agent work (teams and workflows) over the graph, harness, and session + layers; depends one-way on those crates and stays host-free. - **`tinyagents-integration-tests`** — cross-crate tests and the runnable examples referenced below (not published, workspace-internal). diff --git a/ROADMAP.md b/ROADMAP.md index 2dbec98e..4e9833f2 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -25,7 +25,11 @@ build toward a production-grade Rust agent runtime. - expand live (network-gated) provider contract tests as new OpenAI-compatible endpoints are added - track and close the internal SDK feature-parity backlog in - [`docs/sdk-gaps.md`](docs/sdk-gaps.md) + [`docs/sdk-gaps/README.md`](docs/sdk-gaps/README.md) +- execute the phased plan in + [`docs/runtime-comparison/plan.md`](docs/runtime-comparison/plan.md), which + ranks the correctness fixes and feature gaps found by comparing TinyAgents + with LangGraph, Pydantic AI and pi ## Parallel Agents And Sub-Agents diff --git a/build.log b/build.log new file mode 100644 index 00000000..1c26b587 --- /dev/null +++ b/build.log @@ -0,0 +1,4 @@ + Compiling tinyagents-harness v2.1.2 (/home/enamakel/work/tinyagents/worktrees/runtime-comparison/crates/tinyagents-harness) + Compiling tinyagents-language v2.1.2 (/home/enamakel/work/tinyagents/worktrees/runtime-comparison/crates/tinyagents-language) + Compiling tinyagents-graph v2.1.2 (/home/enamakel/work/tinyagents/worktrees/runtime-comparison/crates/tinyagents-graph) + Finished `dev` profile [unoptimized + debuginfo] target(s) in 7.79s diff --git a/crates/tinyagents-definition/Cargo.toml b/crates/tinyagents-definition/Cargo.toml index 64a15420..000c439b 100644 --- a/crates/tinyagents-definition/Cargo.toml +++ b/crates/tinyagents-definition/Cargo.toml @@ -5,14 +5,16 @@ version.workspace = true edition.workspace = true license.workspace = true repository.workspace = true +rust-version.workspace = true description = "Host-owned agent definition contract." [dependencies] -async-trait = "0.1" -serde = { version = "1", features = ["derive"] } +async-trait = { workspace = true } +serde = { workspace = true } [dev-dependencies] -tokio = { version = "1", features = ["macros", "rt-multi-thread"] } +serde_json = { workspace = true } +tokio = { workspace = true, default-features = true, features = ["macros", "rt-multi-thread"] } [lints] workspace = true diff --git a/crates/tinyagents-graph/Cargo.toml b/crates/tinyagents-graph/Cargo.toml index adebe6cd..71f5b2dc 100644 --- a/crates/tinyagents-graph/Cargo.toml +++ b/crates/tinyagents-graph/Cargo.toml @@ -5,29 +5,35 @@ version.workspace = true edition.workspace = true license.workspace = true repository.workspace = true +rust-version.workspace = true description = "Durable typed state graphs for TinyAgents." [dependencies] -anyhow = "1" -async-trait = "0.1" -futures = "0.3" -rusqlite = { version = "0.40", features = ["bundled"], optional = true } -serde = { version = "1", features = ["derive"] } -serde_json = "1" -tinyagents-harness = { path = "../tinyagents-harness", version = "2.1.2", default-features = false } -tinyagents-tracing = { path = "../tinyagents-tracing", version = "2.1.2", default-features = false } +anyhow = { workspace = true } +async-trait = { workspace = true } +futures = { workspace = true } +rusqlite = { workspace = true, optional = true } +serde = { workspace = true } +serde_json = { workspace = true } +tinyagents-harness = { path = "../tinyagents-harness", version = "2.1.2", default-features = false, features = [ + "langfuse", +] } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.3.0" } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.3.0" } -tokio = { version = "1", default-features = false, features = ["sync", "time", "macros", "rt", "rt-multi-thread", "fs"] } +tokio = { workspace = true, features = ["sync", "time", "macros", "rt", "rt-multi-thread", "fs"] } +tracing = { workspace = true } [features] default = [] sqlite = ["dep:rusqlite"] -tracing = ["tinyagents-harness/tracing", "tinyagents-tracing/tracing"] +# Tracing instrumentation is now always compiled in (via the `tracing` crate +# dependency above). This feature is retained as a no-op so downstream +# feature forwards keep compiling. +tracing = ["tinyagents-harness/tracing"] [dev-dependencies] -tempfile = "3" -tokio = { version = "1", features = ["macros", "rt-multi-thread", "time", "test-util"] } +tempfile = { workspace = true } +tokio = { workspace = true, default-features = true, features = ["macros", "rt-multi-thread", "time", "test-util"] } [lints] workspace = true diff --git a/crates/tinyagents-graph/src/agent_loop/compile.rs b/crates/tinyagents-graph/src/agent_loop/compile.rs new file mode 100644 index 00000000..7ad4b9ac --- /dev/null +++ b/crates/tinyagents-graph/src/agent_loop/compile.rs @@ -0,0 +1,171 @@ +//! [`compile_loop`]: assembles the `plan -> model -> tools -> settle` graph +//! over a [`LoopRuntime`]. +//! +//! See the module doc on [`super`] for the loop's documented scope and the +//! shape of [`LoopState`]/[`LoopUpdate`]. + +use std::sync::Arc; + +use crate::builder::{GraphBuilder, NodeContext}; +use crate::compiled::CompiledGraph; +use tinyagents_harness::error::Result; +use tinyagents_harness::runtime::AgentHarness; + +use super::runtime::{self, LoopRuntime}; +use super::types::{LoopState, LoopUpdate, node}; + +/// Compiles the `plan -> model -> tools -> settle` agent loop into a +/// [`CompiledGraph`] bound to `rt`. +/// +/// # Shape +/// +/// [`LoopState`] doubles as its own [`LoopUpdate`] (`GraphBuilder::overwrite`, +/// see that type's docs): every node returns the whole next state rather than +/// a partial patch, so the reducer is a plain overwrite and there is no +/// separate merge step to keep in sync with the state shape. +/// +/// Every node routes with an explicit [`crate::Command::goto`] +/// (`mark_command_routing`) rather than static/conditional edges — the +/// `plan`/`model`/`tools` nodes can each jump to more than one destination +/// depending on the turn's outcome (a tool-free response skips `tools`; a +/// [`tinyagents_harness::context::MiddlewareControl::JumpTo`] can loop back +/// to `plan` from `tools`, or straight to `settle` from `model`/`tools`), so a +/// fixed edge table cannot express the routing — see `runtime::apply_control` +/// for the full mapping from [`tinyagents_harness::context::MiddlewareControl`] +/// to a `goto`. `plan`, `model`, and `tools` are also marked as interrupt +/// points for the export (see this function's body, below): a steering +/// pause (from `plan`) or a +/// [`tinyagents_harness::context::MiddlewareControl::Interrupt`] (from any of +/// the three) surfaces as a real [`crate::Interrupt`], checkpointed by the +/// graph executor exactly like any other durable interrupt — this is how +/// A5's "approvals surfacing as graph interrupts" requirement is met: an +/// approval gate is just a middleware that requests +/// `MiddlewareControl::Interrupt`, and the graph rendition pauses/resumes it +/// through the same `CompiledGraph::resume` path a hand-authored +/// human-in-the-loop node would use. +/// +/// # Tool batch execution +/// +/// The `tools` node calls +/// [`tinyagents_harness::agent_loop::phases::execute_tool_batch`] once per +/// activation, for the *whole* turn's tool calls in one node run rather than +/// one graph node per tool call. This preserves the direct loop's ordering, +/// concurrency-eligibility, and budget/limit semantics exactly (see that +/// function's docs) without re-deriving them as a `Send`-fanout over +/// per-call nodes, which would have to reimplement the direct loop's +/// serial-admission / serial-or-concurrent-dispatch decision as graph +/// topology instead of reusing it. +/// +/// # Recursion / limits +/// +/// [`tinyagents_harness::limits::RunLimits`] is enforced the same way the +/// direct loop enforces it — via `ctx.record_model_call()`/tool budget +/// checks inside the node bodies, fed from the same +/// [`tinyagents_harness::context::RunContext::limits`] — so a caller who also +/// wants the *graph's own* recursion guard (independent of the harness's +/// limits) can additionally call +/// [`crate::compiled::CompiledGraph::with_recursion_policy`]/ +/// [`crate::compiled::CompiledGraph::with_run_deadline`] on the returned +/// graph, exactly as any other compiled graph. +pub fn compile_loop( + rt: Arc>, +) -> Result> +where + State: Send + Sync + 'static, + Ctx: Send + Sync + 'static, +{ + let mut builder = GraphBuilder::::overwrite() + .with_name("tinyagents.agent_loop") + .add_node(node::PLAN, { + let rt = rt.clone(); + move |loop_state: LoopState, _ctx: NodeContext| { + let rt = rt.clone(); + async move { + let harness: Arc> = rt.harness.clone(); + let mut ctx_guard = rt.ctx.lock().await; + runtime::plan_node(&harness, &mut ctx_guard, loop_state).await + } + } + }) + .add_node(node::MODEL, { + let rt = rt.clone(); + move |loop_state: LoopState, _ctx: NodeContext| { + let rt = rt.clone(); + async move { + let harness = rt.harness.clone(); + let app_state = rt.app_state.clone(); + let mut ctx_guard = rt.ctx.lock().await; + let mut run_guard = rt.run.lock().await; + let mut status_guard = rt.status.lock().await; + runtime::model_node( + &harness, + &app_state, + &mut ctx_guard, + &mut run_guard, + &mut status_guard, + loop_state, + ) + .await + } + } + }) + .add_node(node::TOOLS, { + let rt = rt.clone(); + move |loop_state: LoopState, _ctx: NodeContext| { + let rt = rt.clone(); + async move { + let harness = rt.harness.clone(); + let app_state = rt.app_state.clone(); + let mut ctx_guard = rt.ctx.lock().await; + let mut run_guard = rt.run.lock().await; + let mut status_guard = rt.status.lock().await; + runtime::tools_node( + &harness, + &app_state, + &mut ctx_guard, + &mut run_guard, + &mut status_guard, + loop_state, + ) + .await + } + } + }) + .add_node(node::SETTLE, { + let rt = rt.clone(); + move |loop_state: LoopState, _ctx: NodeContext| { + let rt = rt.clone(); + async move { + let harness = rt.harness.clone(); + let mut run_guard = rt.run.lock().await; + runtime::settle_node(&harness, &mut run_guard, loop_state).await + } + } + }) + .set_entry(node::PLAN) + .mark_command_routing(node::PLAN) + .mark_command_routing(node::MODEL) + .mark_command_routing(node::TOOLS) + .mark_command_routing(node::SETTLE); + + // `plan`/`model`/`tools` each already return a real, node-emitted + // `NodeResult::Interrupt` when they need to pause (a steering pause from + // `plan`, a `MiddlewareControl::Interrupt` from any of the three — see + // the module doc above); that alone is a genuine, checkpointed executor + // pause, with no help from `GraphBuilder::mark_interrupt` needed. + // `mark_interrupt` is *not* used here because — unlike when this graph + // was first written — it is no longer a behavior-free export marker: it + // now aliases `GraphBuilder::interrupt_before`, which would make the + // executor pause *every* activation of these nodes on its own, before + // the node (and its middleware) ever runs, double-pausing on top of the + // node's own interrupt and desyncing resume's interrupt-acknowledgement + // bookkeeping across turns. Setting the `NodeMeta` interrupt marker + // directly (`pub(crate)`, reachable from this sibling module) restores + // the original export-only annotation without opting into that runtime + // pause. + for node in [node::PLAN, node::MODEL, node::TOOLS] { + builder.node_meta.entry(node.into()).or_default().interrupt = true; + } + + builder.compile() +} diff --git a/crates/tinyagents-graph/src/agent_loop/driver.rs b/crates/tinyagents-graph/src/agent_loop/driver.rs new file mode 100644 index 00000000..90e47e20 --- /dev/null +++ b/crates/tinyagents-graph/src/agent_loop/driver.rs @@ -0,0 +1,246 @@ +//! [`GraphLoopDriver`]: plugs the compiled loop's node bodies into +//! [`AgentHarness::invoke`] (and friends) via +//! [`tinyagents_harness::agent_loop::phases::LoopDriver`] + +//! [`AgentHarness::with_loop_driver`], selected by +//! [`tinyagents_harness::runtime::RunPolicy::execution`]`::Graph`. +//! +//! See the module doc on [`super`] ("`GraphLoopDriver` vs. +//! `compile_loop`/`LoopIter`") for why this drives +//! [`super::runtime::plan_node`]/`model_node`/`tools_node`/`settle_node` +//! directly over borrowed `&mut` state in a hand-rolled loop instead of +//! building a [`crate::CompiledGraph`]. + +use async_trait::async_trait; + +use tinyagents_harness::agent_loop::phases::LoopDriver; +use tinyagents_harness::context::RunContext; +use tinyagents_harness::error::{Result, TinyAgentsError}; +use tinyagents_harness::events::{AgentEvent, HarnessRunStatus}; +use tinyagents_harness::ids::HarnessPhase; +use tinyagents_harness::middleware::AgentRun; +use tinyagents_harness::runtime::AgentHarness; +use tinyagents_harness::steering::PauseState; +use tinyinference_llm::message::Message; + +use crate::command::{NodeResult, RouteTarget}; + +use super::runtime; +use super::types::{LoopState, node}; + +/// Drives [`AgentHarness::invoke`] through the same node bodies +/// [`super::compile_loop`] wires into a [`crate::CompiledGraph`], without +/// itself building one. Install with +/// [`AgentHarness::with_loop_driver`]`(Arc::new(GraphLoopDriver::new()))` and +/// [`tinyagents_harness::runtime::RunPolicy::execution`]`::Graph`. +/// +/// Stateless — one instance can be shared (via `Arc`) across every harness +/// that wants the graph engine. +#[derive(Debug, Default, Clone, Copy)] +pub struct GraphLoopDriver; + +impl GraphLoopDriver { + /// Creates a driver. Stateless: nothing to configure. + pub fn new() -> Self { + Self + } +} + +#[async_trait] +impl LoopDriver for GraphLoopDriver +where + State: Send + Sync, + Ctx: Send + Sync, +{ + async fn drive( + &self, + harness: &AgentHarness, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + input: Vec, + streaming: bool, + ) -> Result<()> { + // Mirrors `run_loop`'s own top-of-run bookkeeping (see that + // function's docs on why the limit tracker restarts here rather + // than at `RunContext::new`). + ctx.limits.restart(); + runtime::reconcile_call_limits(ctx, harness.policy()); + ctx.streaming = streaming; + + let record = ctx.emit(AgentEvent::RunStarted { + run_id: ctx.run_id().clone(), + thread_id: ctx.thread_id().cloned(), + }); + status.set_last_event(record.id); + status.mark_running(HarnessPhase::Idle); + + harness.middleware().run_before_agent(ctx, state).await?; + + let mut loop_state = LoopState { + messages: input, + ..LoopState::default() + }; + let mut current: &str = node::PLAN; + + let outcome = loop { + // Keeps `run.messages` a running snapshot of the transcript as + // of the start of each node call, so a node that errors or + // interrupts mid-call (and so never hands `loop_state` back) + // still leaves `run.messages` close to current — mirroring + // `run_loop`'s "transcript survives every exit path" guarantee + // as closely as this driver's by-value node contract allows. + // The one gap: mutations a node makes to its *own* copy of + // `loop_state.messages` before erroring/interrupting (for + // example `plan_node`'s steering-injected message on a pause) + // are not reflected until that node returns successfully. + run.messages = loop_state.messages.clone(); + let result = match current { + node::PLAN => runtime::plan_node(harness, ctx, loop_state).await, + node::MODEL => { + runtime::model_node(harness, state, ctx, run, status, loop_state).await + } + node::TOOLS => { + runtime::tools_node(harness, state, ctx, run, status, loop_state).await + } + node::SETTLE => runtime::settle_node(harness, run, loop_state).await, + other => { + break Err(TinyAgentsError::Validation(format!( + "GraphLoopDriver: unknown loop node `{other}`" + ))); + } + }; + + match result { + Ok(NodeResult::Interrupt(interrupt)) + if interrupt.id.ends_with("-steering-pause") => + { + // No `CompiledGraph` is in play here, so there is no + // checkpoint to pause against — mirror the direct loop's + // steering pause instead: latch `run.paused` and finish + // this call with `Ok(())`, exactly like + // `run_loop`'s `LoopExit::Paused` handling. See the + // module doc on `super` for why this is not a resumable + // graph interrupt. + break Ok(Some(interrupt)); + } + Ok(NodeResult::Interrupt(interrupt)) => { + // A `MiddlewareControl::Interrupt` (an approval gate, not + // a steering pause). The direct loop surfaces this as + // `TinyAgentsError::Interrupted` (see + // `agent_loop::run_loop`'s `apply_pending_control`), not + // a pause — matched here so `RunPolicy::execution == + // Graph` behaves identically to `Direct` for this + // control. Only `compile_loop`/`LoopIter`'s real + // `CompiledGraph` upgrades this into a resumable graph + // interrupt (A5's "approvals surfacing as graph + // interrupts"). + let message = interrupt + .payload + .get("message") + .and_then(|value| value.as_str()) + .unwrap_or("interrupted") + .to_string(); + break Err(TinyAgentsError::Interrupted { + node: interrupt.node.to_string(), + message, + }); + } + Ok(NodeResult::Update(updated)) => { + // Every node body returns `Command`/`Interrupt` (see + // `runtime`'s node docs); a bare `Update` would mean the + // loop cannot determine where to go next. + let _ = updated; + break Err(TinyAgentsError::Validation( + "GraphLoopDriver: loop node returned an un-routed update".to_string(), + )); + } + Ok(NodeResult::Command(command)) => { + loop_state = match command.update { + Some(update) => update, + None => { + break Err(TinyAgentsError::Validation( + "GraphLoopDriver: loop node's command carried no update" + .to_string(), + )); + } + }; + let Some(target) = command.goto.first() else { + break Err(TinyAgentsError::Validation( + "GraphLoopDriver: loop node's command carried no route".to_string(), + )); + }; + let RouteTarget::Node(node_id) = target else { + break Err(TinyAgentsError::Validation( + "GraphLoopDriver: loop node routed via `Send`, which this driver \ + does not support" + .to_string(), + )); + }; + if node_id.as_str() == crate::builder::END { + break Ok(None); + } + current = match node_id.as_str() { + node::PLAN => node::PLAN, + node::MODEL => node::MODEL, + node::TOOLS => node::TOOLS, + node::SETTLE => node::SETTLE, + other => { + break Err(TinyAgentsError::Validation(format!( + "GraphLoopDriver: unknown loop node `{other}`" + ))); + } + }; + continue; + } + Err(error) => break Err(error), + } + }; + + status.mark_running(HarnessPhase::Middleware); + harness + .middleware() + .run_after_agent(ctx, state, run) + .await?; + + // `status.mark_completed`/`mark_interrupted`/`mark_failed` and (on + // error) `AgentEvent::RunFailed` are applied centrally by + // `agent_loop::entry::drive_collecting` after this call returns, + // identically for the direct loop and this driver — see that + // function's doc comment. This driver only emits the terminal event + // that (like the direct loop's `run_loop_body`) is its own + // responsibility to raise: `RunCompleted` on a clean finish, or + // latching `run.paused` (mirroring a steering pause) on an + // interrupt. + match outcome { + Ok(None) => { + let record = ctx.emit(AgentEvent::RunCompleted { + run_id: ctx.run_id().clone(), + }); + status.set_last_event(record.id); + Ok(()) + } + Ok(Some(interrupt)) => { + let reason = interrupt + .payload + .get("reason") + .or_else(|| interrupt.payload.get("message")) + .and_then(|value| value.as_str()) + .map(str::to_string); + let record = ctx.emit(AgentEvent::ControlApplied { + control: "paused".to_string(), + detail: reason + .clone() + .unwrap_or_else(|| format!("paused at node `{}`", interrupt.node)), + }); + status.set_last_event(record.id); + run.paused = Some(PauseState { + reason, + paused_at_checkpoint: 0, + }); + Ok(()) + } + Err(error) => Err(error), + } + } +} diff --git a/crates/tinyagents-graph/src/agent_loop/iter.rs b/crates/tinyagents-graph/src/agent_loop/iter.rs new file mode 100644 index 00000000..9e374d1a --- /dev/null +++ b/crates/tinyagents-graph/src/agent_loop/iter.rs @@ -0,0 +1,285 @@ +//! [`AgentLoopGraphExt::iter`]/[`LoopIter`]: steps the compiled loop one node +//! at a time. +//! +//! See the module doc on [`super`] for how this relates to [`super::compile_loop`] +//! and [`super::GraphLoopDriver`]. + +use std::sync::Arc; + +use tinyagents_harness::context::RunContext; +use tinyagents_harness::error::{Result, TinyAgentsError}; +use tinyagents_harness::events::HarnessRunStatus; +use tinyagents_harness::ids::ComponentId; +use tinyagents_harness::middleware::AgentRun; +use tinyagents_harness::runtime::AgentHarness; +use tinyinference_llm::message::Message; + +use crate::command::NodeResult; + +use super::runtime::{self, LoopRuntime}; +use super::types::{LoopState, node}; + +/// One completed activation reported by [`LoopIter::next`]. +#[derive(Clone, Debug)] +pub struct LoopStep { + /// The node that just ran. + pub node: String, + /// The node [`LoopIter::next`] will run next, unless + /// [`LoopIter::override_next`] changes it first. `None` once the run has + /// finished (reached `END`). + pub next: Option, + /// Whether this activation interrupted the run (a steering pause or a + /// [`tinyagents_harness::context::MiddlewareControl::Interrupt`]) rather + /// than completing normally. A caller that wants to resume calls + /// [`LoopIter::next`] again — the interrupted node re-runs, exactly like + /// [`crate::compiled::CompiledGraph::resume`] re-running an interrupted + /// node — after applying whatever unblocks it (for example draining a + /// [`tinyagents_harness::steering::SteeringHandle`]). + pub interrupted: bool, +} + +/// Extension trait adding [`Self::iter`] to `Arc<`[`AgentHarness`]`>`, mirroring pydantic-ai's `Agent.iter` (`docs/runtime-comparison/pydantic-ai.md` +/// §4). +/// +/// Implemented for `Arc>` rather than +/// `AgentHarness` directly because the returned [`LoopIter`] +/// keeps driving the loop across many `.await` points spanning its own +/// lifetime (unlike a single [`AgentHarness::invoke`] call), so it needs to +/// own a durable handle to the harness — see the module doc on [`super`] +/// ("`GraphLoopDriver` vs. `compile_loop`/`LoopIter`") for why that rules out +/// a borrowed `&AgentHarness` the way [`super::GraphLoopDriver`] gets one. +pub trait AgentLoopGraphExt { + /// Starts a steppable run: builds the compiled `plan -> model -> tools -> + /// settle` graph (see [`super::compile_loop`]) and returns a + /// [`LoopIter`] positioned at its entry (`plan`), seeded with `input` as + /// the starting transcript. `app_state` is the harness's shared, + /// read-only application state (the same value an + /// [`AgentHarness::invoke`] caller would pass as `state`); `ctx` supplies + /// the run identity and dependencies exactly as for `invoke`. + fn iter( + self, + app_state: Arc, + ctx: RunContext, + input: Vec, + ) -> Result>; +} + +impl AgentLoopGraphExt for Arc> +where + State: Send + Sync + 'static, + Ctx: Send + Sync + 'static, +{ + fn iter( + self, + app_state: Arc, + ctx: RunContext, + input: Vec, + ) -> Result> { + LoopIter::new(self, app_state, ctx, input) + } +} + +/// Steps the compiled agent loop one node activation at a time. +/// +/// Built via [`AgentLoopGraphExt::iter`]. Each [`Self::next`] call runs +/// exactly one node body (the same [`super::runtime::plan_node`]/`model_node`/ +/// `tools_node`/`settle_node` [`super::compile_loop`] wires into a +/// [`crate::CompiledGraph`]) against this iterator's own [`LoopRuntime`], so +/// stepping through a `LoopIter` observes the identical transcript, usage, +/// and routing decisions a full [`crate::CompiledGraph::run`] over the same +/// graph would — just one activation at a time, with [`Self::override_next`] +/// able to redirect the very next activation (for tests, debugging, or a +/// host that wants to splice in extra bookkeeping between turns). +pub struct LoopIter { + rt: Arc>, + state: LoopState, + next: Option, + overridden: Option, + finished: bool, +} + +impl LoopIter +where + State: Send + Sync + 'static, + Ctx: Send + Sync + 'static, +{ + pub(crate) fn new( + harness: Arc>, + app_state: Arc, + ctx: RunContext, + input: Vec, + ) -> Result { + let run_id = ctx.run_id().clone(); + let status = HarnessRunStatus::new(run_id, ComponentId::new("agent_loop.iter")); + let rt = Arc::new(LoopRuntime::new( + harness, + app_state, + ctx, + AgentRun::default(), + status, + false, + )); + Ok(Self { + rt, + state: LoopState { + messages: input, + ..LoopState::default() + }, + next: Some(node::PLAN.to_string()), + overridden: None, + finished: false, + }) + } + + /// The committed [`LoopState`] as of the last completed [`Self::next`] + /// call (or the seeded input, before the first call). + pub fn state(&self) -> &LoopState { + &self.state + } + + /// The accumulated [`AgentRun`] mirrored alongside the loop state — + /// useful for reading `usage`/`model_calls`/`tool_calls` without waiting + /// for `settle` to populate [`LoopState::final_text`]. + pub async fn run(&self) -> AgentRun { + self.rt.run.lock().await.clone() + } + + /// The node [`Self::next`] will run on its next call, or `None` once the + /// run has finished. + pub fn next_node(&self) -> Option<&str> { + self.next.as_deref() + } + + /// Redirects the very next [`Self::next`] call to `node` instead of + /// wherever the last activation routed to. Consumed by that one call — + /// subsequent calls follow the graph's normal routing again unless + /// overridden again. A no-op once the run has finished. + pub fn override_next(&mut self, node: impl Into) { + if !self.finished { + self.overridden = Some(node.into()); + } + } + + /// Runs exactly one node activation and reports what happened, or `Ok(None)` + /// if the run had already finished. + pub async fn next(&mut self) -> Result> { + if self.finished { + return Ok(None); + } + let current = self + .overridden + .take() + .or_else(|| self.next.clone()) + .ok_or_else(|| { + TinyAgentsError::Validation("LoopIter::next: no node to run next".to_string()) + })?; + + let loop_state = std::mem::take(&mut self.state); + let harness = self.rt.harness.clone(); + let app_state = self.rt.app_state.clone(); + let mut ctx_guard = self.rt.ctx.lock().await; + let mut run_guard = self.rt.run.lock().await; + let mut status_guard = self.rt.status.lock().await; + + let result = match current.as_str() { + node::PLAN => runtime::plan_node(&harness, &mut ctx_guard, loop_state).await?, + node::MODEL => { + runtime::model_node( + &harness, + &app_state, + &mut ctx_guard, + &mut run_guard, + &mut status_guard, + loop_state, + ) + .await? + } + node::TOOLS => { + runtime::tools_node( + &harness, + &app_state, + &mut ctx_guard, + &mut run_guard, + &mut status_guard, + loop_state, + ) + .await? + } + node::SETTLE => runtime::settle_node(&harness, &mut run_guard, loop_state).await?, + other => { + return Err(TinyAgentsError::Validation(format!( + "LoopIter::next: unknown loop node `{other}`" + ))); + } + }; + drop((ctx_guard, run_guard, status_guard)); + + match result { + NodeResult::Interrupt(_interrupt) => Ok(Some(LoopStep { + node: current.clone(), + // The interrupted node is the natural resume target: calling + // `next()` again re-runs it, mirroring + // `CompiledGraph::resume`. + next: Some(current), + interrupted: true, + })), + NodeResult::Command(command) => { + self.state = command.update.ok_or_else(|| { + TinyAgentsError::Validation( + "LoopIter::next: loop node's command carried no update".to_string(), + ) + })?; + let target = command.goto.first().ok_or_else(|| { + TinyAgentsError::Validation( + "LoopIter::next: loop node's command carried no route".to_string(), + ) + })?; + let crate::command::RouteTarget::Node(node_id) = target else { + return Err(TinyAgentsError::Validation( + "LoopIter::next: loop node routed via `Send`, which `LoopIter` does not \ + support" + .to_string(), + )); + }; + let next_str = node_id.as_str().to_string(); + if next_str == crate::builder::END { + self.finished = true; + self.next = None; + Ok(Some(LoopStep { + node: current, + next: None, + interrupted: false, + })) + } else { + self.next = Some(next_str.clone()); + Ok(Some(LoopStep { + node: current, + next: Some(next_str), + interrupted: false, + })) + } + } + NodeResult::Update(_) => Err(TinyAgentsError::Validation( + "LoopIter::next: loop node returned an un-routed update".to_string(), + )), + } + } + + /// Steps to completion (or the first unresolved interrupt), returning the + /// final [`LoopState`]. + /// + /// Stops — without erroring — the moment [`Self::next`] reports an + /// interrupted step, leaving [`Self::next_node`] pointing at the + /// interrupted node so a caller can resolve whatever paused it (drain + /// steering, record an approval) and call [`Self::run_to_end`] again to + /// continue. + pub async fn run_to_end(&mut self) -> Result<&LoopState> { + while let Some(step) = self.next().await? { + if step.interrupted { + break; + } + } + Ok(&self.state) + } +} diff --git a/crates/tinyagents-graph/src/agent_loop/mod.rs b/crates/tinyagents-graph/src/agent_loop/mod.rs new file mode 100644 index 00000000..9f1f8966 --- /dev/null +++ b/crates/tinyagents-graph/src/agent_loop/mod.rs @@ -0,0 +1,122 @@ +//! Compiled-graph rendition of the harness agent loop (A5). +//! +//! `docs/runtime-comparison/langgraph.md` §4 ("Agent loop as a graph") and +//! `docs/runtime-comparison/pydantic-ai.md` §4 (the `iter` API) both ask the +//! same question of tinyagents: can the default `plan -> model -> tools` +//! agent loop — normally a monolithic Rust function +//! (`tinyagents_harness::agent_loop::run_loop`) — also be expressed as an +//! ordinary [`crate::CompiledGraph`], so it gets the graph runtime's +//! checkpoint/resume, step-by-step `iter()`, and interrupt machinery for +//! free, instead of each of those being reinvented (or left unavailable) on +//! the harness's own loop? +//! +//! This module is "yes": [`compile_loop`] builds a real +//! `CompiledGraph` with four nodes +//! (`plan`/`model`/`tools`/`settle`, see [`types::node`]), and +//! [`AgentLoopGraphExt::iter`] drives it step-by-step through [`LoopIter`]. +//! [`GraphLoopDriver`] additionally plugs the same phase logic into +//! [`tinyagents_harness::runtime::AgentHarness::invoke`] (and friends) via +//! [`tinyagents_harness::agent_loop::phases::LoopDriver`] + +//! [`tinyagents_harness::runtime::AgentHarness::with_loop_driver`], selected +//! by [`tinyagents_harness::runtime::RunPolicy::execution`]`::Graph`. +//! +//! # Dependency direction +//! +//! `tinyagents-graph` depends on `tinyagents-harness`, never the reverse, so +//! this compiled-graph rendition lives here rather than in the harness +//! crate — the harness only exposes the seam +//! ([`tinyagents_harness::agent_loop::phases`]) this module plugs into. See +//! that module's doc comment for the harness-side half of the boundary. +//! +//! # Scope +//! +//! This is a **subset** of `run_loop`'s behavior, not a byte-for-byte +//! reimplementation. It covers the common path exercised by +//! `crates/tinyagents-integration-tests/tests/loop_as_graph.rs`: tool +//! calling, structured output (`ResponseFormat::Auto`/`JsonSchema`, provider- +//! schema and tool-call-fallback strategies), the output-validation retry +//! loop (A3), run limits, `MiddlewareControl` routing (`JumpTo`, +//! `StopWithFinal`, `Interrupt`, `UpdateState`), and steering +//! (cancel/pause/inject). It intentionally does **not** cover: host-model +//! routing (`HostCapabilities`), cross-provider handoff transforms +//! (`agent_loop::handoff_transform`), the deferred-tool discovery bridge, +//! truncated-empty-response recovery/retry, `RunPolicy::retry`/`fallback` +//! (a registered `ModelMiddleware` still runs, so a retry-on-error +//! middleware still applies, but the built-in retry/fallback loop the direct +//! model call performs does not), response caching, and +//! `StructuredStrategy::Prompted`/`ToolCallUnion` (A6's +//! `structured_strategy_override`) or `EndStrategy::Early`/`Exhaustive` (A6); +//! `resolve_structured_plan` also resolves the profile-driven `Auto` choice +//! against the *default* model binding rather than the turn's actually- +//! resolved model (see that function's docs). [`RunPolicy::execution`] +//! defaults to [`tinyagents_harness::runtime::LoopExecution::Direct`], so +//! every existing caller is unaffected unless it opts in. +//! +//! # Tool batch execution shape +//! +//! The `tools` node runs a turn's whole tool-call batch in one node +//! activation via [`tinyagents_harness::agent_loop::phases::execute_tool_batch`] +//! rather than one graph node (or one `Send` fan-out branch) per tool call. +//! This was the deliberate choice over a per-call fan-out: the direct loop's +//! serial-admission / serial-or-concurrent-dispatch decision (see +//! `tinyagents_harness::agent_loop::tools`'s module docs) is genuinely +//! call-count- and middleware-dependent, and re-deriving it as graph +//! topology would either have to duplicate that decision as a router (two +//! sources of truth to keep in sync) or lose the exact ordering/budget +//! guarantees the direct loop promises. Reusing the harness function as-is +//! keeps the tool batch's ordering, concurrency-eligibility, and budget/limit +//! semantics identical to the direct loop by construction, at the cost of a +//! coarser graph (a `tools` node activation is opaque to the executor's own +//! per-task checkpointing — a batch is atomic from the graph's point of +//! view, resuming re-runs the whole batch rather than only its unfinished +//! calls). +//! +//! # `GraphLoopDriver` vs. `compile_loop`/`LoopIter` +//! +//! These are two different integration points sharing the same node bodies +//! (`runtime::plan_node`/`model_node`/`tools_node`/`settle_node`), not one +//! layered on the other: +//! +//! - [`compile_loop`]/[`LoopIter`] build and step a real +//! `CompiledGraph` over an owned [`runtime::LoopRuntime`] +//! (`Arc`'d harness/state, owned `RunContext`) — this is the path that gets +//! checkpointing, `resume`, and `iter()`. +//! - [`GraphLoopDriver`] is handed only a transient `&mut RunContext`/`&mut +//! AgentRun`/`&mut HarnessRunStatus` bundle by +//! [`tinyagents_harness::agent_loop::phases::LoopDriver::drive`] (that +//! trait's signature mirrors `run_loop`'s own borrowed-state contract) and +//! an `&AgentHarness`/`&State` it does not own. Building a `LoopRuntime` +//! (which needs `Arc>` to satisfy the `'static` +//! bound `GraphBuilder::add_node`'s closures require) from a bare `&AgentHarness` +//! is not possible in safe Rust without either forcing every caller of +//! `with_loop_driver` to already hold an `Arc` (which would +//! make installing a driver on a not-yet-`Arc`'d harness impossible — a +//! real usability regression) or unsafely extending the borrow's lifetime +//! (which this workspace denies via `unsafe_code = "deny"`). So +//! [`GraphLoopDriver::drive`] instead calls the exact same node bodies +//! directly, in a hand-rolled loop, against the real borrowed `&mut` +//! state — no `Arc`, no `Mutex`, no `CompiledGraph` involved — which is +//! sound with zero unsafe code precisely because a borrowed async call +//! needs no `'static` bound. The two paths are behaviorally identical +//! (same node functions) but the `GraphLoopDriver` path does not get graph +//! checkpointing: an interrupt it raises still pauses the run (mirroring +//! `run.paused`, exactly like a steering pause on the direct loop) but is +//! not itself a resumable `CompiledGraph` checkpoint — resume it the same +//! way a paused direct-loop run is resumed (feed `run.messages` back in as +//! the next call's `input`), not via `CompiledGraph::resume`. A caller that +//! wants graph-level checkpoint/resume across the interrupt should drive +//! the loop through [`compile_loop`]/[`LoopIter`] directly instead of +//! through `AgentHarness::invoke`. + +pub mod driver; +pub mod iter; +pub mod runtime; +pub mod types; + +mod compile; + +pub use compile::compile_loop; +pub use driver::GraphLoopDriver; +pub use iter::{AgentLoopGraphExt, LoopIter, LoopStep}; +pub use runtime::LoopRuntime; +pub use types::{LoopState, LoopUpdate, node}; diff --git a/crates/tinyagents-graph/src/agent_loop/runtime.rs b/crates/tinyagents-graph/src/agent_loop/runtime.rs new file mode 100644 index 00000000..09083e18 --- /dev/null +++ b/crates/tinyagents-graph/src/agent_loop/runtime.rs @@ -0,0 +1,683 @@ +//! Shared per-run state and the four node bodies (`plan`, `model`, `tools`, +//! `settle`) that [`super::compile_loop`] wires into a graph, and that +//! [`super::LoopIter`] steps through directly. +//! +//! See the module doc on [`super`] for the loop's documented scope. + +use std::sync::Arc; + +use tokio::sync::Mutex; + +use tinyagents_harness::agent_loop::phases; +use tinyagents_harness::context::{LoopTarget, MiddlewareControl, RunContext}; +use tinyagents_harness::error::{Result, TinyAgentsError}; +use tinyagents_harness::events::{AgentEvent, HarnessRunStatus}; +use tinyagents_harness::ids::{CallId, NodeId}; +use tinyagents_harness::middleware::{AgentRun, BoxModelFuture, ModelBaseCall}; +use tinyagents_harness::runtime::AgentHarness; +use tinyagents_harness::steering::{SteeringOutcome, apply_pending_steering}; +use tinyagents_harness::structured::{StructuredExtractor, StructuredStrategy}; + +use crate::command::Interrupt; +use crate::{Command, NodeResult, RouteTarget}; + +use super::types::{LoopState, PendingStructuredPlan, node}; + +use tinyinference_llm::model::{ModelRequest, ModelResponse, ResponseFormat, ToolChoice}; +use tinyinference_llm::tool::{ToolFormat, ToolSchema}; + +/// Per-run state shared by every node closure [`super::compile_loop`] +/// builds, and by [`super::LoopIter`]. +/// +/// Captured behind an `Arc` (with interior mutability for the pieces that +/// need `&mut` access) because [`crate::GraphBuilder::add_node`] handlers are +/// `Fn`, not `FnMut`: the graph executor may in principle invoke a node +/// concurrently with itself across forked branches, so every mutable piece +/// here is guarded by its own [`Mutex`]. In this loop's own topology no two +/// nodes ever run concurrently (it is a strictly sequential chain), so the +/// locks are never contended — they exist to satisfy `Send + Sync + 'static` +/// and the `Fn` bound, not to arbitrate real concurrency. +pub struct LoopRuntime { + pub(crate) harness: Arc>, + pub(crate) app_state: Arc, + pub(crate) ctx: Mutex>, + pub(crate) run: Mutex, + pub(crate) status: Mutex, + /// Not yet consulted: `model_node`'s `DirectModelBase` always dispatches + /// through `ChatModel::invoke`, not `ChatModel::stream` (see the module + /// doc on `super` — streaming is out of scope for this rendition of the + /// loop). Kept so `LoopRuntime::new`'s signature already matches what a + /// future streaming node would need. + #[allow(dead_code)] + pub(crate) streaming: bool, +} + +impl LoopRuntime { + /// Builds a fresh, owned [`LoopRuntime`] for one run: [`super::LoopIter`] + /// (which owns `ctx`/`input` for the run's whole lifetime) is the + /// intended caller. [`super::GraphLoopDriver`] does **not** use this — + /// see its module doc for why it drives the same node bodies directly + /// over borrowed `&mut` state instead of through an owned + /// `LoopRuntime`/`CompiledGraph`. + pub fn new( + harness: Arc>, + app_state: Arc, + mut ctx: RunContext, + run: AgentRun, + status: HarnessRunStatus, + streaming: bool, + ) -> Self { + ctx.limits.restart(); + reconcile_call_limits(&mut ctx, harness.policy()); + Self { + harness, + app_state, + ctx: Mutex::new(ctx), + run: Mutex::new(run), + status: Mutex::new(status), + streaming, + } + } + + /// [`Self::new`] with a fresh, default [`AgentRun`]/[`HarnessRunStatus`] + /// — the common case for starting a brand-new run (as opposed to + /// resuming one, which would seed `run`/`status` from prior state). + pub fn for_run( + harness: Arc>, + app_state: Arc, + ctx: RunContext, + ) -> Self { + let run_id = ctx.run_id().clone(); + let status = HarnessRunStatus::new( + run_id, + tinyagents_harness::ids::ComponentId::new("agent_loop"), + ); + Self::new(harness, app_state, ctx, AgentRun::default(), status, false) + } +} + +/// The innermost model call: a direct, single-attempt dispatch to the +/// resolved [`tinyinference_llm::model::ChatModel`]. +/// +/// Unlike the direct loop's `ModelCallBase` (private to +/// `tinyagents-harness::agent_loop`), this has no response-cache lookup, no +/// `RunPolicy::retry`/`RunPolicy::fallback` loop, and no host-model routing — +/// see the module doc on [`super`] for the full list of scoped-out behavior. +/// It still runs through +/// [`tinyagents_harness::middleware::MiddlewareStack::run_wrapped_model`], so +/// a registered [`tinyagents_harness::middleware::ModelMiddleware`] (for +/// example a retry-on-error wrap middleware) still applies. +struct DirectModelBase<'m, State: Send + Sync> { + model: &'m dyn tinyinference_llm::model::ChatModel, +} + +impl ModelBaseCall + for DirectModelBase<'_, State> +{ + fn call<'a>( + &'a self, + _ctx: &'a mut RunContext, + state: &'a State, + request: ModelRequest, + ) -> BoxModelFuture<'a> { + Box::pin(async move { + self.model + .invoke(state, request) + .await + .map_err(TinyAgentsError::from) + }) + } +} + +/// Resolves the structured-output plan for `response_format`, mirroring the +/// direct loop's `Auto`/`JsonSchema` resolution but without the +/// `Prompted`/`ToolCallUnion` overrides (see the module doc on [`super`]). +fn resolve_structured_plan( + request: &mut ModelRequest, + profile: Option<&tinyinference_llm::model::ModelProfile>, +) -> Option { + match request.response_format.take() { + Some(ResponseFormat::Auto { name, schema }) => { + let strategy = StructuredStrategy::for_profile(profile); + match strategy { + StructuredStrategy::ProviderSchema => { + request.response_format = + Some(ResponseFormat::json_schema(name.clone(), schema.clone())); + } + StructuredStrategy::ToolCall => { + let fallback_schema = ToolSchema { + name: name.clone(), + description: format!("Return the result as `{name}`."), + parameters: schema.clone(), + format: ToolFormat::Json, + }; + request.tools.push(fallback_schema); + if request.tools.len() == 1 { + request.tool_choice = ToolChoice::Tool(name.clone()); + } + } + StructuredStrategy::Prompted { .. } | StructuredStrategy::ToolCallUnion => { + unreachable!("StructuredStrategy::for_profile never returns these") + } + } + Some(PendingStructuredPlan { + strategy, + schema_name: name, + schema, + }) + } + Some(ResponseFormat::JsonSchema { name, schema }) => { + request.response_format = + Some(ResponseFormat::json_schema(name.clone(), schema.clone())); + Some(PendingStructuredPlan { + strategy: StructuredStrategy::ProviderSchema, + schema_name: name, + schema, + }) + } + other => { + request.response_format = other; + None + } + } +} + +/// The `plan` node body: builds the next [`ModelRequest`] from the working +/// transcript, the harness's registered tools, and the policy's response +/// format, and stashes it on [`LoopState::pending_request`]. +/// +/// Takes `harness`/`ctx` as plain borrows rather than the `Arc` +/// the other node bodies (which also need `run`/`status`) use, so this exact +/// function serves two callers with different ownership shapes without +/// duplicating its logic: [`super::compile`]'s graph closures call it against +/// a locked `MutexGuard` inside an owned, `Arc`'d [`LoopRuntime`], and +/// [`super::driver::GraphLoopDriver`] calls it directly against the +/// short-lived `&mut RunContext` [`tinyagents_harness::agent_loop::phases::LoopDriver::drive`] +/// is handed — see that module's doc for why it cannot build a `LoopRuntime` +/// of its own. +pub(crate) async fn plan_node( + harness: &AgentHarness, + ctx: &mut RunContext, + mut loop_state: LoopState, +) -> Result> +where + State: Send + Sync, + Ctx: Send + Sync, +{ + if ctx.cancellation.is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + match apply_pending_steering(ctx, &mut loop_state.messages)? { + SteeringOutcome::Cancel => return Err(TinyAgentsError::Cancelled), + SteeringOutcome::Pause => { + return Ok(NodeResult::Interrupt(Interrupt { + id: format!("{}-steering-pause", ctx.run_id()), + node: NodeId::from(node::PLAN), + payload: serde_json::json!({ "reason": "steering paused the run" }), + task_id: None, + response_schema: None, + })); + } + SteeringOutcome::Continue => {} + } + if ctx.check_deadline().is_err() { + return Err(TinyAgentsError::Timeout(format!( + "run `{}` exceeded its wall-clock deadline", + ctx.run_id() + ))); + } + + let tool_schemas = harness.tools().schemas(); + // Mirrors `run_loop_body`'s `AgentEvent::ToolsAdvertised`, emitted once + // per run there (right after `before_agent`) versus once per `plan` + // activation here — a documented, harmless divergence (see the module + // doc on `super`): the tool set does not change mid-run, so repeating + // the event on every turn only adds extra `tool.advertised` events, it + // never drops or reorders the one the direct loop's own listeners + // expect. + let advertised_record = ctx.emit(AgentEvent::ToolsAdvertised { + direct: tool_schemas.len(), + deferred: 0, + schema_bytes: tinyagents_harness::token_estimation::tool_schema_bytes(&tool_schemas), + }); + let _ = advertised_record; + let mut request = ModelRequest { + messages: loop_state.messages.clone(), + tools: tool_schemas, + ..ModelRequest::default() + }; + if let Some(format) = &harness.policy().default_response_format { + request.response_format = Some(format.clone()); + } + + // The structured plan depends on the resolved model's profile, but the + // model is not resolved until the `model` node (mirroring the direct + // loop's ordering). Resolve against the *default* binding here as a + // reasonable approximation for `ResponseFormat::Auto`'s profile-based + // strategy choice; an explicit `ResponseFormat::JsonSchema` is + // unaffected either way. This is a documented simplification relative to + // the direct loop, which resolves the model first. + let profile = harness + .models() + .resolve_request(&request, None, None) + .and_then(|binding| binding.model.profile().cloned()); + let structured = resolve_structured_plan(&mut request, profile.as_ref()); + + loop_state.pending_request = Some(request); + loop_state.pending_structured = structured; + Ok(goto(loop_state, node::MODEL)) +} + +/// The `model` node body: dispatches the request [`plan_node`] built, +/// records usage, appends the assistant message, and routes to `tools` or +/// `settle`. +pub(crate) async fn model_node( + harness: &AgentHarness, + app_state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + mut loop_state: LoopState, +) -> Result> +where + State: Send + Sync, + Ctx: Send + Sync, +{ + // `ctx.record_model_call()` itself raises a bare `Validation` error on a + // cap hit; `run_loop_body` wraps that into `LimitExceeded` (and honors + // `LimitBehavior::StopWithPartial` by finishing cleanly instead of + // erroring) — mirrored here so a caller sees the identical outcome + // regardless of which engine is driving the run. + if let Err(error) = ctx.record_model_call() { + let record = ctx.emit(AgentEvent::LimitReached { + kind: tinyagents_harness::events::LimitKind::ModelCalls, + }); + status.set_last_event(record.id); + if matches!( + harness.policy().limits.behavior, + tinyagents_harness::limits::LimitBehavior::StopWithPartial + ) { + loop_state.finished = true; + if loop_state.final_text.is_none() { + loop_state.final_text = Some(last_assistant_text(&loop_state.messages)); + } + return Ok(goto(loop_state, node::SETTLE)); + } + return Err(TinyAgentsError::LimitExceeded(error.to_string())); + } + + let request = loop_state + .pending_request + .take() + .ok_or_else(|| TinyAgentsError::Validation("model node ran with no pending plan".into()))?; + + let binding = harness + .models() + .resolve_request(&request, None, None) + .ok_or_else(|| { + TinyAgentsError::ModelNotFound( + request.model.clone().unwrap_or_else(|| "".into()), + ) + })?; + let model_name = binding.resolved.name.clone(); + let call_id = CallId::new(format!("{}-model-{}", ctx.run_id(), run.model_calls + 1)); + + let mut request = request; + harness + .middleware() + .run_before_model(ctx, app_state, &mut request) + .await?; + + let started_record = ctx.emit(AgentEvent::ModelStarted { + call_id: call_id.clone(), + model: model_name.clone(), + }); + status.set_last_event(started_record.id); + + let base = DirectModelBase { + model: binding.model.as_ref(), + }; + let (mut response, wrap_control) = harness + .middleware() + .run_wrapped_model(ctx, app_state, request, &base) + .await? + .into_response_with_control(); + if let Some(control) = wrap_control { + ctx.request_control(control); + } + + run.model_calls += 1; + run.steps += 1; + status.model_calls = run.model_calls; + if let Some(usage) = response.usage { + run.usage.record(usage); + loop_state.usage = run.usage; + let usage_record = ctx.emit(AgentEvent::UsageRecorded { usage }); + status.set_last_event(usage_record.id); + } + + harness + .middleware() + .run_after_model(ctx, app_state, &mut response) + .await?; + + let completed_record = ctx.emit(AgentEvent::ModelCompleted { + call_id: call_id.clone(), + started_at_ms: None, + usage: response.usage, + input: None, + output: None, + }); + status.set_last_event(completed_record.id); + + loop_state.model_calls = run.model_calls; + loop_state.last_call_id = Some(call_id.to_string()); + loop_state + .messages + .push(tinyinference_llm::message::Message::Assistant( + response.message.clone(), + )); + loop_state.turn += 1; + + let tool_calls = response.tool_calls().to_vec(); + loop_state.pending_tool_calls = tool_calls.clone(); + + let route = if tool_calls.is_empty() { + node::SETTLE + } else { + node::TOOLS + }; + + // Computed above `take_control` (rather than the reverse) so a + // `MiddlewareControl::Continue`/`UpdateState` control — which means "no + // override, proceed with whatever the turn would have done anyway" — + // has the real tool-routing decision to fall through to instead of an + // arbitrary default. + if let Some(control) = ctx.take_control() { + return apply_control(ctx, &mut loop_state, control, node::MODEL, route); + } + // Stash the response for `settle` to extract structured output from. + // Reusing `pending_request`'s sibling field would need a new field; keep + // it simple by re-deriving what `settle` needs from `messages` (the + // response text/tool-calls) plus `structured` plan already on + // `loop_state`. `response.finish_reason`/raw provider fields are not + // needed by this reduced-scope settle (see the module doc on `super`). + let _ = model_name; + let _ = ModelOutcomeShadow(&response); + Ok(goto(loop_state, route)) +} + +/// Zero-sized marker used only to keep `response` "used" for readability at +/// the call site above without over-cloning it into `LoopState`. +struct ModelOutcomeShadow<'a>(#[allow(dead_code)] &'a ModelResponse); + +/// The `tools` node body: executes the batch [`model_node`] requested via +/// [`phases::execute_tool_batch`] (the exact same admission / +/// serial-or-concurrent execution / middleware pipeline the direct loop +/// uses — see that function's docs), then routes back to `plan` for the next +/// turn. +pub(crate) async fn tools_node( + harness: &AgentHarness, + app_state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + mut loop_state: LoopState, +) -> Result> +where + State: Send + Sync, + Ctx: Send + Sync, +{ + let calls = std::mem::take(&mut loop_state.pending_tool_calls); + let outcome = phases::execute_tool_batch( + harness, + app_state, + ctx, + run, + status, + &mut loop_state.messages, + calls, + ) + .await?; + loop_state.tool_calls = run.tool_calls; + loop_state.executed_tools = run.executed_tools.clone(); + let _ = outcome; + + if harness.middleware().any_should_stop_after_turn(ctx, run) { + ctx.request_control(MiddlewareControl::JumpTo(LoopTarget::End)); + } + + if let Some(control) = ctx.take_control() { + return apply_control(ctx, &mut loop_state, control, node::TOOLS, node::PLAN); + } + + Ok(goto(loop_state, node::PLAN)) +} + +/// The `settle` node body: extracts/validates structured output when the +/// turn planned one, drives the output-validation retry loop +/// (`RunPolicy::output_retry`), and finishes the run. +pub(crate) async fn settle_node( + harness: &AgentHarness, + run: &mut AgentRun, + mut loop_state: LoopState, +) -> Result> +where + State: Send + Sync, + Ctx: Send + Sync, +{ + if let Some(plan) = loop_state.pending_structured.take() { + let extractor = StructuredExtractor::new( + plan.strategy.clone(), + &plan.schema_name, + plan.schema.clone(), + ); + let last_response = last_response_from_messages(&loop_state.messages); + let outcome = extractor.extract_outcome(&last_response); + let variant = outcome.variant.clone(); + // Note: unlike the direct loop, this does not consult + // `AgentHarness::with_output_validator` (A3's post-extraction + // validator hook) — out of scope for this rendition (see the module + // doc on `super`). + let error = match outcome.value { + Some(value) => { + run.structured = Some(value.clone()); + run.structured_variant = variant.clone(); + loop_state.structured = Some(value); + loop_state.structured_variant = variant; + None + } + None => outcome.error, + }; + if let Some(error) = error { + let max_attempts = harness.policy().output_retry.max_attempts; + if loop_state.output_retry_attempts < max_attempts { + loop_state.output_retry_attempts += 1; + let template = &harness.policy().output_retry.message_template; + let prompt = template.replace("{error}", &error); + loop_state + .messages + .push(tinyinference_llm::message::Message::user(prompt)); + // Back to `plan`, not `model` directly: the direct loop's + // retry re-enters its outer loop, which rebuilds the + // `ModelRequest` from `messages` (now including the repair + // prompt) — `model_node` needs a fresh `pending_request`, + // which only `plan_node` produces. + return Ok(goto(loop_state, node::PLAN)); + } + return Err(TinyAgentsError::StructuredOutput(error)); + } + } + + loop_state.finished = true; + if loop_state.final_text.is_none() { + loop_state.final_text = Some(last_assistant_text(&loop_state.messages)); + } + run.messages = loop_state.messages.clone(); + run.final_response = Some(ModelResponse::assistant( + loop_state.final_text.clone().unwrap_or_default(), + )); + + Ok(NodeResult::Command(Command { + update: Some(loop_state), + goto: vec![RouteTarget::Node(NodeId::from(crate::builder::END))], + resume: None, + resume_by_task: Default::default(), + })) +} + +/// Reconstructs the model response [`settle_node`] needs from the last +/// assistant message on the transcript. A documented simplification: the +/// full [`ModelResponse`] (usage, `finish_reason`, provider `raw`) produced by +/// [`model_node`] is not threaded through to `settle` — only the message +/// (text + tool calls) that [`tinyagents_harness::structured::StructuredExtractor`] +/// actually reads. +fn last_response_from_messages(messages: &[tinyinference_llm::message::Message]) -> ModelResponse { + for message in messages.iter().rev() { + if let tinyinference_llm::message::Message::Assistant(assistant) = message { + return ModelResponse { + message: assistant.clone(), + usage: None, + finish_reason: None, + raw: None, + resolved_model: None, + continue_turn: None, + served_from_cache: false, + correlation: None, + resolved_route: None, + }; + } + } + ModelResponse::assistant(String::new()) +} + +fn last_assistant_text(messages: &[tinyinference_llm::message::Message]) -> String { + messages + .iter() + .rev() + .find(|message| matches!(message, tinyinference_llm::message::Message::Assistant(_))) + .map(tinyinference_llm::message::Message::text) + .unwrap_or_default() +} + +/// Appends a synthetic tool-result message for every still-unanswered tool +/// call on the last assistant message, mirroring the direct loop's +/// `close_unanswered_tool_calls` so a `JumpTo(Model)`/`JumpTo(End)`/ +/// `StopWithFinal` control leaves a replayable transcript. +fn close_unanswered_tool_calls( + messages: &mut Vec, + reason: &str, +) { + let Some(tinyinference_llm::message::Message::Assistant(last)) = messages.last() else { + return; + }; + if last.tool_calls.is_empty() { + return; + } + let synthetic: Vec<_> = last + .tool_calls + .iter() + .map(|call| tinyinference_llm::message::Message::tool(call.id.clone(), reason)) + .collect(); + messages.extend(synthetic); +} + +/// Applies a drained [`MiddlewareControl`], mirroring the direct loop's +/// `apply_pending_control` but expressed as a graph routing decision instead +/// of a `LoopExit`/`ControlEffect`. +fn apply_control( + ctx: &mut RunContext, + loop_state: &mut LoopState, + control: MiddlewareControl, + from_node: &str, + natural_next: &str, +) -> Result> +where + Ctx: Send + Sync, +{ + match control { + // `Continue`/`UpdateState` request no override: route to whatever + // the calling node had already determined the turn's natural next + // step to be (see the call sites in `model_node`/`tools_node`). + MiddlewareControl::Continue => Ok(goto(loop_state.clone(), natural_next)), + MiddlewareControl::UpdateState(update) => { + ctx.push_state_update(update); + Ok(goto(loop_state.clone(), natural_next)) + } + MiddlewareControl::JumpTo(LoopTarget::Tools) => { + let route = if loop_state.pending_tool_calls.is_empty() { + node::SETTLE + } else { + node::TOOLS + }; + Ok(goto(loop_state.clone(), route)) + } + MiddlewareControl::JumpTo(LoopTarget::Model) => { + close_unanswered_tool_calls( + &mut loop_state.messages, + "run jumped back to the model before this tool call was executed", + ); + Ok(goto(loop_state.clone(), node::PLAN)) + } + MiddlewareControl::JumpTo(LoopTarget::End) => { + close_unanswered_tool_calls( + &mut loop_state.messages, + "run stopped before this tool call was executed", + ); + loop_state.finished = true; + if loop_state.final_text.is_none() { + loop_state.final_text = Some(last_assistant_text(&loop_state.messages)); + } + Ok(goto(loop_state.clone(), node::SETTLE)) + } + MiddlewareControl::StopWithFinal(text) => { + close_unanswered_tool_calls( + &mut loop_state.messages, + "run stopped before this tool call was executed", + ); + loop_state.finished = true; + loop_state.final_text = Some(text); + Ok(goto(loop_state.clone(), node::SETTLE)) + } + MiddlewareControl::Interrupt { node, message } => Ok(NodeResult::Interrupt(Interrupt { + id: format!("{from_node}-{node}"), + node: NodeId::from(node.as_str()), + payload: serde_json::json!({ "message": message }), + task_id: None, + response_schema: None, + })), + } +} + +fn goto(loop_state: LoopState, target: &str) -> NodeResult { + NodeResult::Command(Command { + update: Some(loop_state), + goto: vec![RouteTarget::Node(NodeId::from(target))], + resume: None, + resume_by_task: Default::default(), + }) +} + +/// Reconciles `ctx.config`'s per-run call caps against `policy.limits`, +/// exactly like the direct loop's `run_loop_body` does at the top of every +/// run (`resolve_call_cap` + `LimitTracker::sync_call_limits`) — without it, +/// `ctx.limits` stays at whatever `RunContext::new` derived from `config` +/// alone, silently ignoring a `RunPolicy::limits` override, in either +/// direction. Exposed so [`super::driver::GraphLoopDriver`] (which does not +/// build a [`LoopRuntime`], see that module's docs) can apply the exact same +/// reconciliation before it starts stepping nodes. +pub(crate) fn reconcile_call_limits( + ctx: &mut RunContext, + policy: &tinyagents_harness::runtime::RunPolicy, +) { + let effective_model_calls = match ctx.config.max_model_calls { + Some(explicit) => explicit.min(policy.limits.max_model_calls), + None => policy.limits.max_model_calls, + }; + let effective_tool_calls = match ctx.config.max_tool_calls { + Some(explicit) => explicit.min(policy.limits.max_tool_calls), + None => policy.limits.max_tool_calls, + }; + ctx.limits + .sync_call_limits(effective_model_calls, effective_tool_calls); +} diff --git a/crates/tinyagents-graph/src/agent_loop/types.rs b/crates/tinyagents-graph/src/agent_loop/types.rs new file mode 100644 index 00000000..890388ee --- /dev/null +++ b/crates/tinyagents-graph/src/agent_loop/types.rs @@ -0,0 +1,117 @@ +//! Types for the compiled-graph rendition of the agent loop (A5). +//! +//! See the module doc on [`super`] for the full design and its documented +//! scope relative to the harness's direct loop. + +use serde::{Deserialize, Serialize}; + +use tinyagents_harness::structured::StructuredStrategy; +use tinyinference_llm::message::Message; +use tinyinference_llm::model::ModelRequest; +use tinyinference_llm::tool::ToolCall; +use tinyinference_llm::usage::UsageTotals; + +/// The committed graph state driven around the `plan -> model -> tools -> +/// settle` loop. +/// +/// This is a whole-state graph (`Update == State`, see +/// [`crate::GraphBuilder::overwrite`]): every node returns the complete next +/// `LoopState`, not a partial patch, so [`LoopState`] doubles as its own +/// `Update` type (aliased as [`LoopUpdate`]). +/// +/// Everything here is `Serialize`/`Deserialize` so a graph-driven loop can be +/// checkpointed mid-run (including at an [`crate::Interrupt`]) and resumed — +/// unlike the harness's direct-loop [`RunContext`][tinyagents_harness::context::RunContext], +/// which is deliberately non-serializable. +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct LoopState { + /// The working transcript, in order. + pub messages: Vec, + /// Cumulative token usage across every model call so far. + pub usage: UsageTotals, + /// Extracted structured output, once a final turn produced one. + pub structured: Option, + /// Which schema variant matched (see + /// [`tinyagents_harness::middleware::AgentRun::structured_variant`]). + pub structured_variant: Option, + /// Number of loop turns (model-call + optional tool-batch pairs) executed + /// so far. + pub turn: u32, + /// Number of model calls dispatched so far. + pub model_calls: usize, + /// Number of tool invocations executed so far. + pub tool_calls: usize, + /// Names of calls that reached a tool executor, in execution order. + pub executed_tools: Vec, + /// Set once the loop has produced a terminal outcome (finished, not + /// necessarily successfully — see [`Self::final_error`]). + pub finished: bool, + /// The final assistant text, once [`Self::finished`] is set by a normal + /// completion, [`MiddlewareControl::StopWithFinal`][mc], or + /// [`MiddlewareControl::JumpTo`][mc]`(`[`LoopTarget::End`][lt]`)`. + /// + /// [mc]: tinyagents_harness::context::MiddlewareControl + /// [lt]: tinyagents_harness::context::LoopTarget + pub final_text: Option, + /// The plan built by the `plan` node for the `model` node to dispatch. + /// `None` before the first `plan` activation of a turn. + pub(crate) pending_request: Option, + /// The structured-output plan resolved alongside `pending_request`, when + /// the run requested structured output. + pub(crate) pending_structured: Option, + /// The tool calls the `model` node's response requested, for the `tools` + /// node to execute. Empty when the last model response requested none. + pub(crate) pending_tool_calls: Vec, + /// The harness-assigned id of the most recent model call, for + /// correlation. + pub(crate) last_call_id: Option, + /// How many output-validation retries have been spent so far (bounds + /// [`tinyagents_harness::runtime::RunPolicy::output_retry`]). + pub(crate) output_retry_attempts: u8, +} + +impl LoopState { + /// Seeds a fresh [`LoopState`] with `messages` as the starting + /// transcript, everything else at its `Default`. The public constructor + /// for callers outside this crate (this type's remaining fields are + /// crate-private, so a struct-literal `LoopState { messages, ..Default::default() }` + /// is not otherwise expressible from `tinyagents-integration-tests`). + pub fn seed(messages: Vec) -> Self { + Self { + messages, + ..Self::default() + } + } +} + +/// The resolved structured-output plan for the in-flight turn. Kept +/// crate-private and distinct from [`tinyagents_harness::agent_loop::phases::StructuredPlan`] +/// only in that it carries the real [`StructuredStrategy`] (needed to build a +/// [`tinyagents_harness::structured::StructuredExtractor`] without +/// re-deriving it from a string tag). +#[derive(Clone, Debug, Serialize, Deserialize)] +pub(crate) struct PendingStructuredPlan { + pub(crate) strategy: StructuredStrategy, + pub(crate) schema_name: String, + pub(crate) schema: serde_json::Value, +} + +/// [`LoopState`] doubles as its own partial-update type: every node in +/// [`super::compile_loop`]'s graph returns the whole next state (see +/// [`crate::GraphBuilder::overwrite`]). +pub type LoopUpdate = LoopState; + +/// Node ids used by [`super::compile_loop`]'s compiled graph. Exposed so a +/// caller can name a node for [`super::LoopIter::override_next`] or interpret +/// a [`super::LoopStep::node`]. +pub mod node { + /// Builds the next turn's [`super::TurnPlan`][tinyagents_harness::agent_loop::phases::TurnPlan]-shaped request. + pub const PLAN: &str = "plan"; + /// Dispatches the model call built by [`PLAN`]. + pub const MODEL: &str = "model"; + /// Executes the tool calls the last model response requested. + pub const TOOLS: &str = "tools"; + /// Settles the turn: structured extraction/validation and the + /// finish/continue decision. + pub const SETTLE: &str = "settle"; +} diff --git a/crates/tinyagents-graph/src/builder/README.md b/crates/tinyagents-graph/src/builder/README.md index ad9e0ac4..70a05532 100644 --- a/crates/tinyagents-graph/src/builder/README.md +++ b/crates/tinyagents-graph/src/builder/README.md @@ -21,9 +21,14 @@ recursively-generated sub-workflow compiles through, one level down. - `mark_command_routing`, `with_command_destinations` — declares a node routes exclusively via `Command::goto` rather than static/conditional edges; `compile()` rejects nodes that mix the two. -- `with_node_kind`, `with_node_metadata`, `mark_subgraph`, `mark_interrupt`, - `mark_deferred` — behavior-free introspection markers surfaced by - `graph::export`. +- `with_node_kind`, `with_node_metadata`, `mark_subgraph` — behavior-free + introspection markers surfaced by `graph::export`. `mark_deferred` and + `mark_interrupt` also set the marker but are *not* behavior-free: + `mark_deferred` is `NodePolicy::defer`, and `mark_interrupt` is an alias + for `interrupt_before`. +- `interrupt_before`, `interrupt_after` — executor-level pauses at named + nodes (before the handler runs / after it runs but before its result is + applied); see `docs/modules/graph/interrupts.md`. - `with_parallel`, `with_max_concurrency`, `with_node_timeout`, `with_recursion_limit`, `with_graph_id`, `with_name`, `set_defaults` — per-graph configuration, either called directly or bundled via diff --git a/crates/tinyagents-graph/src/builder/mod.rs b/crates/tinyagents-graph/src/builder/mod.rs index 7ca5c44a..ad795a60 100644 --- a/crates/tinyagents-graph/src/builder/mod.rs +++ b/crates/tinyagents-graph/src/builder/mod.rs @@ -11,12 +11,14 @@ //! See `types` for the builder data types. `compile` validates the topology //! and freezes it into an immutable [`crate::CompiledGraph`]. +mod policy; mod types; -pub(crate) use types::{Branch, BuilderNode, NodeMeta}; +pub use policy::{CacheKeyFn, NodeCachePolicy, NodePolicy, OnErrorFn}; +pub(crate) use types::{Branch, BuilderNode, NodeMeta, UpdateCodec}; pub use types::{ - END, ForkId, GraphBuilder, GraphDefaults, NodeContext, NodeFuture, NodeHandler, Route, - RouterFn, START, + END, ForkId, GraphBuilder, GraphDefaults, IdleClock, NodeContext, NodeFuture, NodeHandler, + Route, RouterFn, START, }; use std::collections::{HashMap, HashSet}; @@ -79,6 +81,7 @@ where nodes: HashMap::new(), edges: HashMap::new(), branches: HashMap::new(), + route_label_checks: HashMap::new(), command_nodes: HashSet::new(), waiting: HashMap::new(), barrier_reliefs: Vec::new(), @@ -88,9 +91,41 @@ where max_concurrency: None, node_timeout: None, node_meta: HashMap::new(), + node_policies: HashMap::new(), + node_defaults: None, + interrupt_before: HashSet::new(), + interrupt_after: HashSet::new(), + update_codec: None, } } + /// Attaches a per-node execution [`NodePolicy`] (retry, timeouts, cache, + /// `on_error`, `defer`) to `node`, replacing any policy previously set + /// for it. At run time each field falls back to the + /// [`Self::set_node_defaults`] policy, then to the legacy graph-wide + /// `with_node_timeout`/`with_node_retry` settings — see + /// [`NodePolicy`]'s module docs for the exact precedence. + pub fn with_node_policy( + mut self, + node: impl Into, + policy: NodePolicy, + ) -> Self { + let node = node.into(); + // Keep the export-only marker in sync with the runtime flag. + if policy.defer { + self.node_meta.entry(node.clone()).or_default().deferred = true; + } + self.node_policies.insert(node, policy); + self + } + + /// Sets the graph-wide default [`NodePolicy`] every node falls back to, + /// field by field, when it has no per-node override. + pub fn set_node_defaults(mut self, policy: NodePolicy) -> Self { + self.node_defaults = Some(policy); + self + } + /// Applies a bundle of [`GraphDefaults`] in one call. Only the `Some` fields /// override the builder's current configuration, so this composes with /// explicit `with_*` calls regardless of ordering. @@ -173,17 +208,78 @@ where self } + /// Registers a named [`crate::Reducer`] closure in the + /// process-wide [`crate::channel::ReducerRegistry`], returning the + /// builder for chaining. + /// + /// This is what makes a [`crate::BinaryAggregate`] channel serializable: + /// `BinaryAggregate::named(name)` looks the closure back up by name (see + /// its docs), and a channel built that way persists only `name` in its + /// [`crate::Channel::config`] — decoding a checkpoint later, in this or + /// another process, requires the same name to have been registered + /// first. The built-ins `"append"`, `"last"`, `"sum"`, `"max"`, `"min"`, + /// and `"set_union"` are always available with no registration. + /// + /// The registry is global rather than scoped to this builder because + /// checkpoint decode has no builder in scope at all — see + /// `crate::channel::registry`'s module docs. + pub fn register_reducer( + self, + name: impl Into, + f: impl Fn(serde_json::Value, serde_json::Value) -> Result + + Send + + Sync + + 'static, + ) -> Self { + crate::channel::ReducerRegistry::register(name, f); + self + } + /// Adds an async node returning a [`NodeResult`]. + /// + /// This is a thin by-value adapter over [`Self::add_node_shared`] (M2 in + /// `docs/runtime-comparison/code-review-graph.md`): internally every + /// handler receives the step's state as an `Arc`, and this + /// adapter clones out of it once per invocation so the handler closure + /// keeps taking an owned `State` exactly as before — every existing + /// caller of `add_node` compiles unchanged. A handler that does not need + /// to mutate or move its own copy of `State` should prefer + /// [`Self::add_node_shared`] instead, which hands it the `Arc` + /// directly and clones nothing. pub fn add_node(mut self, id: impl Into, handler: F) -> Self where F: Fn(State, NodeContext) -> Fut + Send + Sync + 'static, Fut: Future>> + Send + 'static, { - let id = id.into(); self.nodes.insert( - id.clone(), + id.into(), + BuilderNode { + handler: Arc::new(move |state: Arc, ctx| { + Box::pin(handler((*state).clone(), ctx)) + }), + }, + ); + self + } + + /// Adds an async node that receives the step's committed state directly + /// as an `Arc`, returning a [`NodeResult`]. + /// + /// The zero-clone counterpart to [`Self::add_node`] (M2): a superstep + /// clones `State` at most once (building the `Arc` the executor threads + /// through that step), and every branch/attempt of a handler added this + /// way shares that allocation via a cheap `Arc::clone` — no per-attempt, + /// per-branch `State` clone at all. Prefer this over [`Self::add_node`] + /// for a large `State` (e.g. a message-history-carrying value) or a node + /// that only reads its state. + pub fn add_node_shared(mut self, id: impl Into, handler: F) -> Self + where + F: Fn(Arc, NodeContext) -> Fut + Send + Sync + 'static, + Fut: Future>> + Send + 'static, + { + self.nodes.insert( + id.into(), BuilderNode { - id, handler: Arc::new(move |state, ctx| Box::pin(handler(state, ctx))), }, ); @@ -192,8 +288,14 @@ where /// Adds a direct edge `from -> to`. Use [`START`]/[`END`] for the virtual /// entry/terminal nodes. + /// + /// Calling this more than once for the same `from` accumulates a static + /// **fan-out**: every registered `to` activates (not just the last one + /// registered), matching the documented "one or more node names" routing + /// contract. Adding the exact same `(from, to)` edge twice is a no-op — + /// the target is not scheduled twice. pub fn add_edge(mut self, from: impl Into, to: impl Into) -> Self { - self.edges.insert(from.into(), to.into()); + Self::push_edge(&mut self.edges, from.into(), to.into()); self } @@ -208,11 +310,20 @@ where { let nodes: Vec = nodes.into_iter().map(Into::into).collect(); for pair in nodes.windows(2) { - self.edges.insert(pair[0].clone(), pair[1].clone()); + Self::push_edge(&mut self.edges, pair[0].clone(), pair[1].clone()); } self } + /// Appends `to` to `from`'s static successor list, deduplicating so the + /// same target is never scheduled twice from one static fan-out. + fn push_edge(edges: &mut HashMap>, from: NodeId, to: NodeId) { + let targets = edges.entry(from).or_default(); + if !targets.contains(&to) { + targets.push(to); + } + } + /// Adds a barrier/waiting edge `from -> to`: like [`Self::add_edge`] but `to` /// only activates once *all* of its registered predecessors (every `from` /// declared via `add_waiting_edge`) have completed — possibly across @@ -223,7 +334,7 @@ where pub fn add_waiting_edge(mut self, from: impl Into, to: impl Into) -> Self { let from = from.into(); let to = to.into(); - self.edges.insert(from.clone(), to.clone()); + Self::push_edge(&mut self.edges, from.clone(), to.clone()); self.waiting.entry(to).or_default().insert(from); self } @@ -295,13 +406,72 @@ where self.branches.insert( from.into(), Branch { - router: Arc::new(move |state| router(state).to_string()), + router: Arc::new(move |state| Route::new(router(state))), routes, }, ); self } + /// Like [`Self::add_conditional_edges`], but additionally declares the + /// **exhaustive** set of labels `router` can ever return. + /// + /// [`Self::compile`] (via [`Self::validate_routes`]) cross-checks + /// `all_labels` against `routes`'s keys and rejects the build if a + /// declared label has no route — catching a typo'd route label (e.g. the + /// router returns `AgentRoute::Toool` because `Toool`/`Tool` are both + /// wired but one is missing from `routes`) before the graph ever runs, + /// instead of only failing at run time with + /// [`crate::TinyAgentsError::MissingRoute`] on whichever branch happens + /// to be taken. + /// + /// `all_labels` shares `router`'s return type `R`, so the compiler (not + /// just this check) ties the declared label set to what the router can + /// actually produce — a typed enum with, e.g., a `strum::EnumIter`-style + /// listing of its own variants is the natural `all_labels` source. + pub fn add_conditional_edges_checked( + mut self, + from: impl Into, + router: F, + routes: I, + all_labels: L, + ) -> Self + where + F: Fn(&State) -> R + Send + Sync + 'static, + R: ToString, + I: IntoIterator, + K: ToString, + V: Into, + L: IntoIterator, + { + let from = from.into(); + let labels: Vec = all_labels.into_iter().map(|l| l.to_string()).collect(); + self = self.add_conditional_edges(from.clone(), router, routes); + self.route_label_checks.insert(from, labels); + self + } + + /// Cross-checks every [`Self::add_conditional_edges_checked`] declaration + /// against its node's actual route table, returning + /// [`crate::TinyAgentsError::MissingRoute`] for the first declared label + /// with no matching route. Called automatically by [`Self::compile`]. + fn validate_routes(&self) -> Result<()> { + for (node, labels) in &self.route_label_checks { + let Some(branch) = self.branches.get(node) else { + continue; + }; + for label in labels { + if !branch.routes.contains_key(label) { + return Err(TinyAgentsError::MissingRoute { + node: node.to_string(), + route: label.clone(), + }); + } + } + } + Ok(()) + } + /// Declares that `node` routes exclusively via [`crate::Command`] /// `goto` (not static or conditional edges). Compile rejects nodes that mix /// command routing with static/conditional edges. @@ -362,15 +532,74 @@ where self } - /// Marks `node` as an interrupt point for the export. - pub fn mark_interrupt(mut self, node: impl Into) -> Self { - self.node_meta.entry(node.into()).or_default().interrupt = true; + /// Marks `node` as an interrupt point: an alias for + /// [`Self::interrupt_before`] (the run pauses before `node` executes) + /// that also sets the export-facing interrupt marker + /// (`NodeInfo::interrupt`). Earlier versions set only the marker; the + /// runtime pause is now real. + pub fn mark_interrupt(self, node: impl Into) -> Self { + self.interrupt_before([node]) + } + + /// Pauses the run *before* each of `nodes` executes. + /// + /// When the executor is about to invoke a listed node, it instead + /// records an [`Interrupt`](crate::Interrupt) for that activation with + /// payload `{"phase": "before"}` (stamped with the task id) and persists + /// an interrupt-boundary checkpoint, without calling the handler. The + /// handler runs exactly once overall: `resume` re-schedules the paused + /// activation and runs it normally, delivering any `Command::resume` + /// value on [`NodeContext::resume`]. Requires a checkpointer and a + /// thread, like any interrupt. Nodes are validated at [`Self::compile`]; + /// the export marks them as interrupt points. + pub fn interrupt_before(mut self, nodes: impl IntoIterator>) -> Self { + for node in nodes { + let node = node.into(); + self.node_meta.entry(node.clone()).or_default().interrupt = true; + self.interrupt_before.insert(node); + } self } - /// Marks `node` as a deferred join for the export. + /// Pauses the run *after* each of `nodes` has run, before its result is + /// applied. + /// + /// The handler runs to completion; its `Update`/`Command` is then held + /// back from committed state — serialized as a deferred-result write + /// (`PendingWrite::interrupt_after`) in the interrupt-boundary + /// checkpoint — and an [`Interrupt`](crate::Interrupt) with payload + /// `{"phase": "after"}` is returned. The paused run's state (and the + /// checkpoint's) therefore does *not* yet include the node's write. On + /// `resume`, the executor replays the stored result — applying the + /// update through the reducer and honouring the node's `goto` — without + /// invoking the handler again, so the handler still runs exactly once. + /// A node that itself returns `NodeResult::Interrupt` is not paused a + /// second time. Requires `Update: Serialize + DeserializeOwned` (the + /// codec for the deferred write), plus a checkpointer and a thread. + pub fn interrupt_after(mut self, nodes: impl IntoIterator>) -> Self + where + Update: serde::Serialize + serde::de::DeserializeOwned, + { + if self.update_codec.is_none() { + self.update_codec = Some(UpdateCodec::serde()); + } + for node in nodes { + let node = node.into(); + self.node_meta.entry(node.clone()).or_default().interrupt = true; + self.interrupt_after.insert(node); + } + self + } + + /// Marks `node` as a deferred join: it is surfaced as deferred in the + /// export *and* scheduled with [`NodePolicy::defer`] semantics — it + /// only runs once nothing else is left in the frontier. Equivalent to + /// `with_node_policy(node, NodePolicy { defer: true, ..existing })`, + /// merging with any policy already set for the node. pub fn mark_deferred(mut self, node: impl Into) -> Self { - self.node_meta.entry(node.into()).or_default().deferred = true; + let node = node.into(); + self.node_meta.entry(node.clone()).or_default().deferred = true; + self.node_policies.entry(node).or_default().defer = true; self } @@ -382,11 +611,26 @@ where )); } - // entry must exist - let entry = self + // entry must exist, and be exactly one node: START does not fan out. + let start_targets = self .edges .get(&NodeId::from(START)) .cloned() + .unwrap_or_default(); + if start_targets.len() > 1 { + return Err(TinyAgentsError::Validation(format!( + "START must route to exactly one entry node, got {}: {}", + start_targets.len(), + start_targets + .iter() + .map(ToString::to_string) + .collect::>() + .join(", ") + ))); + } + let entry = start_targets + .into_iter() + .next() .ok_or(TinyAgentsError::MissingStart)?; if entry.as_str() == END { return Err(TinyAgentsError::Validation( @@ -396,25 +640,30 @@ where self.require_node(&entry)?; // static edges - for (from, to) in &self.edges { + for (from, targets) in &self.edges { if from.as_str() != START { self.require_node(from)?; } - if to.as_str() != END { - self.require_node(to)?; - } - if to.as_str() == START { - return Err(TinyAgentsError::Validation( - "START cannot be an edge target".to_string(), - )); - } if from.as_str() == END { return Err(TinyAgentsError::Validation( "END cannot be an edge source".to_string(), )); } + for to in targets { + if to.as_str() != END { + self.require_node(to)?; + } + if to.as_str() == START { + return Err(TinyAgentsError::Validation( + "START cannot be an edge target".to_string(), + )); + } + } } + // conditional route labels declared exhaustive must all have a route + self.validate_routes()?; + // conditional edges for (from, branch) in &self.branches { self.require_node(from)?; @@ -438,6 +687,11 @@ where } } + // interrupt selectors must name real nodes + for node in self.interrupt_before.iter().chain(&self.interrupt_after) { + self.require_node(node)?; + } + // command-routing nodes must not also have static/conditional edges for node in &self.command_nodes { self.require_node(node)?; @@ -454,6 +708,7 @@ where nodes, edges, branches, + route_label_checks: _, command_nodes, waiting, reducer, @@ -463,6 +718,11 @@ where node_timeout, node_meta, barrier_reliefs, + node_policies, + node_defaults, + interrupt_before, + interrupt_after, + update_codec, } = self; Ok(CompiledGraph::from_parts( @@ -481,7 +741,9 @@ where node_timeout, node_meta, barrier_reliefs, - )) + ) + .with_node_policies(node_policies, node_defaults) + .with_interrupt_selectors(interrupt_before, interrupt_after, update_codec)) } fn require_node(&self, id: &NodeId) -> Result<()> { diff --git a/crates/tinyagents-graph/src/builder/policy.rs b/crates/tinyagents-graph/src/builder/policy.rs new file mode 100644 index 00000000..724cb519 --- /dev/null +++ b/crates/tinyagents-graph/src/builder/policy.rs @@ -0,0 +1,219 @@ +//! Per-node execution policy: retry, timeouts, caching, error recovery, and +//! deferred scheduling. +//! +//! A [`NodePolicy`] is attached to one node via +//! [`GraphBuilder::with_node_policy`](super::GraphBuilder::with_node_policy) +//! or to every node via +//! [`GraphBuilder::set_node_defaults`](super::GraphBuilder::set_node_defaults). +//! At run time the executor resolves an *effective* policy for each +//! activation field by field ([`NodePolicy::resolve`]): a per-node `Some` +//! wins, else the graph-wide default's field, else the older graph-wide +//! `with_node_retry` / `with_node_timeout` settings, else nothing. + +use std::sync::Arc; +use std::time::Duration; + +use crate::command::Command; +use tinyagents_harness::error::TinyAgentsError; +use tinyagents_harness::retry::RetryPolicy; + +/// Computes a task-cache key from the committed state snapshot and the +/// activation's [`crate::NodeContext::send_arg`] (if any). +pub type CacheKeyFn = dyn Fn(&State, Option<&serde_json::Value>) -> String + Send + Sync; + +/// Recovers from a node failure: given the state the node ran against and +/// the error that survived its retry policy, optionally produce a +/// [`Command`] that stands in for the node's result. +pub type OnErrorFn = + dyn Fn(&State, &TinyAgentsError) -> Option> + Send + Sync; + +/// Opt-in result caching for one node. +/// +/// `key` derives the cache key from the state snapshot and `send_arg`; an +/// identical key on a later activation replays the cached `Update` without +/// invoking the handler. `ttl` bounds how long an entry stays valid +/// (`None` = forever, until [`crate::cache::TaskCache::clear`]). +/// +/// The policy itself is bound-free over `Update`. Actually (de)serializing +/// a cached update needs `Update: Serialize + DeserializeOwned`, which is +/// only required by +/// [`CompiledGraph::with_cached_node`](crate::CompiledGraph::with_cached_node) +/// (the entry point that installs the codec) — see its docs. +pub struct NodeCachePolicy { + /// Derives the cache key for one activation. + pub key: Arc>, + /// Optional time-to-live for a cached entry. + pub ttl: Option, +} + +impl NodeCachePolicy { + /// Builds a cache policy from a key function, with no TTL. + pub fn new(key: F) -> Self + where + F: Fn(&State, Option<&serde_json::Value>) -> String + Send + Sync + 'static, + { + Self { + key: Arc::new(key), + ttl: None, + } + } + + /// Sets the entry time-to-live. + pub fn with_ttl(mut self, ttl: Duration) -> Self { + self.ttl = Some(ttl); + self + } +} + +impl Clone for NodeCachePolicy { + fn clone(&self) -> Self { + Self { + key: self.key.clone(), + ttl: self.ttl, + } + } +} + +impl std::fmt::Debug for NodeCachePolicy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("NodeCachePolicy") + .field("ttl", &self.ttl) + .finish_non_exhaustive() + } +} + +/// Execution policy for one node (or, via `set_node_defaults`, every node). +/// +/// Every field is optional and additive; the [`Default`] changes nothing. +/// See the module docs for how per-node and graph-wide policies compose. +pub struct NodePolicy { + /// Retry policy applied around the handler; a + /// [retryable](tinyagents_harness::retry::is_retryable) error is re-run + /// from the node's start up to the policy's attempt cap. + pub retry: Option, + /// Maximum total wall-clock time for one attempt of the handler, + /// regardless of heartbeats. + pub timeout: Option, + /// Maximum time between two [`crate::NodeContext::heartbeat`] calls (or + /// between start and the first one). A handler that never heartbeats + /// sees this as a flat timeout of the same duration. + pub idle_timeout: Option, + /// Opt-in result caching; see [`NodeCachePolicy`]. + pub cache: Option>, + /// Error recovery hook consulted once the retry budget is exhausted (or + /// the error was not retryable). Returning `Some(command)` makes the + /// node complete with that command instead of failing the run. + pub on_error: Option>>, + /// Deferred scheduling: the node only runs once nothing *else* is left + /// in the frontier (a "run when nothing else is ready" synthesis join). + pub defer: bool, +} + +impl Default for NodePolicy { + fn default() -> Self { + Self { + retry: None, + timeout: None, + idle_timeout: None, + cache: None, + on_error: None, + defer: false, + } + } +} + +impl Clone for NodePolicy { + fn clone(&self) -> Self { + Self { + retry: self.retry.clone(), + timeout: self.timeout, + idle_timeout: self.idle_timeout, + cache: self.cache.clone(), + on_error: self.on_error.clone(), + defer: self.defer, + } + } +} + +impl std::fmt::Debug for NodePolicy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("NodePolicy") + .field("retry", &self.retry) + .field("timeout", &self.timeout) + .field("idle_timeout", &self.idle_timeout) + .field("cache", &self.cache) + .field("has_on_error", &self.on_error.is_some()) + .field("defer", &self.defer) + .finish() + } +} + +impl NodePolicy { + /// Sets the retry policy. + pub fn with_retry(mut self, retry: RetryPolicy) -> Self { + self.retry = Some(retry); + self + } + + /// Sets the flat per-attempt timeout. + pub fn with_timeout(mut self, timeout: Duration) -> Self { + self.timeout = Some(timeout); + self + } + + /// Sets the idle (heartbeat) timeout. + pub fn with_idle_timeout(mut self, idle_timeout: Duration) -> Self { + self.idle_timeout = Some(idle_timeout); + self + } + + /// Sets the cache policy. + pub fn with_cache(mut self, cache: NodeCachePolicy) -> Self { + self.cache = Some(cache); + self + } + + /// Sets the error-recovery hook. + pub fn with_on_error(mut self, on_error: F) -> Self + where + F: Fn(&State, &TinyAgentsError) -> Option> + Send + Sync + 'static, + { + self.on_error = Some(Arc::new(on_error)); + self + } + + /// Marks the node deferred. + pub fn deferred(mut self) -> Self { + self.defer = true; + self + } + + /// Resolves the effective policy for one node, field by field. + /// + /// Precedence per field: `per_node` (`Some` wins) → `defaults` → + /// the legacy graph-wide `retry`/`timeout` (from + /// `CompiledGraph::with_node_retry` / `GraphBuilder::with_node_timeout`). + /// `defer` is the logical OR of the two policies' flags. + pub(crate) fn resolve( + per_node: Option<&Self>, + defaults: Option<&Self>, + legacy_retry: Option<&RetryPolicy>, + legacy_timeout: Option, + ) -> Self { + let pick = |f: fn(&Self) -> bool| -> Option<&Self> { + per_node.filter(|p| f(p)).or(defaults.filter(|p| f(p))) + }; + Self { + retry: pick(|p| p.retry.is_some()) + .and_then(|p| p.retry.clone()) + .or_else(|| legacy_retry.cloned()), + timeout: pick(|p| p.timeout.is_some()) + .and_then(|p| p.timeout) + .or(legacy_timeout), + idle_timeout: pick(|p| p.idle_timeout.is_some()).and_then(|p| p.idle_timeout), + cache: pick(|p| p.cache.is_some()).and_then(|p| p.cache.clone()), + on_error: pick(|p| p.on_error.is_some()).and_then(|p| p.on_error.clone()), + defer: per_node.is_some_and(|p| p.defer) || defaults.is_some_and(|p| p.defer), + } + } +} diff --git a/crates/tinyagents-graph/src/builder/test.rs b/crates/tinyagents-graph/src/builder/test.rs index f14f04de..c0821444 100644 --- a/crates/tinyagents-graph/src/builder/test.rs +++ b/crates/tinyagents-graph/src/builder/test.rs @@ -80,6 +80,82 @@ fn compile_rejects_static_and_conditional_on_same_node() { assert!(matches!(err, TinyAgentsError::Validation(_))); } +#[test] +fn add_edge_accumulates_static_fan_out_without_duplicates() { + let builder = GraphBuilder::::overwrite() + .add_node("a", |s: S, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .add_node("b", |s: S, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .add_node("c", |s: S, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .set_entry("a") + .add_edge("a", "b") + .add_edge("a", "c") + .add_edge("a", "b"); // duplicate of an existing edge: must not double-schedule "b" + + let targets = builder.edges.get(&NodeId::from("a")).cloned().unwrap(); + assert_eq!( + targets, + vec![NodeId::from("b"), NodeId::from("c")], + "add_edge must accumulate a static fan-out list and dedupe repeats" + ); +} + +#[test] +fn add_conditional_edges_checked_catches_route_typo_at_build_time() { + let err = GraphBuilder::::overwrite() + .add_node("a", |s: S, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .add_node("b", |s: S, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .set_entry("a") + .add_conditional_edges_checked( + "a", + |_s: &S| "tool".to_string(), + // typo: the route table declares "tol", not "tool" + [("tol", "b")], + ["tool".to_string(), "final".to_string()], + ) + .set_finish("b") + .compile() + .unwrap_err(); + + match err { + TinyAgentsError::MissingRoute { node, route } => { + assert_eq!(node, "a"); + assert_eq!(route, "tool"); + } + other => panic!("expected MissingRoute at build time, got {other:?}"), + } +} + +#[test] +fn add_conditional_edges_checked_accepts_matching_routes() { + let compiled = GraphBuilder::::overwrite() + .add_node("a", |s: S, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .add_node("b", |s: S, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .set_entry("a") + .add_conditional_edges_checked( + "a", + |_s: &S| "tool".to_string(), + [("tool", "b"), ("final", "b")], + ["tool".to_string(), "final".to_string()], + ) + .set_finish("b") + .compile(); + assert!(compiled.is_ok()); +} + #[test] fn compile_succeeds_for_valid_graph() { let compiled = GraphBuilder::::overwrite() diff --git a/crates/tinyagents-graph/src/builder/types.rs b/crates/tinyagents-graph/src/builder/types.rs index 2480fa4f..71eb604f 100644 --- a/crates/tinyagents-graph/src/builder/types.rs +++ b/crates/tinyagents-graph/src/builder/types.rs @@ -14,9 +14,10 @@ use std::sync::Arc; use std::time::Duration; use crate::Result; +use crate::checkpoint::PendingWrite; use crate::command::NodeResult; use crate::reducer::StateReducer; -use tinyagents_harness::ids::{GraphId, NodeId, RunId, ThreadId}; +use tinyagents_harness::ids::{GraphId, NodeId, RunId, TaskId, ThreadId}; /// The reserved virtual entry node. pub const START: &str = "__start__"; @@ -28,12 +29,29 @@ pub type NodeFuture = Pin` (M2 in `docs/runtime-comparison/code-review-graph.md`): a +/// superstep clones `State` at most once (building this `Arc`), and every +/// branch/attempt within that step shares it via a cheap `Arc::clone` +/// instead of re-cloning the whole state. [`super::GraphBuilder::add_node`] +/// (the by-value convenience entry point) is a thin adapter over this +/// signature that clones out of the `Arc` on every invocation; callers that +/// want the zero-clone path use +/// [`super::GraphBuilder::add_node_shared`], whose closure receives the +/// `Arc` directly. pub type NodeHandler = - dyn Fn(State, NodeContext) -> NodeFuture + Send + Sync; + dyn Fn(Arc, NodeContext) -> NodeFuture + Send + Sync; /// A conditional routing function over committed state. Returns a route label /// resolved against the node's route table at the step boundary. -pub type RouterFn = dyn Fn(&State) -> String + Send + Sync; +/// +/// Internally this returns a typed [`Route`] rather than a bare `String` — +/// `Route` is `From`/cheaply stringifies, so this is purely a +/// representation change and does not affect +/// [`super::GraphBuilder::add_conditional_edges`]'s public signature, which +/// still accepts any router closure returning `impl ToString`. +pub type RouterFn = dyn Fn(&State) -> Route + Send + Sync; /// Identifies one branch of a concurrent (fan-out) superstep. /// @@ -62,6 +80,56 @@ impl ForkId { } } +/// The heartbeat channel between a running node handler and the executor's +/// idle-timeout watcher (see [`NodeContext::heartbeat`]). +/// +/// Cheap to clone; every clone of a [`NodeContext`] shares the same clock, +/// which is how a `heartbeat()` call made *inside* the handler future is +/// observed by the timeout race wrapped *around* it. A fresh clock is built +/// per activation; a hand-built context can use the [`Default`]. +#[derive(Clone, Default)] +pub struct IdleClock { + notify: Arc, + beats: Arc, +} + +impl IdleClock { + /// Records a heartbeat, waking the idle-timeout watcher so it re-arms. + pub fn touch(&self) { + self.beats + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + self.notify.notify_one(); + } + + /// Total heartbeats recorded so far. + pub fn beats(&self) -> u64 { + self.beats.load(std::sync::atomic::Ordering::Relaxed) + } + + /// Resolves once `idle` elapses with no heartbeat in between; every + /// [`Self::touch`] restarts the window. Never resolves if heartbeats keep + /// arriving inside the window. With no heartbeat at all this resolves + /// exactly `idle` after it is first polled — a flat timeout. + pub(crate) async fn idle_elapsed(&self, idle: Duration) { + loop { + let sleep = tokio::time::sleep(idle); + tokio::pin!(sleep); + tokio::select! { + _ = &mut sleep => return, + _ = self.notify.notified() => continue, + } + } + } +} + +impl std::fmt::Debug for IdleClock { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("IdleClock") + .field("beats", &self.beats()) + .finish() + } +} + /// Per-task runtime context passed to a durable node handler. /// /// The context exposes run identity, the current step, and — crucially — an @@ -91,7 +159,13 @@ pub struct NodeContext { /// This is how map-reduce / search-fanout branches and external graph /// inputs receive custom data that differs from the graph's shared /// committed state. - pub send_arg: Option, + /// + /// `Arc`-wrapped (M2) so a repeated `Send` fan-out of the same node, and + /// every retry attempt of one activation, share the same allocation + /// instead of deep-cloning the argument per attempt. Serializes + /// transparently as the underlying `serde_json::Value` (serde's blanket + /// `Arc` impl), so on-disk checkpoint records are unaffected. + pub send_arg: Option>, /// The root run id of the recursion tree this node executes within. For a /// top-level run this equals `run_id`; for a subgraph/sub-agent child run it /// is the shared ancestor, so a child a node spawns can preserve the root. @@ -107,6 +181,134 @@ pub struct NodeContext { /// Complete host-owned recursive-agent binding for this execution, if one /// was supplied at the graph entry point. pub agent_binding: Option, + /// Stable identity of this scheduled activation within its superstep + /// (R5). Distinguishes repeated `Send` fan-out activations of the same + /// node — a subgraph node consults this (with [`Self::siblings`]) to + /// namespace its child checkpoint per fan-out branch instead of sharing + /// one namespace across every concurrent activation of the node (I1). + pub task_id: TaskId, + /// The number of activations of [`Self::node_id`] in this same + /// superstep's active set (I1). `1` for an ordinary (non-fan-out) + /// activation; greater than `1` means a `Send` fan-out scheduled several + /// concurrent activations of this node this step. + pub siblings: usize, + /// The channel versions (I5/R3) as this node's state snapshot sees them + /// — [`crate::channel::ChannelState::channel_versions`] for a channel + /// graph, or a single `{"state": n}` entry for a plain whole-state + /// graph. Compared against [`Self::versions_seen`] by + /// [`Self::changed_since_last_run`]. + pub channel_versions: BTreeMap, + /// This node's own snapshot of [`Self::channel_versions`] as of the last + /// time it ran (empty on a node's first-ever activation in the thread). + pub versions_seen: BTreeMap, + /// Heartbeat channel for the node's idle timeout + /// ([`crate::NodePolicy::idle_timeout`]); see [`Self::heartbeat`]. + pub idle_clock: IdleClock, + /// This task's [`Self::durable_task`] memo buffer, shared by every clone + /// of the context (a retried attempt sees the first attempt's memos). + /// Pre-seeded by the executor with the durable-task writes the task's + /// checkpoint already holds (so a re-run after an interrupt, crash, or + /// retry hits instead of re-executing), and appended to on every miss; + /// the executor folds it back into the boundary checkpoint's + /// `pending_writes` when the task stalls. Empty on a hand-built context. + pub(crate) durable_writes: Arc>>, +} + +impl NodeContext { + /// Runs `fut` at most once per `(task, key)` across re-runs of this task. + /// + /// A node handler is re-run from its start after an interrupt/resume, + /// a failure/retry, or an in-process node retry, so any side effect it + /// performs (an API call, a payment, a counter increment) would repeat. + /// Wrapping that side effect in `durable_task` memoises its output in + /// this task's checkpoint write ledger, keyed by the task id and `key`: + /// the first execution awaits `fut`, serializes its `Ok` output as a + /// [`PendingWrite::durable_task`] memo, and returns it; a later re-run + /// of the same task finds the memo and returns the stored value + /// **without polling `fut` at all**. Memos are only ever recorded for a + /// successful `fut`; an `Err` is returned unmemoised so a retry re-runs + /// the step. Keys are independent: a handler that memoised `"a"` and + /// then failed before `"b"` replays `"a"` and runs `"b"` fresh. + /// + /// The memo is scoped to this task in this thread — it is not a + /// cross-run cache (see [`crate::TaskCache`] for that) — and it is only + /// durable across process restarts when the graph runs on a + /// checkpointed thread; without a checkpointer it still dedupes within + /// one run (in-process retries). Two calls with the same `key` inside + /// one handler execution return the same stored value. + pub async fn durable_task(&self, key: &str, fut: F) -> Result + where + T: serde::Serialize + serde::de::DeserializeOwned, + F: Future>, + { + let hit = self + .lock_durable_writes() + .iter() + .find(|w| w.durable_task_key() == Some(key)) + .map(|w| w.payload.clone()); + if let Some(payload) = hit { + return serde_json::from_value(payload).map_err(|err| { + crate::TinyAgentsError::Serialization(serde::de::Error::custom(format!( + "durable_task `{key}` of task `{}` (node `{}`) holds a memo that does not decode as the requested type: {err}", + self.task_id.as_str(), + self.node_id + ))) + }); + } + let value = fut.await?; + let payload = serde_json::to_value(&value)?; + let mut writes = self.lock_durable_writes(); + let idx = writes.iter().map(|w| w.idx).max().unwrap_or(0).max(0) + 1; + writes.push(PendingWrite::durable_task( + self.node_id.clone(), + self.task_id.clone(), + idx, + key, + payload, + )); + Ok(value) + } + + /// Locks the durable-task memo buffer, tolerating a poisoned lock (the + /// buffer is a plain `Vec` push/scan, so a panic mid-hold leaves it + /// consistent). + pub(crate) fn lock_durable_writes(&self) -> std::sync::MutexGuard<'_, Vec> { + self.durable_writes + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + /// A snapshot of this task's durable-task memo writes (pre-seeded plus + /// any recorded by [`Self::durable_task`] so far). + pub(crate) fn durable_writes_snapshot(&self) -> Vec { + self.lock_durable_writes().clone() + } + + /// Signals liveness to the executor's idle-timeout watcher, restarting + /// the node's [`crate::NodePolicy::idle_timeout`] window. Cheap (an + /// atomic increment plus a notify); a no-op for a node with no idle + /// timeout configured. Does not affect the flat + /// [`crate::NodePolicy::timeout`] ceiling. + pub fn heartbeat(&self) { + self.idle_clock.touch(); + } + + /// This activation's stable task identity (R5). See the field docs on + /// [`Self::task_id`]. + pub fn task_id(&self) -> &TaskId { + &self.task_id + } + + /// Whether `channel` has been written since the last time this node ran + /// (I5/R3): compares [`Self::channel_versions`] (current) against + /// [`Self::versions_seen`] (this node's own last-observed snapshot). A + /// channel this node has never seen before (including a node's very + /// first activation) counts as changed whenever it has ever been + /// written at all. + pub fn changed_since_last_run(&self, channel: &str) -> bool { + self.channel_versions.get(channel).copied().unwrap_or(0) + != self.versions_seen.get(channel).copied().unwrap_or(0) + } } impl std::fmt::Debug for NodeContext { @@ -125,6 +327,12 @@ impl std::fmt::Debug for NodeContext { .field("recursion_frames", &self.recursion_frames) .field("has_child_runs", &self.child_runs.is_some()) .field("has_agent_binding", &self.agent_binding.is_some()) + .field("task_id", &self.task_id) + .field("siblings", &self.siblings) + .field("channel_versions", &self.channel_versions) + .field("versions_seen", &self.versions_seen) + .field("idle_clock", &self.idle_clock) + .field("durable_writes", &self.lock_durable_writes().len()) .finish() } } @@ -155,17 +363,57 @@ pub(crate) struct NodeMeta { pub(crate) metadata: BTreeMap, } -/// A compiled-in node: id plus its handler. +/// Encodes an `Update` for checkpoint storage. +pub(crate) type EncodeUpdateFn = + dyn Fn(&Update) -> serde_json::Result + Send + Sync; +/// Decodes a stored value back into an `Update`. +pub(crate) type DecodeUpdateFn = + dyn Fn(serde_json::Value) -> serde_json::Result + Send + Sync; + +/// A type-erased `Update` serde codec, installed by +/// [`GraphBuilder::interrupt_after`](super::GraphBuilder::interrupt_after) +/// (the one builder entry point that requires `Update: Serialize + +/// DeserializeOwned`) so the executor can persist a node's deferred result +/// in the checkpoint write ledger and replay it on resume, while +/// [`GraphBuilder`]/[`crate::CompiledGraph`] themselves stay bound-free +/// over `Update`. Same pattern as +/// [`crate::CompiledGraph::with_cached_node`]. +pub(crate) struct UpdateCodec { + pub(crate) encode: Arc>, + pub(crate) decode: Arc>, +} + +impl UpdateCodec +where + Update: serde::Serialize + serde::de::DeserializeOwned + 'static, +{ + /// Builds the serde codec for `Update`. + pub(crate) fn serde() -> Self { + Self { + encode: Arc::new(|update: &Update| serde_json::to_value(update)), + decode: Arc::new(|value: serde_json::Value| serde_json::from_value(value)), + } + } +} + +impl Clone for UpdateCodec { + fn clone(&self) -> Self { + Self { + encode: self.encode.clone(), + decode: self.decode.clone(), + } + } +} + +/// A compiled-in node: its handler. The node's id lives as the key of the +/// `nodes` map it is stored in ([`crate::compiled::CompiledGraph::nodes`]). pub(crate) struct BuilderNode { - #[allow(dead_code)] - pub(crate) id: NodeId, pub(crate) handler: Arc>, } impl Clone for BuilderNode { fn clone(&self) -> Self { Self { - id: self.id.clone(), handler: self.handler.clone(), } } @@ -199,6 +447,24 @@ impl std::fmt::Display for Route { } } +impl From for String { + fn from(route: Route) -> Self { + route.0 + } +} + +impl From for Route { + fn from(label: String) -> Self { + Self(label) + } +} + +impl From<&str> for Route { + fn from(label: &str) -> Self { + Self(label.to_string()) + } +} + /// Tunable per-graph defaults applied to a [`GraphBuilder`] in one call via /// [`GraphBuilder::set_defaults`]. /// @@ -248,8 +514,18 @@ pub struct GraphBuilder { /// (the `graph_id` remains the stable identifier). pub(crate) name: Option, pub(crate) nodes: HashMap>, - pub(crate) edges: HashMap, + /// Static/waiting edges: source node -> its ordered, deduplicated list of + /// successor targets. A node may have more than one static successor + /// (fan-out): every target in the list activates, not just one. + pub(crate) edges: HashMap>, pub(crate) branches: HashMap>, + /// Exhaustive route-label declarations registered via + /// [`super::GraphBuilder::add_conditional_edges_checked`]: node -> every + /// label its router may produce. [`super::GraphBuilder::validate_routes`] + /// cross-checks these against the node's actual route table at build + /// time, catching a typo'd label before it can fail a run with + /// [`crate::TinyAgentsError::MissingRoute`]. + pub(crate) route_label_checks: HashMap>, pub(crate) command_nodes: HashSet, /// Barrier/waiting edges: target node -> set of predecessor nodes that must /// all have completed (across steps) before the target activates. @@ -267,4 +543,19 @@ pub struct GraphBuilder { pub(crate) node_timeout: Option, /// Behavior-free per-node markers/metadata surfaced by the topology export. pub(crate) node_meta: HashMap, + /// Per-node execution policies (retry/timeout/cache/on_error/defer); see + /// [`super::NodePolicy`]. + pub(crate) node_policies: HashMap>, + /// Graph-wide default execution policy every node falls back to, field + /// by field, when it has no per-node override. + pub(crate) node_defaults: Option>, + /// Nodes the executor pauses *before* running (see + /// [`super::GraphBuilder::interrupt_before`]). + pub(crate) interrupt_before: HashSet, + /// Nodes the executor pauses *after* running, before applying their + /// result (see [`super::GraphBuilder::interrupt_after`]). + pub(crate) interrupt_after: HashSet, + /// The `Update` codec `interrupt_after` needs to persist a deferred + /// result; `None` until the first `interrupt_after` call. + pub(crate) update_codec: Option>, } diff --git a/crates/tinyagents-graph/src/cache/memory.rs b/crates/tinyagents-graph/src/cache/memory.rs new file mode 100644 index 00000000..066a6ff9 --- /dev/null +++ b/crates/tinyagents-graph/src/cache/memory.rs @@ -0,0 +1,66 @@ +//! Process-local, TTL-aware [`TaskCache`]. + +use std::collections::HashMap; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; + +use super::{TaskCache, TaskCacheKey}; +use crate::{Result, TinyAgentsError}; +use tinyagents_harness::ids::GraphId; + +/// One stored entry: the value plus its absolute expiry (if any). +type Entry = (serde_json::Value, Option); + +/// An in-memory [`TaskCache`]: a mutex-guarded map honoring per-entry TTL +/// on read. Expired entries are dropped lazily when read; `clear` (or +/// dropping the cache) is what otherwise reclaims them. +#[derive(Default)] +pub struct InMemoryTaskCache { + entries: Mutex>, +} + +impl InMemoryTaskCache { + /// Creates an empty cache. + pub fn new() -> Self { + Self::default() + } + + fn lock(&self) -> Result>> { + self.entries + .lock() + .map_err(|_| TinyAgentsError::Graph("task cache lock poisoned".to_string())) + } +} + +#[async_trait] +impl TaskCache for InMemoryTaskCache { + async fn get(&self, key: &TaskCacheKey) -> Result> { + let mut entries = self.lock()?; + match entries.get(key) { + Some((_, Some(expires_at))) if *expires_at <= Instant::now() => { + entries.remove(key); + Ok(None) + } + Some((value, _)) => Ok(Some(value.clone())), + None => Ok(None), + } + } + + async fn put( + &self, + key: &TaskCacheKey, + value: serde_json::Value, + ttl: Option, + ) -> Result<()> { + let expires_at = ttl.map(|ttl| Instant::now() + ttl); + self.lock()?.insert(key.clone(), (value, expires_at)); + Ok(()) + } + + async fn clear(&self, graph_id: &GraphId) -> Result<()> { + self.lock()?.retain(|key, _| &key.graph_id != graph_id); + Ok(()) + } +} diff --git a/crates/tinyagents-graph/src/cache/mod.rs b/crates/tinyagents-graph/src/cache/mod.rs new file mode 100644 index 00000000..7ebc60c5 --- /dev/null +++ b/crates/tinyagents-graph/src/cache/mod.rs @@ -0,0 +1,26 @@ +//! Task-level result caching for graph nodes. +//! +//! A [`TaskCache`] stores a node's serialized `Update` under a +//! [`TaskCacheKey`] (graph, node, and a caller-computed hash of the inputs) +//! so a later activation with the same key can skip the handler entirely. +//! Caching is opt-in per node through +//! [`NodeCachePolicy`](crate::NodeCachePolicy) and wired into a graph with +//! [`CompiledGraph::with_task_cache`](crate::CompiledGraph::with_task_cache) +//! plus [`CompiledGraph::with_cached_node`](crate::CompiledGraph::with_cached_node). +//! +//! Two backends ship here: [`InMemoryTaskCache`] (process-local, TTL-aware) +//! and, behind the `sqlite` feature, [`SqliteTaskCache`]. + +mod memory; +#[cfg(feature = "sqlite")] +mod sqlite; +mod types; + +pub use memory::InMemoryTaskCache; +#[cfg(feature = "sqlite")] +pub use sqlite::SqliteTaskCache; +pub(crate) use types::CachedNode; +pub use types::{TaskCache, TaskCacheKey}; + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-graph/src/cache/sqlite.rs b/crates/tinyagents-graph/src/cache/sqlite.rs new file mode 100644 index 00000000..e5a0182e --- /dev/null +++ b/crates/tinyagents-graph/src/cache/sqlite.rs @@ -0,0 +1,164 @@ +//! SQLite-backed [`TaskCache`] behind the optional `sqlite` cargo feature. +//! +//! One row per [`TaskCacheKey`] in a `task_cache` table; `put` upserts, and +//! `get` reports a row whose `expires_at` (unix millis) has passed as a miss +//! and deletes it. The connection is opened with the same pragmas as +//! [`crate::SqliteCheckpointer`] (WAL, `synchronous = NORMAL`, busy +//! timeout) via the shared `prepare_connection` helper, and every call runs +//! on `spawn_blocking` so the executor is never blocked on disk I/O. + +use std::path::Path; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use async_trait::async_trait; +use rusqlite::{Connection, OptionalExtension, params}; + +use super::{TaskCache, TaskCacheKey}; +use crate::checkpoint::prepare_connection; +use crate::{Result, TinyAgentsError}; +use tinyagents_harness::ids::GraphId; + +const SCHEMA: &str = " +CREATE TABLE IF NOT EXISTS task_cache ( + graph_id TEXT NOT NULL, + node_id TEXT NOT NULL, + hash TEXT NOT NULL, + value TEXT NOT NULL, + expires_at INTEGER, + PRIMARY KEY (graph_id, node_id, hash) +); +"; + +/// A [`TaskCache`] persisted in a SQLite database. Cheap to clone; clones +/// share one connection (and so the same data, including for `:memory:`). +#[derive(Clone)] +pub struct SqliteTaskCache { + conn: Arc>, +} + +fn sqlite_err(context: &str, err: impl std::fmt::Display) -> TinyAgentsError { + TinyAgentsError::Graph(format!("sqlite task cache: {context}: {err}")) +} + +impl SqliteTaskCache { + /// Opens (creating if needed) a cache database at `path`. + pub fn open(path: impl AsRef) -> Result { + let conn = Connection::open(path.as_ref()).map_err(|e| sqlite_err("open database", e))?; + Self::from_connection(conn) + } + + /// Opens an ephemeral in-memory cache. + pub fn in_memory() -> Result { + let conn = Connection::open_in_memory().map_err(|e| sqlite_err("open in-memory", e))?; + Self::from_connection(conn) + } + + /// Wraps a caller-owned [`Connection`], applying the shared pragmas and + /// ensuring the (idempotent) schema exists. + pub fn from_connection(conn: Connection) -> Result { + prepare_connection(&conn)?; + conn.execute_batch(SCHEMA) + .map_err(|e| sqlite_err("create schema", e))?; + Ok(Self { + conn: Arc::new(Mutex::new(conn)), + }) + } + + /// Runs `f` against the locked connection on the blocking pool. + async fn with_conn(&self, context: &'static str, f: F) -> Result + where + T: Send + 'static, + F: FnOnce(&Connection) -> Result + Send + 'static, + { + let conn = self.conn.clone(); + tokio::task::spawn_blocking(move || { + let conn = conn + .lock() + .map_err(|_| sqlite_err(context, "connection lock poisoned"))?; + f(&conn) + }) + .await + .map_err(|e| sqlite_err(context, e))? + } +} + +fn now_ms() -> i64 { + tinyagents_harness::ids::now_ms() as i64 +} + +#[async_trait] +impl TaskCache for SqliteTaskCache { + async fn get(&self, key: &TaskCacheKey) -> Result> { + let key = key.clone(); + self.with_conn("get", move |conn| { + let row: Option<(String, Option)> = conn + .query_row( + "SELECT value, expires_at FROM task_cache + WHERE graph_id = ?1 AND node_id = ?2 AND hash = ?3", + params![key.graph_id.as_str(), key.node_id.as_str(), key.hash], + |row| Ok((row.get(0)?, row.get(1)?)), + ) + .optional() + .map_err(|e| sqlite_err("read task_cache", e))?; + let Some((value, expires_at)) = row else { + return Ok(None); + }; + if expires_at.is_some_and(|at| at <= now_ms()) { + conn.execute( + "DELETE FROM task_cache WHERE graph_id = ?1 AND node_id = ?2 AND hash = ?3", + params![key.graph_id.as_str(), key.node_id.as_str(), key.hash], + ) + .map_err(|e| sqlite_err("delete expired task_cache row", e))?; + return Ok(None); + } + serde_json::from_str(&value) + .map(Some) + .map_err(|e| sqlite_err("decode cached value", e)) + }) + .await + } + + async fn put( + &self, + key: &TaskCacheKey, + value: serde_json::Value, + ttl: Option, + ) -> Result<()> { + let key = key.clone(); + let expires_at = ttl.map(|ttl| now_ms().saturating_add(ttl.as_millis() as i64)); + self.with_conn("put", move |conn| { + let encoded = + serde_json::to_string(&value).map_err(|e| sqlite_err("encode cached value", e))?; + conn.execute( + "INSERT INTO task_cache (graph_id, node_id, hash, value, expires_at) + VALUES (?1, ?2, ?3, ?4, ?5) + ON CONFLICT(graph_id, node_id, hash) + DO UPDATE SET value = excluded.value, expires_at = excluded.expires_at", + params![ + key.graph_id.as_str(), + key.node_id.as_str(), + key.hash, + encoded, + expires_at + ], + ) + .map_err(|e| sqlite_err("upsert task_cache", e))?; + Ok(()) + }) + .await + } + + async fn clear(&self, graph_id: &GraphId) -> Result<()> { + let graph_id = graph_id.as_str().to_string(); + self.with_conn("clear", move |conn| { + conn.execute( + "DELETE FROM task_cache WHERE graph_id = ?1", + params![graph_id], + ) + .map_err(|e| sqlite_err("clear task_cache", e))?; + Ok(()) + }) + .await + } +} diff --git a/crates/tinyagents-graph/src/cache/test.rs b/crates/tinyagents-graph/src/cache/test.rs new file mode 100644 index 00000000..0ca73572 --- /dev/null +++ b/crates/tinyagents-graph/src/cache/test.rs @@ -0,0 +1,122 @@ +//! Unit tests for the task cache backends: in-memory round trip, TTL +//! expiry, per-graph `clear`, and (behind `sqlite`) the SQLite backend. + +use super::*; +use serde_json::json; +use std::time::Duration; +use tinyagents_harness::ids::{GraphId, NodeId}; + +fn key(graph: &str, node: &str, hash: &str) -> TaskCacheKey { + TaskCacheKey::new(GraphId::new(graph), NodeId::from(node), hash) +} + +#[tokio::test] +async fn in_memory_put_then_get_round_trips() { + let cache = InMemoryTaskCache::new(); + let k = key("g", "n", "h1"); + assert_eq!(cache.get(&k).await.unwrap(), None); + cache.put(&k, json!({"v": 1}), None).await.unwrap(); + assert_eq!(cache.get(&k).await.unwrap(), Some(json!({"v": 1}))); + // A different hash is a different entry. + assert_eq!(cache.get(&key("g", "n", "h2")).await.unwrap(), None); +} + +#[tokio::test] +async fn in_memory_entry_expires_after_ttl() { + let cache = InMemoryTaskCache::new(); + let k = key("g", "n", "h"); + cache + .put(&k, json!(1), Some(Duration::from_millis(30))) + .await + .unwrap(); + assert_eq!(cache.get(&k).await.unwrap(), Some(json!(1))); + tokio::time::sleep(Duration::from_millis(60)).await; + assert_eq!(cache.get(&k).await.unwrap(), None, "expired after TTL"); +} + +#[tokio::test] +async fn in_memory_clear_only_drops_one_graph() { + let cache = InMemoryTaskCache::new(); + cache + .put(&key("a", "n", "h"), json!(1), None) + .await + .unwrap(); + cache + .put(&key("b", "n", "h"), json!(2), None) + .await + .unwrap(); + cache.clear(&GraphId::new("a")).await.unwrap(); + assert_eq!(cache.get(&key("a", "n", "h")).await.unwrap(), None); + assert_eq!( + cache.get(&key("b", "n", "h")).await.unwrap(), + Some(json!(2)) + ); +} + +#[cfg(feature = "sqlite")] +mod sqlite_backend { + use super::*; + + #[tokio::test] + async fn sqlite_put_then_get_round_trips() { + let cache = SqliteTaskCache::in_memory().unwrap(); + let k = key("g", "n", "h1"); + assert_eq!(cache.get(&k).await.unwrap(), None); + cache.put(&k, json!({"v": [1, 2]}), None).await.unwrap(); + assert_eq!(cache.get(&k).await.unwrap(), Some(json!({"v": [1, 2]}))); + // Overwriting the same key replaces the value. + cache.put(&k, json!("new"), None).await.unwrap(); + assert_eq!(cache.get(&k).await.unwrap(), Some(json!("new"))); + assert_eq!(cache.get(&key("g", "n", "h2")).await.unwrap(), None); + } + + #[tokio::test] + async fn sqlite_entry_expires_after_ttl() { + let cache = SqliteTaskCache::in_memory().unwrap(); + let k = key("g", "n", "h"); + cache + .put(&k, json!(1), Some(Duration::from_millis(30))) + .await + .unwrap(); + assert_eq!(cache.get(&k).await.unwrap(), Some(json!(1))); + tokio::time::sleep(Duration::from_millis(60)).await; + assert_eq!(cache.get(&k).await.unwrap(), None, "expired after TTL"); + } + + #[tokio::test] + async fn sqlite_clear_only_drops_one_graph() { + let cache = SqliteTaskCache::in_memory().unwrap(); + cache + .put(&key("a", "n", "h"), json!(1), None) + .await + .unwrap(); + cache + .put(&key("b", "n", "h"), json!(2), None) + .await + .unwrap(); + cache.clear(&GraphId::new("a")).await.unwrap(); + assert_eq!(cache.get(&key("a", "n", "h")).await.unwrap(), None); + assert_eq!( + cache.get(&key("b", "n", "h")).await.unwrap(), + Some(json!(2)) + ); + } + + #[tokio::test] + async fn sqlite_file_backed_cache_persists_across_handles() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("cache.db"); + { + let cache = SqliteTaskCache::open(&path).unwrap(); + cache + .put(&key("g", "n", "h"), json!(7), None) + .await + .unwrap(); + } + let reopened = SqliteTaskCache::open(&path).unwrap(); + assert_eq!( + reopened.get(&key("g", "n", "h")).await.unwrap(), + Some(json!(7)) + ); + } +} diff --git a/crates/tinyagents-graph/src/cache/types.rs b/crates/tinyagents-graph/src/cache/types.rs new file mode 100644 index 00000000..8c919987 --- /dev/null +++ b/crates/tinyagents-graph/src/cache/types.rs @@ -0,0 +1,91 @@ +//! Task cache trait and key type. + +use std::time::Duration; + +use async_trait::async_trait; + +use crate::Result; +use tinyagents_harness::ids::{GraphId, NodeId}; + +/// Identifies one cached task result: which graph and node produced it, +/// and the caller-computed `hash` of the inputs it was computed from +/// (see [`crate::NodeCachePolicy::key`]). +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct TaskCacheKey { + /// The graph the node belongs to. + pub graph_id: GraphId, + /// The node whose result is cached. + pub node_id: NodeId, + /// The input hash the cache policy's key function produced. + pub hash: String, +} + +impl TaskCacheKey { + /// Builds a key from its three parts. + pub fn new(graph_id: GraphId, node_id: NodeId, hash: impl Into) -> Self { + Self { + graph_id, + node_id, + hash: hash.into(), + } + } +} + +/// A store for cached node results, keyed by [`TaskCacheKey`]. +/// +/// Implementations must honor `ttl` on [`Self::put`]: an entry older than +/// its TTL is reported as a miss by [`Self::get`]. The executor treats every +/// cache error as a miss (reads) or logs and continues (writes) — caching +/// is an optimization, never a correctness requirement. +#[async_trait] +pub trait TaskCache: Send + Sync { + /// Looks up a live (non-expired) entry. + async fn get(&self, key: &TaskCacheKey) -> Result>; + /// Stores `value` under `key`, expiring after `ttl` when given. + async fn put( + &self, + key: &TaskCacheKey, + value: serde_json::Value, + ttl: Option, + ) -> Result<()>; + /// Drops every entry belonging to `graph_id`. + async fn clear(&self, graph_id: &GraphId) -> Result<()>; +} + +/// Encodes an `Update` for cache storage. +type EncodeFn = dyn Fn(&Update) -> serde_json::Result + Send + Sync; +/// Decodes a stored value back into an `Update`. +type DecodeFn = dyn Fn(serde_json::Value) -> serde_json::Result + Send + Sync; + +/// One node's resolved cache policy plus the type-erased `Update` codec, +/// installed by [`crate::CompiledGraph::with_cached_node`]. +/// +/// [`crate::NodeCachePolicy`] itself is bound-free over `Update` (it only +/// ever touches `State`); actually storing a cached value needs +/// `Update: Serialize + DeserializeOwned`, which `with_cached_node` supplies +/// locally when it builds `encode`/`decode` — this struct carries the +/// resulting closures rather than the bound itself, so [`CompiledGraph`] +/// stays generic over any `Update`. +/// +/// [`CompiledGraph`]: crate::CompiledGraph +pub(crate) struct CachedNode { + /// Derives the cache key for one activation (state + optional send arg). + pub(crate) key: std::sync::Arc>, + /// Optional time-to-live for a cached entry. + pub(crate) ttl: Option, + /// Encodes an `Update` for storage. + pub(crate) encode: std::sync::Arc>, + /// Decodes a stored value back into an `Update`. + pub(crate) decode: std::sync::Arc>, +} + +impl Clone for CachedNode { + fn clone(&self) -> Self { + Self { + key: self.key.clone(), + ttl: self.ttl, + encode: self.encode.clone(), + decode: self.decode.clone(), + } + } +} diff --git a/crates/tinyagents-graph/src/channel/mod.rs b/crates/tinyagents-graph/src/channel/mod.rs index 75a962cd..183d4df9 100644 --- a/crates/tinyagents-graph/src/channel/mod.rs +++ b/crates/tinyagents-graph/src/channel/mod.rs @@ -36,16 +36,19 @@ //! ephemeral clearing) — so existing whole-state habits keep working and //! conflict detection is strictly opt-in. +mod registry; mod types; +pub use registry::ReducerRegistry; pub use types::{ - Barrier, BinaryAggregate, Channel, ChannelSet, ChannelState, ChannelUpdate, Delta, Ephemeral, - LastValue, Messages, NamedBarrier, Topic, Untracked, + Barrier, BinaryAggregate, Channel, ChannelSet, ChannelState, ChannelUpdate, ChannelWrite, + Delta, Ephemeral, LastValue, Messages, NamedBarrier, Topic, Untracked, }; use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; +use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::reducer::StateReducer; @@ -249,6 +252,10 @@ impl Channel for Barrier { Ok(Value::Array(list)) } + fn config(&self) -> Value { + serde_json::json!({ "expected": self.expected }) + } + fn allows_concurrent(&self) -> bool { true } @@ -301,6 +308,10 @@ impl Channel for NamedBarrier { Ok(Value::Object(map)) } + fn config(&self) -> Value { + serde_json::json!({ "expected": self.expected }) + } + fn allows_concurrent(&self) -> bool { true } @@ -320,22 +331,51 @@ impl Channel for NamedBarrier { impl BinaryAggregate { /// Creates an aggregate channel from a binary fold closure. The first write /// becomes the value directly; later writes are `fold(current, incoming)`. + /// + /// Unnamed: [`Channel::config`] carries no reducer name, so a channel + /// built this way merges correctly at runtime but cannot round-trip + /// through a durable checkpointer. Use [`BinaryAggregate::named`] (backed + /// by [`ReducerRegistry`]) for a channel that must survive a checkpoint + /// decode. pub fn new(fold: F) -> Self where F: Fn(Value, Value) -> Result + Send + Sync + 'static, { Self { fold: Arc::new(fold), + reducer_name: None, } } - /// Builds an aggregate channel from a [`crate::Reducer`]. + /// Builds an aggregate channel from a [`crate::Reducer`]. Also + /// unnamed — see [`BinaryAggregate::new`]. pub fn from_reducer(reducer: R) -> Self where R: crate::Reducer + 'static, { Self::new(move |current, incoming| reducer.reduce(current, incoming)) } + + /// Builds an aggregate channel from the reducer registered under `name` + /// in the process-wide [`ReducerRegistry`] (register it first with + /// [`crate::GraphBuilder::register_reducer`], or use one of the built-ins + /// — `"append"`, `"last"`, `"sum"`, `"max"`, `"min"`, `"set_union"`). + /// + /// Unlike [`BinaryAggregate::new`], this channel's [`Channel::config`] + /// persists `name`, so it round-trips through a durable checkpointer: + /// decoding looks `name` back up in the registry (present in the + /// resuming process — the same call site that ran this graph before must + /// have registered it) and fails with + /// `TinyAgentsError::Checkpoint("unknown reducer ...")` if it is not + /// there. + pub fn named(name: impl Into) -> Result { + let name = name.into(); + let fold = ReducerRegistry::require(&name)?; + Ok(Self { + fold, + reducer_name: Some(name), + }) + } } impl Channel for BinaryAggregate { @@ -350,6 +390,10 @@ impl Channel for BinaryAggregate { } } + fn config(&self) -> Value { + ReducerRegistry::config_for(self.reducer_name.as_deref()) + } + fn allows_concurrent(&self) -> bool { true } @@ -359,6 +403,56 @@ impl Channel for BinaryAggregate { } } +/// Reconstructs a boxed [`Channel`] from its persisted `{kind, config}` pair +/// (the counterpart of [`Channel::config`]), used by [`ChannelSet`]'s +/// [`serde::Deserialize`] impl to hydrate a checkpoint's channel schema with +/// no external context — see `channel/registry.rs`'s module docs for why +/// `binary_aggregate` alone needs the process-wide [`ReducerRegistry`] to do +/// this. +fn channel_from_config(kind: &str, config: &Value) -> Result> { + match kind { + "last_value" => Ok(Box::new(LastValue)), + "topic" => Ok(Box::new(Topic)), + "delta" => Ok(Box::new(Delta)), + "messages" => Ok(Box::new(Messages)), + "ephemeral" => Ok(Box::new(Ephemeral)), + "untracked" => Ok(Box::new(Untracked)), + "barrier" => { + let expected = config.get("expected").and_then(Value::as_u64).unwrap_or(0) as usize; + Ok(Box::new(Barrier::new(expected))) + } + "named_barrier" => { + let expected: Vec = config + .get("expected") + .and_then(Value::as_array) + .map(|items| { + items + .iter() + .filter_map(|v| v.as_str().map(str::to_string)) + .collect() + }) + .unwrap_or_default(); + Ok(Box::new(NamedBarrier::new(expected))) + } + "binary_aggregate" => { + let name = config + .get("reducer") + .and_then(Value::as_str) + .ok_or_else(|| { + TinyAgentsError::Checkpoint( + "binary_aggregate channel requires a named reducer to decode; build it \ + with `BinaryAggregate::named` so its config persists a reducer name" + .to_string(), + ) + })?; + Ok(Box::new(BinaryAggregate::named(name)?)) + } + other => Err(TinyAgentsError::Checkpoint(format!( + "unknown channel kind `{other}`" + ))), + } +} + // --- ChannelSet --- impl ChannelSet { @@ -382,6 +476,27 @@ impl ChannelSet { self.channels.insert(name.into(), Box::new(channel)); } + /// Marks an already-registered append-style channel (typically [`Topic`] + /// or a `"append"`/`"set_union"` [`BinaryAggregate`]) for delta-history + /// tracking: every write to `name` also records the raw incoming value + /// into [`ChannelState::step_deltas`] for that step, which the + /// checkpoint-construction call sites persist into + /// [`crate::checkpoint::Checkpoint::channel_deltas`]. Every + /// `snapshot_every` writes (minimum `1`) an additional full-value + /// snapshot marker (`{"$snapshot": }`) is recorded alongside the + /// delta, so a consumer walking the history can fast-forward without + /// replaying every write from genesis. + /// + /// Returns the set for chaining. A no-op marker on a channel name that + /// is never registered with [`ChannelSet::with_channel`]/ + /// [`ChannelSet::add_channel`] has no effect (there is nothing to track + /// writes for). + pub fn with_delta(mut self, name: impl Into, snapshot_every: u32) -> Self { + self.delta_channels + .insert(name.into(), snapshot_every.max(1)); + self + } + /// Returns the current value of `name`, if any has been written. pub fn get(&self, name: &str) -> Option<&Value> { self.values.get(name) @@ -430,6 +545,35 @@ impl ChannelSet { Ok(()) } + /// The single dispatch point for one channel write, folding an ordinary + /// [`ChannelWrite::Merge`] through [`ChannelSet::apply_update`] or + /// replacing the value outright for a [`ChannelWrite::Overwrite`] (which + /// bypasses the channel's merge rule and becomes the new baseline for + /// any merge/delta tracking that follows). Returns the channel's value + /// after the write. + /// + /// This is the one write path every channel-graph write funnels through + /// — a normal executor superstep boundary + /// ([`ChannelState::merge`]/[`crate::channel::ChannelUpdate`]), + /// `CompiledGraph::update_state`, and `CompiledGraph::fork_state`'s copy + /// — so replay and a manual update can never disagree about what a + /// write means (I5/R3; see `docs/modules/graph/state-channels.md`). + pub fn apply_channel_write(&mut self, name: &str, write: &ChannelWrite) -> Result { + match write { + ChannelWrite::Merge(value) => { + self.apply_update(name, value.clone())?; + Ok(self.values.get(name).cloned().unwrap_or(Value::Null)) + } + ChannelWrite::Overwrite(value) => { + // Validate the channel exists (same contract as `apply_update`) + // before mutating. + self.channel(name)?; + self.values.insert(name.to_string(), value.clone()); + Ok(value.clone()) + } + } + } + /// Returns the tracked channel values as an ordered map, excluding /// [`Untracked`] channels. This is the durable/inspectable state view. pub fn snapshot(&self) -> BTreeMap { @@ -467,6 +611,80 @@ impl ChannelSet { } } +/// One channel's wire representation: `{ kind, config, value }` (the +/// counterpart of [`Channel::config`]/[`channel_from_config`]). +#[derive(Serialize, Deserialize)] +struct ChannelEntry { + kind: String, + #[serde(default, skip_serializing_if = "Value::is_null")] + config: Value, + #[serde(default, skip_serializing_if = "Option::is_none")] + value: Option, +} + +/// [`ChannelSet`]'s full wire representation: its channel schema/values plus +/// the [`ChannelSet::with_delta`] registrations, so a decoded set round-trips +/// which channels are delta-tracked (not just their current values). +#[derive(Serialize, Deserialize)] +struct ChannelSetWire { + channels: BTreeMap, + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + delta: HashMap, +} + +impl serde::Serialize for ChannelSet { + fn serialize( + &self, + serializer: S, + ) -> std::result::Result { + let channels: BTreeMap = self + .channels + .iter() + .map(|(name, channel)| { + ( + name.clone(), + ChannelEntry { + kind: channel.kind().to_string(), + config: channel.config(), + value: channel + .is_tracked() + .then(|| self.values.get(name).cloned()) + .flatten(), + }, + ) + }) + .collect(); + ChannelSetWire { + channels, + delta: self.delta_channels.clone(), + } + .serialize(serializer) + } +} + +impl<'de> serde::Deserialize<'de> for ChannelSet { + fn deserialize>( + deserializer: D, + ) -> std::result::Result { + let wire = ChannelSetWire::deserialize(deserializer)?; + let mut channels: HashMap> = HashMap::new(); + let mut values: HashMap = HashMap::new(); + for (name, entry) in wire.channels { + let channel = channel_from_config(&entry.kind, &entry.config) + .map_err(serde::de::Error::custom)?; + channels.insert(name.clone(), channel); + if let Some(value) = entry.value { + values.insert(name, value); + } + } + Ok(ChannelSet { + channels, + values, + delta_channels: wire.delta, + }) + } +} + // --- ChannelUpdate --- impl ChannelUpdate { @@ -475,9 +693,19 @@ impl ChannelUpdate { Self::default() } - /// Adds a `(name, value)` write, returning the update for chaining. + /// Adds a `(name, value)` merged write, returning the update for + /// chaining. pub fn set(mut self, name: impl Into, value: impl Into) -> Self { - self.writes.push((name.into(), value.into())); + self.writes + .push((name.into(), ChannelWrite::Merge(value.into()))); + self + } + + /// Adds a `(name, value)` write that bypasses the channel's merge rule + /// and replaces its value outright (see [`ChannelWrite::Overwrite`]). + pub fn overwrite(mut self, name: impl Into, value: impl Into) -> Self { + self.writes + .push((name.into(), ChannelWrite::Overwrite(value.into()))); self } @@ -514,6 +742,13 @@ impl ChannelState { self } + /// Marks an already-registered channel for delta-history tracking; see + /// [`ChannelSet::with_delta`]. + pub fn with_delta(mut self, name: impl Into, snapshot_every: u32) -> Self { + self.set = self.set.with_delta(name, snapshot_every); + self + } + /// Borrows the underlying [`ChannelSet`]. pub fn channels(&self) -> &ChannelSet { &self.set @@ -548,12 +783,14 @@ impl ChannelState { Some(step) if step != self.current_step => { self.current_step = step; self.step_writes.clear(); + self.step_deltas.clear(); self.set.clear_ephemeral(); } Some(_) => {} None => { // Unstamped updates are independent: no cross-update detection. self.step_writes.clear(); + self.step_deltas.clear(); } } @@ -579,14 +816,88 @@ impl ChannelState { } let touched: HashSet = distinct.iter().map(|n| n.to_string()).collect(); - for name in touched { - *self.step_writes.entry(name).or_insert(0) += 1; + for name in &touched { + *self.step_writes.entry(name.clone()).or_insert(0) += 1; + // Channel versions (I5/R3): bumped once per distinct channel + // name touched by this update, regardless of write kind + // (Merge/Overwrite) — see the module docs on + // `Checkpoint::channel_versions`. + *self.channel_versions.entry(name.clone()).or_insert(0) += 1; } - for (name, value) in update.writes { - self.set.apply_update(&name, value)?; + // `apply_channel_write` (on `ChannelSet`) is the single write-path + // dispatch point every channel-graph write funnels through — see its + // docs. This loop is that path's boundary-fold caller; `update_state` + // and `fork_state` reach the same dispatch point through + // `compiled::channel_bookkeeping`/direct `ChannelSet` access so + // replay and a manual write cannot diverge. + for (name, write) in update.writes { + let is_overwrite = write.is_overwrite(); + let value = write.value().clone(); + self.set.apply_channel_write(&name, &write)?; + if let Some(&snapshot_every) = self.set.delta_channels.get(&name) { + let entry = self.step_deltas.entry(name.clone()).or_default(); + if is_overwrite { + // Overwrite rebases the delta/append history: prior + // accumulated deltas for this channel no longer describe + // the current baseline. + entry.clear(); + } + entry.push(value); + let version = self.channel_versions.get(&name).copied().unwrap_or(0); + if snapshot_every > 0 && version % u64::from(snapshot_every) == 0 { + let full = self.set.get(&name).cloned().unwrap_or(Value::Null); + entry.push(serde_json::json!({ "$snapshot": full })); + } + } } Ok(self) } + + /// Cumulative per-channel version counters (I5/R3). See + /// [`crate::checkpoint::Checkpoint::channel_versions`]. + pub fn channel_versions(&self) -> &BTreeMap { + &self.channel_versions + } + + /// This step's accumulated raw write values for every + /// [`ChannelSet::with_delta`]-tracked channel, reset when the stamped + /// step advances. See [`crate::checkpoint::Checkpoint::channel_deltas`]. + pub fn step_deltas(&self) -> &BTreeMap> { + &self.step_deltas + } +} + +/// Extracts `(channel_versions, channel_deltas)` to embed into a freshly +/// built [`crate::checkpoint::Checkpoint`], downcasting `state` to +/// [`ChannelState`] when the graph uses the channel model. +/// +/// For any other `State` type (a plain whole-state graph) a single +/// `"state"` channel is reported at `fallback_version`, with no deltas — see +/// the module docs on [`crate::checkpoint::Checkpoint::channel_versions`]. +/// +/// Shared by every checkpoint-construction call site (the executor's normal/ +/// failure/cancel boundaries in `compiled::boundary`, and +/// `compiled::state_api`'s `update_state`) so a normal superstep boundary +/// and a manual write can never disagree about what they persist here (the +/// "one write path" contract — I5/R3). +pub fn channel_bookkeeping( + state: &State, + fallback_version: u64, +) -> ( + BTreeMap, + BTreeMap>, +) { + match (state as &dyn std::any::Any).downcast_ref::() { + Some(channel_state) => ( + channel_state.channel_versions().clone(), + channel_state.step_deltas().clone(), + ), + None => { + let mut versions = BTreeMap::new(); + versions.insert("state".to_string(), fallback_version); + (versions, BTreeMap::new()) + } + } } /// `ChannelState` is its own [`StateReducer`]: the `&self` receiver is unused diff --git a/crates/tinyagents-graph/src/channel/registry.rs b/crates/tinyagents-graph/src/channel/registry.rs new file mode 100644 index 00000000..de9f0dcb --- /dev/null +++ b/crates/tinyagents-graph/src/channel/registry.rs @@ -0,0 +1,157 @@ +//! The process-wide named reducer registry backing [`crate::BinaryAggregate`] +//! channels. +//! +//! A [`crate::BinaryAggregate`] channel's merge rule is a closure, which is +//! not serializable — but a durable [`crate::Checkpoint`] must be able to +//! decode a `ChannelState` with no context beyond the bytes on disk (it +//! implements plain `DeserializeOwned`, not a seeded deserialize). This +//! registry is the bridge: a reducer is registered once under a stable name +//! (the built-ins below, or a caller's own via +//! [`crate::GraphBuilder::register_reducer`]), the channel's +//! [`crate::Channel::config`] persists only that *name*, and decoding a +//! `binary_aggregate` channel looks the closure back up by name — see +//! [`crate::BinaryAggregate::named`] and `channel_from_config` in `mod.rs`. +//! +//! The registry is global (not scoped to one [`crate::GraphBuilder`]) because +//! decoding happens with no builder in scope at all — only the checkpoint +//! bytes. A name registered anywhere in the process is visible to every +//! decode; decoding a name nobody has registered fails with +//! `TinyAgentsError::Checkpoint("unknown reducer ...")` rather than silently +//! losing the reducer's behavior. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, OnceLock}; + +use serde_json::{Value, json}; + +use crate::{Result, TinyAgentsError}; + +pub(crate) type ReduceFn = Arc Result + Send + Sync>; + +fn numeric_add(a: &Value, b: &Value) -> Result { + let err = || TinyAgentsError::Graph("`sum` reducer requires numeric values".to_string()); + if a.is_i64() && b.is_i64() { + return Ok(Value::from(a.as_i64().unwrap() + b.as_i64().unwrap())); + } + let sum = a.as_f64().ok_or_else(err)? + b.as_f64().ok_or_else(err)?; + Ok(Value::from(sum)) +} + +fn storage() -> &'static Mutex> { + static REGISTRY: OnceLock>> = OnceLock::new(); + REGISTRY.get_or_init(|| { + let mut map: HashMap = HashMap::new(); + map.insert( + "append".to_string(), + Arc::new(|current: Value, incoming: Value| { + let mut list = match current { + Value::Array(items) => items, + other => vec![other], + }; + match incoming { + Value::Array(items) => list.extend(items), + other => list.push(other), + } + Ok(Value::Array(list)) + }), + ); + map.insert( + "last".to_string(), + Arc::new(|_current: Value, incoming: Value| Ok(incoming)), + ); + map.insert( + "sum".to_string(), + Arc::new(|current: Value, incoming: Value| numeric_add(¤t, &incoming)), + ); + map.insert( + "max".to_string(), + Arc::new(|current: Value, incoming: Value| { + Ok(if incoming.as_f64() > current.as_f64() { + incoming + } else { + current + }) + }), + ); + map.insert( + "min".to_string(), + Arc::new(|current: Value, incoming: Value| { + Ok(if incoming.as_f64() < current.as_f64() { + incoming + } else { + current + }) + }), + ); + map.insert( + "set_union".to_string(), + Arc::new(|current: Value, incoming: Value| { + let mut list = match current { + Value::Array(items) => items, + other => vec![other], + }; + let incoming = match incoming { + Value::Array(items) => items, + other => vec![other], + }; + for item in incoming { + if !list.contains(&item) { + list.push(item); + } + } + Ok(Value::Array(list)) + }), + ); + Mutex::new(map) + }) +} + +/// The process-wide named registry of [`crate::BinaryAggregate`] reducer +/// closures. See the module docs for why this exists and why it is global. +/// +/// Pre-registered built-ins: `"append"`, `"last"`, `"sum"`, `"max"`, +/// `"min"`, `"set_union"`. +pub struct ReducerRegistry; + +impl ReducerRegistry { + /// Registers `f` under `name`, overwriting any previous registration of + /// that name (including a built-in). Prefer + /// [`crate::GraphBuilder::register_reducer`], which delegates here. + pub fn register( + name: impl Into, + f: impl Fn(Value, Value) -> Result + Send + Sync + 'static, + ) { + let mut guard = storage() + .lock() + .unwrap_or_else(|poison| poison.into_inner()); + guard.insert(name.into(), Arc::new(f)); + } + + /// Looks up the reducer registered under `name`. + pub(crate) fn get(name: &str) -> Option { + let guard = storage() + .lock() + .unwrap_or_else(|poison| poison.into_inner()); + guard.get(name).cloned() + } + + /// Looks up `name`, producing the standard + /// `TinyAgentsError::Checkpoint("unknown reducer ...")` error a + /// checkpoint decode raises for a name nobody registered. + pub(crate) fn require(name: &str) -> Result { + Self::get(name).ok_or_else(|| { + TinyAgentsError::Checkpoint(format!( + "unknown reducer `{name}`: no closure is registered under this name in this \ + process (register it with GraphBuilder::register_reducer before decoding this \ + checkpoint)" + )) + }) + } + + /// The `{"reducer": name}` config payload for a named `BinaryAggregate` + /// channel, or `{"reducer": null}` for an unnamed one (which + /// [`crate::channel::channel_from_config`] then rejects on decode). + pub(crate) fn config_for(name: Option<&str>) -> Value { + json!({ "reducer": name }) + } +} diff --git a/crates/tinyagents-graph/src/channel/test.rs b/crates/tinyagents-graph/src/channel/test.rs index d8407a50..07c563d1 100644 --- a/crates/tinyagents-graph/src/channel/test.rs +++ b/crates/tinyagents-graph/src/channel/test.rs @@ -373,3 +373,617 @@ fn single_update_repeat_write_is_last_wins() { .unwrap(); assert_eq!(merged.get("v"), Some(&json!(2))); } + +// --- Serializable channels + ReducerRegistry (I5/R3) --- + +#[test] +fn channel_config_round_trips_for_every_built_in_kind() { + let set = ChannelSet::new() + .with_channel("last", LastValue) + .with_channel("topic", Topic) + .with_channel("delta", Delta) + .with_channel("messages", Messages) + .with_channel("barrier", Barrier::new(2)) + .with_channel("named", NamedBarrier::new(["a", "b"])) + .with_channel("agg", BinaryAggregate::named("sum").unwrap()); + let json = serde_json::to_value(&set).unwrap(); + let decoded: ChannelSet = serde_json::from_value(json).unwrap(); + assert!(decoded.contains("last")); + assert!(decoded.contains("topic")); + assert!(decoded.contains("delta")); + assert!(decoded.contains("messages")); + assert!(decoded.contains("barrier")); + assert!(decoded.contains("named")); + assert!(decoded.contains("agg")); + // A decoded barrier keeps its `expected` readiness threshold. + assert!(decoded.allows_concurrent("barrier").unwrap()); +} + +#[test] +fn untracked_channel_value_is_not_serialized() { + let mut set = ChannelSet::new() + .with_channel("durable", LastValue) + .with_channel("scratch", Untracked); + set.apply_update("durable", json!("saved")).unwrap(); + set.apply_update("scratch", json!("discarded")).unwrap(); + + let json = serde_json::to_value(&set).unwrap(); + assert_eq!(json["channels"]["durable"]["value"], json!("saved")); + assert!( + json["channels"]["scratch"].get("value").is_none(), + "untracked values must not enter checkpoint wire data" + ); + + let decoded: ChannelSet = serde_json::from_value(json).unwrap(); + assert_eq!(decoded.get("durable"), Some(&json!("saved"))); + assert_eq!(decoded.get("scratch"), None); +} + +#[test] +fn binary_aggregate_named_unknown_reducer_errors() { + let err = BinaryAggregate::named("does-not-exist-anywhere").unwrap_err(); + assert!(matches!(err, TinyAgentsError::Checkpoint(_))); + assert!(err.to_string().contains("unknown reducer")); +} + +#[test] +fn decoding_binary_aggregate_without_reducer_name_errors() { + // A `BinaryAggregate::new` closure carries no name, so its config is + // `{"reducer": null}` — decoding must fail rather than silently drop the + // merge rule. + let set = ChannelSet::new().with_channel( + "agg", + BinaryAggregate::new(|a: Value, b: Value| { + Ok(json!(a.as_i64().unwrap() + b.as_i64().unwrap())) + }), + ); + let json = serde_json::to_value(&set).unwrap(); + let decoded: std::result::Result = serde_json::from_value(json); + assert!(decoded.is_err()); +} + +#[test] +fn named_binary_aggregate_round_trips_and_merges_after_decode() { + crate::ReducerRegistry::register("channel-test-double", |current: Value, incoming: Value| { + Ok(json!( + current.as_i64().unwrap_or(1) * incoming.as_i64().unwrap() + )) + }); + let set = ChannelSet::new().with_channel( + "product", + BinaryAggregate::named("channel-test-double").unwrap(), + ); + let json = serde_json::to_value(&set).unwrap(); + let mut decoded: ChannelSet = serde_json::from_value(json).unwrap(); + decoded.apply_update("product", json!(3)).unwrap(); + decoded.apply_update("product", json!(4)).unwrap(); + assert_eq!(decoded.get("product"), Some(&json!(12))); +} + +#[test] +fn register_reducer_is_reachable_through_the_graph_builder() { + let _graph = GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .register_reducer("channel-test-via-builder", |_a: Value, b: Value| Ok(b)) + .add_node("noop", |_s: ChannelState, _c: NodeContext| async move { + Ok(NodeResult::Update(ChannelUpdate::new())) + }) + .set_entry("noop") + .set_finish("noop") + .compile() + .unwrap(); + // Registered globally, so `named` can find it without the builder. + assert!(BinaryAggregate::named("channel-test-via-builder").is_ok()); +} + +#[tokio::test] +async fn channel_state_graph_round_trips_through_file_checkpointer() { + use crate::checkpoint::{Checkpointer, FileCheckpointer}; + use crate::command::{Command, Interrupt}; + use std::sync::Arc; + + fn graph() -> crate::compiled::CompiledGraph { + GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .add_node("collect", |_s: ChannelState, c: NodeContext| async move { + match c.resume { + Some(value) => { + let bump = value.get("bump").and_then(Value::as_i64).unwrap_or(0); + Ok(NodeResult::Update( + ChannelUpdate::new() + .set("total", bump) + .set("log", "collected") + .at_step(c.step), + )) + } + None => Ok(NodeResult::Interrupt(Interrupt::new( + "collect", + json!({ "ask": "bump?" }), + ))), + } + }) + .set_entry("collect") + .set_finish("collect") + .compile() + .unwrap() + } + + let initial = ChannelState::new() + .with_channel("total", BinaryAggregate::named("sum").unwrap()) + .with_channel("log", Topic); + + let dir = tempfile::tempdir().unwrap(); + { + let cp: Arc> = + Arc::new(FileCheckpointer::::new(dir.path())); + let g = graph().with_checkpointer(cp); + let paused = g.run_with_thread("ch-thread", initial).await.unwrap(); + assert!(paused.is_interrupted()); + } + // Fresh checkpointer over the same directory, simulating a process + // restart: the `binary_aggregate` channel must decode using the (still + // process-wide registered) `"sum"` reducer and the resumed run must + // merge through it correctly. + { + let cp: Arc> = + Arc::new(FileCheckpointer::::new(dir.path())); + let g = graph().with_checkpointer(cp); + let exec = g + .resume("ch-thread", Command::resume(json!({ "bump": 7 }))) + .await + .unwrap(); + assert_eq!(exec.state.get("total"), Some(&json!(7))); + assert_eq!(exec.state.get("log"), Some(&json!(["collected"]))); + } +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn channel_state_graph_round_trips_through_sqlite_checkpointer() { + use crate::checkpoint::{Checkpointer, SqliteCheckpointer}; + use crate::command::{Command, Interrupt}; + use std::sync::Arc; + + fn graph() -> crate::compiled::CompiledGraph { + GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .add_node("collect", |_s: ChannelState, c: NodeContext| async move { + match c.resume { + Some(value) => { + let bump = value.get("bump").and_then(Value::as_i64).unwrap_or(0); + Ok(NodeResult::Update( + ChannelUpdate::new().set("total", bump).at_step(c.step), + )) + } + None => Ok(NodeResult::Interrupt(Interrupt::new( + "collect", + json!({ "ask": "bump?" }), + ))), + } + }) + .set_entry("collect") + .set_finish("collect") + .compile() + .unwrap() + } + + let initial = ChannelState::new().with_channel("total", BinaryAggregate::named("sum").unwrap()); + + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("checkpoints.sqlite"); + { + let cp: Arc> = + Arc::new(SqliteCheckpointer::::open(&db_path).unwrap()); + let g = graph().with_checkpointer(cp); + let paused = g + .run_with_thread("ch-thread-sqlite", initial) + .await + .unwrap(); + assert!(paused.is_interrupted()); + } + { + let cp: Arc> = + Arc::new(SqliteCheckpointer::::open(&db_path).unwrap()); + let g = graph().with_checkpointer(cp); + let exec = g + .resume("ch-thread-sqlite", Command::resume(json!({ "bump": 9 }))) + .await + .unwrap(); + assert_eq!(exec.state.get("total"), Some(&json!(9))); + } +} + +// --- Channel versions (I5/R3) --- + +#[test] +fn channel_versions_bump_once_per_distinct_channel_per_update() { + let state = ChannelState::new() + .with_channel("a", LastValue) + .with_channel("b", LastValue); + let state = state + .merge( + ChannelUpdate::new() + .set("a", 1) + .set("a", 2) // same update, same channel: still one bump. + .set("b", 1) + .at_step(1), + ) + .unwrap(); + assert_eq!(state.channel_versions().get("a"), Some(&1)); + assert_eq!(state.channel_versions().get("b"), Some(&1)); + + let state = state + .merge(ChannelUpdate::new().set("a", 3).at_step(2)) + .unwrap(); + assert_eq!(state.channel_versions().get("a"), Some(&2)); + // `b` was not touched this step, so its version does not move. + assert_eq!(state.channel_versions().get("b"), Some(&1)); +} + +#[tokio::test] +async fn node_context_changed_since_last_run_tracks_channel_writes() { + use std::sync::{Arc, Mutex}; + + let observed: Arc>> = Arc::new(Mutex::new(Vec::new())); + let observed_a = observed.clone(); + let observed_b = observed.clone(); + + let graph = GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .add_node("write", |_s: ChannelState, c: NodeContext| async move { + Ok(NodeResult::Update( + ChannelUpdate::new().set("v", 1).at_step(c.step), + )) + }) + .add_node("observe_a", move |_s: ChannelState, c: NodeContext| { + let observed_a = observed_a.clone(); + async move { + observed_a + .lock() + .unwrap() + .push(c.changed_since_last_run("v")); + Ok(NodeResult::Update(ChannelUpdate::new())) + } + }) + .add_node("observe_b", move |_s: ChannelState, c: NodeContext| { + let observed_b = observed_b.clone(); + async move { + observed_b + .lock() + .unwrap() + .push(c.changed_since_last_run("v")); + Ok(NodeResult::Update(ChannelUpdate::new())) + } + }) + .add_edge("write", "observe_a") + .add_edge("observe_a", "observe_b") + .add_edge("observe_b", "observe_a") + .set_entry("write") + .set_finish("observe_a") + .with_recursion_limit(4) + .compile() + .unwrap(); + + let initial = ChannelState::new().with_channel("v", LastValue); + let _ = graph.run(initial).await; + // `observe_a`'s first-ever run sees `v` as changed (it has never seen + // it), then its second run (after `observe_b`, which made no writes) + // sees no further change. + let seen_a = observed.lock().unwrap().clone(); + assert_eq!(seen_a.first(), Some(&true)); +} + +// --- Delta-channel history + Overwrite (I5/R3) --- + +#[test] +fn overwrite_bypasses_merge_and_rebases_baseline() { + let state = ChannelState::new().with_channel("log", Topic); + let state = state + .merge(ChannelUpdate::new().set("log", "a").at_step(1)) + .unwrap(); + assert_eq!(state.get("log"), Some(&json!(["a"]))); + + // Overwrite replaces the value outright... + let state = state + .merge( + ChannelUpdate::new() + .overwrite("log", json!(["reset"])) + .at_step(2), + ) + .unwrap(); + assert_eq!(state.get("log"), Some(&json!(["reset"]))); + + // ...and subsequent merges build on the new baseline, not the old one. + let state = state + .merge(ChannelUpdate::new().set("log", "b").at_step(3)) + .unwrap(); + assert_eq!(state.get("log"), Some(&json!(["reset", "b"]))); +} + +#[test] +fn delta_tracked_channel_accumulates_step_deltas_and_overwrite_rebases_them() { + let set = ChannelSet::new() + .with_channel("log", Topic) + .with_delta("log", 1000); + let state = ChannelState::new(); + let state = ChannelState { set, ..state }; + let state = state + .merge(ChannelUpdate::new().set("log", "a").at_step(1)) + .unwrap(); + assert_eq!(state.step_deltas().get("log"), Some(&vec![json!("a")])); + + let state = state + .merge(ChannelUpdate::new().set("log", "b").at_step(2)) + .unwrap(); + // Deltas are per-step only: step 2's entry does not include step 1's. + assert_eq!(state.step_deltas().get("log"), Some(&vec![json!("b")])); + + let state = state + .merge( + ChannelUpdate::new() + .overwrite("log", json!(["reset"])) + .at_step(3), + ) + .unwrap(); + // The overwrite rebases the delta baseline: only the overwrite's own + // value is recorded for this step. + assert_eq!( + state.step_deltas().get("log"), + Some(&vec![json!(["reset"])]) + ); +} + +#[tokio::test] +async fn delta_history_replays_from_checkpoints() { + use crate::checkpoint::{CheckpointConfig, Checkpointer, InMemoryCheckpointer}; + use std::sync::Arc; + + let set = ChannelSet::new() + .with_channel("log", Topic) + .with_delta("log", 1000); + let base = ChannelState { + set, + ..ChannelState::new() + }; + + let cp: Arc> = Arc::new(InMemoryCheckpointer::new()); + let graph = GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .add_node("append", |_s: ChannelState, c: NodeContext| async move { + Ok(NodeResult::Update( + ChannelUpdate::new() + .set("log", format!("item-{}", c.step)) + .at_step(c.step), + )) + }) + .add_edge("append", "append") + .set_entry("append") + .set_finish("append") + .with_recursion_limit(3) + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + // `with_recursion_limit(3)` plus the self-edge lets `append` run three + // supersteps before the limit trips. + let _ = graph.run_with_thread("delta-thread", base).await; + + let config = CheckpointConfig::latest("delta-thread"); + let history = cp.delta_history(&config, "log").await.unwrap(); + assert_eq!( + history, + vec![json!("item-1"), json!("item-2"), json!("item-3")] + ); +} + +/// A long-running delta-tracked append channel's *per-checkpoint* byte size +/// must grow ~linearly with step count (each checkpoint carries only its own +/// step's delta, not a cumulative history) — comparing the whole checkpoint +/// record's serialized size at step 100 vs step 200 bounds the ratio well +/// under the quadratic blowup a cumulative (or naive full-replay) design +/// would produce. +#[tokio::test] +async fn delta_channel_checkpoint_bytes_grow_linearly_over_two_hundred_steps() { + use crate::checkpoint::{Checkpointer, InMemoryCheckpointer}; + use std::sync::Arc; + + let set = ChannelSet::new() + .with_channel("log", Topic) + .with_delta("log", 100_000); // no periodic full-snapshot marker in range. + let base = ChannelState { + set, + ..ChannelState::new() + }; + + let cp: Arc> = Arc::new(InMemoryCheckpointer::new()); + let graph = GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .add_node("append", |_s: ChannelState, c: NodeContext| async move { + Ok(NodeResult::Update( + ChannelUpdate::new() + .set("log", format!("item-{}", c.step)) + .at_step(c.step), + )) + }) + .add_edge("append", "append") + .set_entry("append") + .set_finish("append") + .with_recursion_limit(205) + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let _ = graph.run_with_thread("linear-thread", base).await; + + let metas = cp.list("linear-thread").await.unwrap(); + let id_at = |step: usize| -> String { + metas + .iter() + .find(|m| m.step == step) + .unwrap_or_else(|| panic!("no checkpoint at step {step}")) + .checkpoint_id + .clone() + }; + async fn bytes_at(cp: &Arc>, id: &str) -> usize { + let checkpoint = cp + .get("linear-thread", Some(id)) + .await + .unwrap() + .expect("checkpoint exists"); + serde_json::to_vec(&checkpoint).unwrap().len() + } + let bytes_at_100 = bytes_at(&cp, &id_at(100)).await; + let bytes_at_200 = bytes_at(&cp, &id_at(200)).await; + + let ratio = bytes_at_200 as f64 / bytes_at_100 as f64; + assert!( + ratio < 2.5, + "checkpoint bytes should grow ~linearly (step100={bytes_at_100}, step200={bytes_at_200}, ratio={ratio})" + ); +} + +// --- `update_state`/`fork_state` share the delta/version write path --- + +#[tokio::test] +async fn update_state_after_delta_writes_round_trips() { + use crate::checkpoint::{CheckpointConfig, Checkpointer, InMemoryCheckpointer}; + use std::sync::Arc; + + let set = ChannelSet::new() + .with_channel("log", Topic) + .with_delta("log", 1000); + let base = ChannelState { + set, + ..ChannelState::new() + }; + + let cp: Arc> = Arc::new(InMemoryCheckpointer::new()); + let graph = GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .add_node("append", |_s: ChannelState, c: NodeContext| async move { + Ok(NodeResult::Update( + ChannelUpdate::new() + .set("log", format!("item-{}", c.step)) + .at_step(c.step), + )) + }) + .set_entry("append") + .set_finish("append") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let exec = graph.run_with_thread("update-thread", base).await.unwrap(); + assert_eq!(exec.state.get("log"), Some(&json!(["item-1"]))); + + // A manual write layers on top through the same channel write path. + let config = graph + .update_state( + "update-thread", + ChannelUpdate::new().set("log", "manual"), + None, + ) + .await + .unwrap(); + + let checkpoint = cp + .get(&config.thread_id, config.checkpoint_id.as_deref()) + .await + .unwrap() + .expect("checkpoint exists"); + assert_eq!( + checkpoint.state.get("log"), + Some(&json!(["item-1", "manual"])) + ); + // The manual write's own delta is recorded, honoring the same + // per-checkpoint delta-tracking contract a normal boundary uses. + assert_eq!( + checkpoint.channel_deltas.get("log"), + Some(&vec![json!("manual")]) + ); + + // The delta history across the whole lineage includes both the normal + // boundary's write and the manual one, in order. + let history_config = CheckpointConfig::latest("update-thread"); + let history = cp.delta_history(&history_config, "log").await.unwrap(); + assert_eq!(history, vec![json!("item-1"), json!("manual")]); +} + +#[tokio::test] +async fn update_state_overwrite_resets_baseline_for_subsequent_appends() { + use crate::checkpoint::{Checkpointer, InMemoryCheckpointer}; + use std::sync::Arc; + + let base = ChannelState::new().with_channel("log", Topic); + let cp: Arc> = Arc::new(InMemoryCheckpointer::new()); + let graph = GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .add_node("noop", |_s: ChannelState, _c: NodeContext| async move { + Ok(NodeResult::Update(ChannelUpdate::new().set("log", "a"))) + }) + .set_entry("noop") + .set_finish("noop") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + graph.run_with_thread("reset-thread", base).await.unwrap(); + graph + .update_state( + "reset-thread", + ChannelUpdate::new().overwrite("log", json!(["reset"])), + None, + ) + .await + .unwrap(); + let config = graph + .update_state("reset-thread", ChannelUpdate::new().set("log", "b"), None) + .await + .unwrap(); + + let checkpoint = cp + .get(&config.thread_id, config.checkpoint_id.as_deref()) + .await + .unwrap() + .expect("checkpoint exists"); + assert_eq!(checkpoint.state.get("log"), Some(&json!(["reset", "b"]))); +} + +#[tokio::test] +async fn state_history_reconstruction_equals_live_state_at_every_step() { + use crate::checkpoint::{Checkpointer, InMemoryCheckpointer}; + use std::sync::Arc; + + let base = ChannelState::new().with_channel("log", Topic); + let cp: Arc> = Arc::new(InMemoryCheckpointer::new()); + let graph = GraphBuilder::::new() + .set_reducer(ChannelState::new()) + .add_node("append", |_s: ChannelState, c: NodeContext| async move { + Ok(NodeResult::Update( + ChannelUpdate::new() + .set("log", format!("item-{}", c.step)) + .at_step(c.step), + )) + }) + .add_edge("append", "append") + .set_entry("append") + .set_finish("append") + .with_recursion_limit(6) + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let _ = graph.run_with_thread("history-thread", base).await; + + let history = cp.state_history("history-thread", &[], None).await.unwrap(); + for tuple in &history { + let step = tuple.checkpoint.to_metadata().step; + let expected: Vec = (1..=step).map(|n| json!(format!("item-{n}"))).collect(); + assert_eq!( + tuple.checkpoint.state.get("log"), + Some(&Value::Array(expected)), + "checkpoint state at step {step} did not match the expected live sequence" + ); + } + assert!(!history.is_empty()); +} diff --git a/crates/tinyagents-graph/src/channel/types.rs b/crates/tinyagents-graph/src/channel/types.rs index d8d47a9d..e35605cc 100644 --- a/crates/tinyagents-graph/src/channel/types.rs +++ b/crates/tinyagents-graph/src/channel/types.rs @@ -22,6 +22,7 @@ use std::collections::{BTreeMap, HashMap}; +use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::Result; @@ -55,6 +56,18 @@ pub trait Channel: Send + Sync { /// O(existing) allocation per write. fn merge(&self, current: Option, incoming: Value) -> Result; + /// The channel's construction config, serialized so [`ChannelSet`] can + /// round-trip `{ kind, config, value }` through a durable checkpointer + /// without knowing the concrete channel type. Paired with + /// [`channel_from_config`] on decode. Channels with no configuration + /// (the default) serialize `Value::Null`; [`Barrier`]/[`NamedBarrier`] + /// carry their `expected` set, and [`BinaryAggregate`] carries the + /// registered reducer name (see [`BinaryAggregate::named`] and + /// [`crate::channel::ReducerRegistry`]). + fn config(&self) -> Value { + Value::Null + } + /// Whether more than one concurrent branch may write this channel within a /// single superstep. Aggregates (append/fold/accumulate/barrier) return /// `true`; overwrite-style channels return `false` and trigger @@ -163,6 +176,14 @@ pub struct NamedBarrier { #[derive(Clone)] pub struct BinaryAggregate { pub(crate) fold: std::sync::Arc Result + Send + Sync>, + /// The reducer's registered name (see [`BinaryAggregate::named`]), when + /// it was constructed from the [`crate::channel::ReducerRegistry`]. + /// `None` for a channel built from a bare closure via + /// [`BinaryAggregate::new`]/[`BinaryAggregate::from_reducer`] — such a + /// channel merges correctly at runtime but cannot round-trip through a + /// durable checkpointer (its [`Channel::config`] carries no reducer name + /// to decode from). + pub(crate) reducer_name: Option, } impl std::fmt::Debug for BinaryAggregate { @@ -183,6 +204,10 @@ impl std::fmt::Debug for BinaryAggregate { pub struct ChannelSet { pub(crate) channels: HashMap>, pub(crate) values: HashMap, + /// Delta-tracked channels registered via [`ChannelSet::with_delta`]: + /// name -> how often (in writes to that channel) a full-value snapshot + /// marker is additionally recorded alongside the per-write delta. + pub(crate) delta_channels: HashMap, } impl Clone for ChannelSet { @@ -194,6 +219,7 @@ impl Clone for ChannelSet { .map(|(k, v)| (k.clone(), v.clone_box())) .collect(), values: self.values.clone(), + delta_channels: self.delta_channels.clone(), } } } @@ -212,15 +238,48 @@ impl std::fmt::Debug for ChannelSet { } } -/// A batch of `(channel_name, value)` writes returned by a node. +/// One write within a [`ChannelUpdate`]: an ordinary reducer-merged write, or +/// an [`Overwrite`](ChannelWrite::Overwrite) that bypasses the channel's +/// merge rule entirely and replaces the value outright. /// -/// Build one with [`ChannelUpdate::new`] and chain [`ChannelUpdate::set`]. Tag -/// it with [`ChannelUpdate::at_step`] (passing `ctx.step`) to opt into -/// same-step concurrent-write conflict detection and ephemeral clearing — see -/// the module docs and [`ChannelState`]. +/// `Overwrite` is what lets an append-style ([`Topic`], a delta-tracked +/// channel) reset its baseline: the replaced value becomes what subsequent +/// merges build on, and — for a channel registered with +/// [`ChannelSet::with_delta`] — it also rebases that channel's accumulated +/// delta history (see [`ChannelState::step_deltas`]). +#[derive(Clone, Debug)] +pub enum ChannelWrite { + /// Folds `Value` into the channel's current value via its merge rule. + Merge(Value), + /// Replaces the channel's current value with `Value`, bypassing the + /// merge rule. + Overwrite(Value), +} + +impl ChannelWrite { + /// The raw value carried by either variant. + pub fn value(&self) -> &Value { + match self { + ChannelWrite::Merge(v) | ChannelWrite::Overwrite(v) => v, + } + } + + /// Whether this write is an [`Overwrite`](ChannelWrite::Overwrite). + pub fn is_overwrite(&self) -> bool { + matches!(self, ChannelWrite::Overwrite(_)) + } +} + +/// A batch of `(channel_name, write)` writes returned by a node. +/// +/// Build one with [`ChannelUpdate::new`] and chain [`ChannelUpdate::set`] (an +/// ordinary merged write) or [`ChannelUpdate::overwrite`] (bypasses the merge +/// rule). Tag it with [`ChannelUpdate::at_step`] (passing `ctx.step`) to opt +/// into same-step concurrent-write conflict detection and ephemeral clearing +/// — see the module docs and [`ChannelState`]. #[derive(Clone, Debug, Default)] pub struct ChannelUpdate { - pub(crate) writes: Vec<(String, Value)>, + pub(crate) writes: Vec<(String, ChannelWrite)>, pub(crate) step: Option, } @@ -236,13 +295,31 @@ pub struct ChannelUpdate { /// The reducer's `&self` receiver is unused — the merge rules travel inside the /// running state's [`ChannelSet`] — so any `ChannelState` value (for example /// [`ChannelState::default`]) can be passed to `set_reducer`. -#[derive(Clone, Debug, Default)] +#[derive(Clone, Debug, Default, Serialize, Deserialize)] pub struct ChannelState { pub(crate) set: ChannelSet, /// The step number of the writes currently accumulated in `step_writes`; /// `0` before the first stamped update is seen. + #[serde(default)] pub(crate) current_step: usize, /// Per-channel write counts within `current_step`, used to detect /// concurrent writes to non-aggregate channels. + #[serde(default)] pub(crate) step_writes: HashMap, + /// Cumulative per-channel version counter, bumped once per distinct + /// channel name touched by each folded [`ChannelUpdate`] (I5/R3: see + /// `docs/modules/graph/state-channels.md`'s "channel versions" section). + /// Persisted on [`crate::checkpoint::Checkpoint::channel_versions`] at + /// every boundary so a resumed run's node-visible versions stay + /// continuous across a restart. + #[serde(default)] + pub(crate) channel_versions: BTreeMap, + /// This step's accumulated raw write values for every channel + /// registered via [`ChannelSet::with_delta`], reset whenever the + /// stamped step advances (mirrors `step_writes`). Read by the + /// checkpoint-construction call sites (`compiled::boundary`, + /// `compiled::state_api`) into + /// [`crate::checkpoint::Checkpoint::channel_deltas`]. + #[serde(default)] + pub(crate) step_deltas: BTreeMap>, } diff --git a/crates/tinyagents-graph/src/checkpoint/file.rs b/crates/tinyagents-graph/src/checkpoint/file.rs index c56a77ea..9c052f56 100644 --- a/crates/tinyagents-graph/src/checkpoint/file.rs +++ b/crates/tinyagents-graph/src/checkpoint/file.rs @@ -22,19 +22,99 @@ use async_trait::async_trait; use serde::Serialize; use serde::de::DeserializeOwned; -/// Minimal projection used to read a checkpoint's id without deserializing its -/// `State` payload, so `get` can pick the target line and decode only that one. +/// Minimal projection used to read a checkpoint's addressing/lineage/metadata +/// fields without deserializing its `State` payload. +/// +/// Every field here is state-independent, so this decodes successfully for +/// *any* `Checkpoint` line regardless of what `State` is. It backs +/// three paths that never need the full state: `get`/`get_scoped` picking +/// their target line, and `list` projecting [`CheckpointMetadata`] for every +/// line in a thread. #[derive(serde::Deserialize)] -struct CheckpointIdHeader { +struct CheckpointHeader { + #[serde(default = "checkpoint_header_version_v1")] + version: u32, checkpoint_id: String, + #[serde(default)] + run_id: Option, + #[serde(default)] + parent_checkpoint_id: Option, + #[serde(default)] + namespace: Vec, + /// v2 pending-task set; empty on a v1 record (see `next_nodes`/ + /// `pending_activations` below). + #[serde(default)] + tasks: Vec, + /// v1: node-id-only projection of pending work. + #[serde(default)] + next_nodes: Vec, + /// v1: richer pending-activation superset of `next_nodes`. + #[serde(default)] + pending_activations: Option>, + /// Only the count matters ([`CheckpointMetadata::has_interrupts`]), so + /// each element is decoded as an opaque, ignored JSON value rather than + /// the full `Interrupt` type. + #[serde(default)] + interrupts: Vec, + #[serde(default)] + metadata: serde_json::Value, +} + +fn checkpoint_header_version_v1() -> u32 { + 1 +} + +impl CheckpointHeader { + /// Projects this header onto [`CheckpointMetadata`], mirroring + /// [`Checkpoint::to_metadata`] field-for-field (source/step parsed out of + /// the same free-form `metadata` value, and the pending-task set resolved + /// the same v2-else-v1 way `Checkpoint::effective_tasks` does — a header + /// decode never sees `State`, so it cannot just deserialize the full + /// record and call `to_metadata` on it). `thread_id` is supplied by the + /// caller rather than decoded, since every header on a thread's file + /// carries the same value the caller already knows. + fn into_metadata(self, thread_id: &str) -> CheckpointMetadata { + let source = self + .metadata + .get("source") + .and_then(|v| v.as_str()) + .and_then(CheckpointSource::parse) + .unwrap_or(CheckpointSource::Loop); + let step = self + .metadata + .get("step") + .and_then(|v| v.as_u64()) + .unwrap_or(0) as usize; + let next_nodes = if self.version >= super::CHECKPOINT_FORMAT_VERSION { + self.tasks.into_iter().map(|t| t.node).collect() + } else { + match self.pending_activations { + Some(pending) if !pending.is_empty() => { + pending.into_iter().map(|t| t.node).collect() + } + _ => self.next_nodes, + } + }; + CheckpointMetadata { + thread_id: thread_id.to_string(), + checkpoint_id: self.checkpoint_id, + run_id: self.run_id, + parent_checkpoint_id: self.parent_checkpoint_id, + namespace: self.namespace, + next_nodes, + has_interrupts: !self.interrupts.is_empty(), + source, + step, + } + } } use super::{ - Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointTuple, Checkpointer, PendingWrite, - decode_json_err, merge_writes, + Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, CheckpointTuple, + Checkpointer, PendingActivation, PendingWrite, decode_json_err, merge_writes, }; use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::CheckpointId; +use tinyagents_harness::ids::{CheckpointId, NodeId}; /// File extension for per-thread checkpoint logs. const THREAD_EXT: &str = "jsonl"; @@ -47,6 +127,16 @@ const THREAD_EXT: &str = "jsonl"; /// file keeps the checkpoint log exactly as it was. const WRITES_SUFFIX: &str = ".writes.jsonl"; +/// Filename suffix for a thread's execution-lease sidecar (C3/R4). +const LEASE_SUFFIX: &str = ".lease"; + +/// One thread's execution lease, as persisted in its `.lease` sidecar. +#[derive(Clone, serde::Serialize, serde::Deserialize)] +struct LeaseRecord { + owner: String, + expires_at_ms: u64, +} + /// Process-wide counter making temp-file names unique so concurrent atomic /// rewrites of the same thread never collide on their scratch file. static TMP_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); @@ -122,6 +212,12 @@ impl FileCheckpointer { .join(format!("{}{WRITES_SUFFIX}", escape_thread_id(thread_id))) } + /// Resolves the execution-lease sidecar path for `thread_id` (C3/R4). + fn lease_path(&self, thread_id: &str) -> PathBuf { + self.base_dir + .join(format!("{}{LEASE_SUFFIX}", escape_thread_id(thread_id))) + } + fn legacy_thread_path(&self, thread_id: &str) -> PathBuf { self.base_dir.join(format!( "{}.{THREAD_EXT}", @@ -148,6 +244,35 @@ impl FileCheckpointer { serde_json::from_str::(line) }) } + + /// Reads every record's [`CheckpointHeader`] in `thread_id`'s file and + /// projects each onto [`CheckpointMetadata`], without ever decoding a + /// line's `State` payload. + /// + /// This is what makes [`Checkpointer::list`] cheap on a large thread: the + /// old implementation went through [`FileCheckpointer::read_records`], + /// which fully deserializes `Checkpoint` — including `state` — + /// for every line just to summarize it. `State` is not `DeserializeOwned` + /// bounded here (unlike `read_records`), since a header decode never + /// touches it. + /// + /// Returns an empty vec when the thread file does not exist. + fn read_headers(&self, thread_id: &str) -> Result> { + let path = self.thread_path(thread_id); + let text = match fs::read_to_string(&path) { + Ok(t) => t, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(e) => return Err(io_err("open thread file", e)), + }; + let headers: Vec = + decode_lines(&text, &format!("thread `{thread_id}`"), |line| { + serde_json::from_str::(line) + })?; + Ok(headers + .into_iter() + .map(|h| h.into_metadata(thread_id)) + .collect()) + } } impl Clone for FileCheckpointer { @@ -251,9 +376,72 @@ where Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), Err(e) => return Err(io_err("open thread file", e)), }; - decode_lines(&text, &format!("thread `{thread_id}`"), |line| { + let mut records = decode_lines(&text, &format!("thread `{thread_id}`"), |line| { serde_json::from_str::>(line) - }) + })?; + for record in &mut records { + record.normalize(); + } + Ok(records) + } + + /// Loads a checkpoint for `thread_id`, optionally scoped to `namespace`. + /// + /// Streams lines and fully decodes only the single target line, instead of + /// deserializing every record's `State` just to pick one — the same + /// header-then-full-decode shape [`FileCheckpointer::read_headers`] uses + /// for `list`. Selection matches the historical `rev().find` / + /// `next_back` semantics: the last matching line (or the last line + /// overall, for `checkpoint_id == None`) wins. `namespace` is `None` for + /// [`Checkpointer::get`] (no scoping) and `Some` for + /// [`Checkpointer::get_scoped`]. + fn get_sync( + &self, + thread_id: &str, + checkpoint_id: Option<&str>, + namespace: Option<&[String]>, + ) -> Result>> { + let path = self.thread_path(thread_id); + let file = match File::open(&path) { + Ok(f) => f, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(e) => return Err(io_err("open thread file", e)), + }; + let reader = BufReader::new(file); + let mut target: Option = None; + for line in reader.lines() { + let line = line.map_err(|e| io_err("read line", e))?; + if line.trim().is_empty() { + continue; + } + // Decode only the header to test the match, not `State` — unless + // there is nothing to filter on (no id, no namespace), in which + // case every line matches and decoding one would be wasted work. + if checkpoint_id.is_some() || namespace.is_some() { + let header: CheckpointHeader = serde_json::from_str(&line) + .map_err(|e| decode_json_err("file checkpointer", "header", e))?; + if let Some(namespace) = namespace + && header.namespace.as_slice() != namespace + { + continue; + } + if let Some(id) = checkpoint_id + && header.checkpoint_id != id + { + continue; + } + } + target = Some(line); + } + match target { + Some(line) => { + let mut checkpoint: Checkpoint = serde_json::from_str(&line) + .map_err(|e| decode_json_err("file checkpointer", "record", e))?; + checkpoint.normalize(); + Ok(Some(checkpoint)) + } + None => Ok(None), + } } } @@ -283,7 +471,7 @@ where match decode(line) { Ok(record) => out.push(record), Err(e) if !complete && i == last_index => { - tinyagents_tracing::warn!( + tracing::warn!( "[checkpoint:file] {what}: discarding torn trailing line \ ({} bytes, no terminating newline): {e}", line.len() @@ -366,6 +554,72 @@ fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> { Ok(()) } +/// Appends `bytes` to `path` (creating it if necessary) and fsyncs, without +/// the temp-file-plus-rename dance [`write_atomic`] pays for a full rewrite. +/// +/// `put_writes` used to be read-modify-**rewrite**: every superstep re-read +/// the whole sidecar, merged in the new writes, and rewrote the entire file +/// through `write_atomic` — one fsync'd temp file and rename per superstep, +/// no matter how small the delta. The sidecar is append-only content by +/// construction (each line is independently addressed by the +/// `(namespace, checkpoint_id, task_id, idx)` it carries), so a superstep +/// only ever needs to add lines, never touch existing ones — appending is +/// the same `OpenOptions::append(true)` + single `write_all` + `sync_all` +/// shape [`Checkpointer::put`] already uses for the (also append-only) +/// checkpoint log itself, and carries the same durability guarantee: the +/// fsync means a "persisted" write is durable on stable storage before this +/// returns, not just sitting in the page cache. +/// +/// Safe to call repeatedly within this process: POSIX/Windows both make a +/// single `write_all` under `O_APPEND`/`FILE_APPEND_DATA` atomic with respect +/// to other appenders (no line here ever exceeds a few hundred bytes, well +/// under any platform's atomic-write threshold), so concurrent in-process +/// callers interleave whole lines, never partial ones — the same assumption +/// `read_lines`'s [`decode_lines`] torn-trailing-line tolerance already +/// covers for a genuine crash mid-write. +fn append_atomic(path: &Path, bytes: &[u8]) -> Result<()> { + if let Some(dir) = path.parent() { + fs::create_dir_all(dir).map_err(|e| io_err("create base dir", e))?; + } + let mut file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .map_err(|e| io_err("open file for append", e))?; + file.write_all(bytes) + .map_err(|e| io_err("append record", e))?; + file.sync_all().map_err(|e| io_err("fsync record", e)) +} + +/// Reconstructs the pending-writes ledger for one `(checkpoint_id, namespace)` +/// from a thread's append-only sidecar records, applying +/// [`merge_writes`]'s replace-vs-ignore identity rule to the matching entries +/// in file order. +/// +/// This is the read-side counterpart of the append-only format: `put_writes` +/// appends a line only when a write is new or (for a control-plane upsert) +/// changes the stored value, so the same `(task_id, idx)` identity can appear +/// on more than one line over a checkpoint's lifetime — the ledger for that +/// checkpoint is not "every matching line" but "every matching line, folded +/// through the same identity rule that decided whether to append it". Folding +/// here rather than trusting `records` to already be deduplicated is what +/// keeps `get_writes` correct regardless of how many times a control-plane +/// value was overwritten. +fn fold_write_records( + records: impl IntoIterator, + checkpoint_id: &str, + namespace: &[String], +) -> Vec { + let mut acc = Vec::new(); + for record in records { + if record.checkpoint_id != checkpoint_id || record.namespace != namespace { + continue; + } + merge_writes(&mut acc, std::slice::from_ref(&record.write)); + } + acc +} + #[async_trait] impl Checkpointer for FileCheckpointer where @@ -407,57 +661,54 @@ where thread_id: &str, checkpoint_id: Option<&str>, ) -> Result>> { - // Stream lines and fully decode only the single target line, instead of - // deserializing every record's `State` just to pick one. Selection - // matches the previous `rev().find` / `next_back` semantics: the last - // matching line (or the last line, for `None`) wins. - let path = self.thread_path(thread_id); - let file = match File::open(&path) { - Ok(f) => f, - Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None), - Err(e) => return Err(io_err("open thread file", e)), - }; - let reader = BufReader::new(file); - let mut target: Option = None; - for line in reader.lines() { - let line = line.map_err(|e| io_err("read line", e))?; - if line.trim().is_empty() { - continue; - } - match checkpoint_id { - Some(id) => { - // Decode only the id header to test the match, not `State`. - let header: CheckpointIdHeader = serde_json::from_str(&line) - .map_err(|e| decode_json_err("file checkpointer", "header", e))?; - if header.checkpoint_id == id { - target = Some(line); - } - } - None => target = Some(line), - } - } - match target { - Some(line) => { - Ok(Some(serde_json::from_str(&line).map_err(|e| { - decode_json_err("file checkpointer", "record", e) - })?)) - } - None => Ok(None), - } + let this = self.clone(); + let thread_id = thread_id.to_string(); + let checkpoint_id = checkpoint_id.map(str::to_string); + tokio::task::spawn_blocking(move || { + this.get_sync(&thread_id, checkpoint_id.as_deref(), None) + }) + .await + .map_err(|e| io_err("join blocking get task", e))? + } + + async fn get_scoped( + &self, + thread_id: &str, + checkpoint_id: Option<&str>, + namespace: &[String], + ) -> Result>> { + // A direct scan, like `get`, instead of the trait default's + // `list` (a full metadata projection of the whole thread) followed by + // a second full pass through `get` — one file read instead of two, + // and the namespace filter is applied on the header alongside the id + // filter rather than as a separate `list` step. + let this = self.clone(); + let thread_id = thread_id.to_string(); + let checkpoint_id = checkpoint_id.map(str::to_string); + let namespace = namespace.to_vec(); + tokio::task::spawn_blocking(move || { + this.get_sync(&thread_id, checkpoint_id.as_deref(), Some(&namespace)) + }) + .await + .map_err(|e| io_err("join blocking get_scoped task", e))? } async fn list(&self, thread_id: &str) -> Result> { - Ok(self - .read_records(thread_id)? - .iter() - .map(Checkpoint::to_metadata) - .collect()) + let this = self.clone(); + let thread_id = thread_id.to_string(); + tokio::task::spawn_blocking(move || this.read_headers(&thread_id)) + .await + .map_err(|e| io_err("join blocking list task", e))? } async fn get_thread(&self, thread_id: &str) -> Result>> { // Single-pass bulk read: parse the thread file once, instead of the // default's one whole-file `get` scan per listed id (O(H²)). - self.read_records(thread_id) + let this = self.clone(); + let thread_id = thread_id.to_string(); + tokio::task::spawn_blocking(move || this.read_records(&thread_id)) + .await + .map_err(|e| io_err("join blocking get_thread task", e))? } async fn state_history( @@ -466,156 +717,185 @@ where namespace: &[String], limit: Option, ) -> Result>> { - // Read the whole thread once, then walk the parent lineage in memory - // (O(H)), instead of re-reading and re-parsing the file per hop (O(H²)). - let records = self.read_records(thread_id)?; - if records.is_empty() { - return Ok(Vec::new()); - } - - // id -> checkpoint, last write wins for duplicate ids (matching `get`, - // which takes the last matching record). Track the latest checkpoint in - // the target namespace as the walk's starting point. - let mut by_id: std::collections::HashMap> = - std::collections::HashMap::with_capacity(records.len()); - let mut cursor: Option = None; - for record in records { - if record.namespace.as_slice() == namespace { - cursor = Some(record.checkpoint_id.clone()); + let this = self.clone(); + let thread_id = thread_id.to_string(); + let namespace = namespace.to_vec(); + tokio::task::spawn_blocking(move || -> Result>> { + // Read the whole thread once, then walk the parent lineage in memory + // (O(H)), instead of re-reading and re-parsing the file per hop (O(H²)). + let records = this.read_records(&thread_id)?; + if records.is_empty() { + return Ok(Vec::new()); } - by_id.insert(record.checkpoint_id.clone(), record); - } - let mut out = Vec::new(); - while let Some(id) = cursor { - if let Some(limit) = limit - && out.len() >= limit - { - break; + // id -> checkpoint, last write wins for duplicate ids (matching `get`, + // which takes the last matching record). Track the latest checkpoint in + // the target namespace as the walk's starting point. + let mut by_id: std::collections::HashMap> = + std::collections::HashMap::with_capacity(records.len()); + let mut cursor: Option = None; + for record in records { + if record.namespace == namespace { + cursor = Some(record.checkpoint_id.clone()); + } + by_id.insert(record.checkpoint_id.clone(), record); } - // `remove` doubles as a cycle guard: each id is visited at most once. - let Some(checkpoint) = by_id.remove(&id) else { - break; - }; - // A checkpoint outside the target namespace is not visible under - // namespace-scoped lookup, so the lineage walk stops (matching the - // `get_scoped`-based default). - if checkpoint.namespace.as_slice() != namespace { - break; + + let mut out = Vec::new(); + while let Some(id) = cursor { + if let Some(limit) = limit + && out.len() >= limit + { + break; + } + // `remove` doubles as a cycle guard: each id is visited at most once. + let Some(checkpoint) = by_id.remove(&id) else { + break; + }; + // A checkpoint outside the target namespace is not visible under + // namespace-scoped lookup, so the lineage walk stops (matching the + // `get_scoped`-based default). + if checkpoint.namespace != namespace { + break; + } + cursor = checkpoint.parent_checkpoint_id.clone(); + out.push(tuple_from_checkpoint(checkpoint)); } - cursor = checkpoint.parent_checkpoint_id.clone(); - out.push(tuple_from_checkpoint(checkpoint)); - } - Ok(out) + Ok(out) + }) + .await + .map_err(|e| io_err("join blocking state_history task", e))? } async fn list_threads(&self) -> Result> { - let entries = match fs::read_dir(&self.base_dir) { - Ok(e) => e, - Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), - Err(e) => return Err(io_err("read base dir", e)), - }; - let mut threads = Vec::new(); - for entry in entries { - let entry = entry.map_err(|e| io_err("read dir entry", e))?; - let path = entry.path(); - // Match on the filename suffix rather than `Path::extension()`. - // The empty thread id escapes to the empty string, so its file is - // literally `.jsonl` — a dotfile whose `extension()` is `None`, - // which made that thread invisible to listing (and to everything - // built on listing) while `get`/`put` addressed it perfectly well. - let name = entry.file_name(); - let Some(name) = name.to_str() else { - continue; + let base_dir = self.base_dir.clone(); + tokio::task::spawn_blocking(move || -> Result> { + let entries = match fs::read_dir(&base_dir) { + Ok(e) => e, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(e) => return Err(io_err("read base dir", e)), }; - if !name.ends_with(&format!(".{THREAD_EXT}")) || name.ends_with(WRITES_SUFFIX) { - continue; - } - // Recover the canonical thread id from the first record rather than - // un-escaping the filename, so the value always matches what was - // persisted. - let file = File::open(&path).map_err(|e| io_err("open thread file", e))?; - let mut reader = BufReader::new(file); - let mut first = String::new(); - loop { - first.clear(); - let read = reader - .read_line(&mut first) - .map_err(|e| io_err("read line", e))?; - if read == 0 { - break; // empty file — skip - } - if first.trim().is_empty() { + let mut threads = Vec::new(); + for entry in entries { + let entry = entry.map_err(|e| io_err("read dir entry", e))?; + let path = entry.path(); + // Match on the filename suffix rather than `Path::extension()`. + // The empty thread id escapes to the empty string, so its file is + // literally `.jsonl` — a dotfile whose `extension()` is `None`, + // which made that thread invisible to listing (and to everything + // built on listing) while `get`/`put` addressed it perfectly well. + let name = entry.file_name(); + let Some(name) = name.to_str() else { + continue; + }; + if !name.ends_with(&format!(".{THREAD_EXT}")) || name.ends_with(WRITES_SUFFIX) { continue; } - // One unreadable file must not take down the whole listing. - // `list_threads` decodes the first line of *every* file, so an - // error here made a single poisoned thread break listing — - // and therefore every operation built on it — globally. - match serde_json::from_str::>(&first) { - Ok(record) => threads.push(record.thread_id), - Err(e) => tinyagents_tracing::warn!( - "[checkpoint:file] list_threads: skipping unreadable thread file {}: {e}", - path.display() - ), + // Recover the canonical thread id from the first record rather than + // un-escaping the filename, so the value always matches what was + // persisted. + let file = File::open(&path).map_err(|e| io_err("open thread file", e))?; + let mut reader = BufReader::new(file); + let mut first = String::new(); + loop { + first.clear(); + let read = reader + .read_line(&mut first) + .map_err(|e| io_err("read line", e))?; + if read == 0 { + break; // empty file — skip + } + if first.trim().is_empty() { + continue; + } + // One unreadable file must not take down the whole listing. + // `list_threads` decodes the first line of *every* file, so an + // error here made a single poisoned thread break listing — + // and therefore every operation built on it — globally. + match serde_json::from_str::>(&first) { + Ok(record) => threads.push(record.thread_id), + Err(e) => tracing::warn!( + "[checkpoint:file] list_threads: skipping unreadable thread file {}: {e}", + path.display() + ), + } + break; } - break; } - } - Ok(threads) + Ok(threads) + }) + .await + .map_err(|e| io_err("join blocking list_threads task", e))? } async fn delete_thread(&self, thread_id: &str) -> Result<()> { // The write sidecar goes with the thread: leaving it behind would let a // later thread of the same id inherit a dead ledger. - for path in [self.thread_path(thread_id), self.writes_path(thread_id)] { - match fs::remove_file(&path) { - Ok(()) => {} - Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} - Err(e) => return Err(io_err("delete thread file", e)), + let this = self.clone(); + let thread_id = thread_id.to_string(); + tokio::task::spawn_blocking(move || -> Result<()> { + for path in [this.thread_path(&thread_id), this.writes_path(&thread_id)] { + match fs::remove_file(&path) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => return Err(io_err("delete thread file", e)), + } } - } - Ok(()) + Ok(()) + }) + .await + .map_err(|e| io_err("join blocking delete_thread task", e))? } async fn delete_checkpoints(&self, thread_id: &str, ids: &[String]) -> Result { if ids.is_empty() { return Ok(0); } - let drop: std::collections::HashSet<&str> = ids.iter().map(String::as_str).collect(); - let mut records = self.read_records(thread_id)?; - let before = records.len(); - records.retain(|c| !drop.contains(c.checkpoint_id.as_str())); - let removed = before - records.len(); - if removed > 0 { - self.write_records(thread_id, &records)?; - // Drop the deleted checkpoints' write ledgers with them. - let writes_path = self.writes_path(thread_id); - let write_records = Self::read_write_records(&writes_path, thread_id)?; - let kept: Vec<&WriteRecord> = write_records - .iter() - .filter(|r| !drop.contains(r.checkpoint_id.as_str())) - .collect(); - if kept.len() != write_records.len() { - let mut buf = String::new(); - for record in kept { - let line = - serde_json::to_string(record).map_err(|e| io_err("encode write", e))?; - buf.push_str(&line); - buf.push('\n'); - } - if buf.is_empty() { - match fs::remove_file(&writes_path) { - Ok(()) => {} - Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} - Err(e) => return Err(io_err("remove empty writes file", e)), + let this = self.clone(); + let thread_id = thread_id.to_string(); + let ids = ids.to_vec(); + tokio::task::spawn_blocking(move || -> Result { + let drop: std::collections::HashSet<&str> = ids.iter().map(String::as_str).collect(); + let mut records = this.read_records(&thread_id)?; + let before = records.len(); + records.retain(|c| !drop.contains(c.checkpoint_id.as_str())); + let removed = before - records.len(); + if removed > 0 { + this.write_records(&thread_id, &records)?; + // Drop the deleted checkpoints' write ledgers with them. This + // compaction path still fully rewrites the sidecar (unlike + // `put_writes`'s append-only steady state): it runs only on an + // explicit prune/delete, not once per superstep, so the + // rewrite cost is paid where it is actually incurred. + let writes_path = this.writes_path(&thread_id); + let write_records = Self::read_write_records(&writes_path, &thread_id)?; + let kept: Vec<&WriteRecord> = write_records + .iter() + .filter(|r| !drop.contains(r.checkpoint_id.as_str())) + .collect(); + if kept.len() != write_records.len() { + let mut buf = String::new(); + for record in kept { + let line = + serde_json::to_string(record).map_err(|e| io_err("encode write", e))?; + buf.push_str(&line); + buf.push('\n'); + } + if buf.is_empty() { + match fs::remove_file(&writes_path) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => return Err(io_err("remove empty writes file", e)), + } + } else { + write_atomic(&writes_path, buf.as_bytes())?; } - } else { - write_atomic(&writes_path, buf.as_bytes())?; } } - } - Ok(removed) + Ok(removed) + }) + .await + .map_err(|e| io_err("join blocking delete_checkpoints task", e))? } async fn put_writes(&self, config: &CheckpointConfig, writes: &[PendingWrite]) -> Result<()> { @@ -623,51 +903,210 @@ where if writes.is_empty() { return Ok(()); } - let path = self.writes_path(&config.thread_id); - let mut records = Self::read_write_records(&path, &config.thread_id)?; - - // Split out this checkpoint's ledger, merge, then rebuild the file. - let (mut mine, others): (Vec, Vec) = records - .drain(..) - .partition(|r| r.checkpoint_id == checkpoint_id && r.namespace == config.namespace); - let mut existing: Vec = mine.drain(..).map(|r| r.write).collect(); - let changed = merge_writes(&mut existing, writes); - - let mut buf = String::new(); - for record in others.iter() { - let line = serde_json::to_string(record).map_err(|e| io_err("encode write", e))?; - buf.push_str(&line); - buf.push('\n'); - } - for write in existing { - let record = WriteRecord { - namespace: config.namespace.clone(), - checkpoint_id: checkpoint_id.clone(), - write, - }; - let line = serde_json::to_string(&record).map_err(|e| io_err("encode write", e))?; - buf.push_str(&line); - buf.push('\n'); - } - fs::create_dir_all(&self.base_dir).map_err(|e| io_err("create base dir", e))?; - write_atomic(&path, buf.as_bytes())?; - tinyagents_tracing::debug!( - "[checkpoint:file] put_writes thread={} checkpoint={checkpoint_id} offered={} stored={changed}", - config.thread_id, - writes.len() - ); - Ok(()) + let this = self.clone(); + let config = config.clone(); + let writes = writes.to_vec(); + tokio::task::spawn_blocking(move || -> Result<()> { + let path = this.writes_path(&config.thread_id); + let existing_records = Self::read_write_records(&path, &config.thread_id)?; + let mut existing = + fold_write_records(existing_records, &checkpoint_id, &config.namespace); + + // Decide, per incoming write, whether it needs to be appended — + // mirroring `merge_writes`'s replace-vs-ignore rule by hand + // rather than calling it, because this call site (unlike every + // other `merge_writes` caller) also needs to know *which* + // entries changed, so only those get appended instead of + // rewriting the whole ledger. A duplicate data write + // (`idx >= 0`, already-seen `(task_id, idx)`) is a no-op and + // appends nothing; a control-plane write (`idx < 0`) always + // appends its latest value, and a later line for the same + // identity is what `fold_write_records` uses to pick the winner + // on read. + let mut to_append: Vec = Vec::new(); + let mut changed = 0usize; + for write in &writes { + match existing.iter().position(|w| w.identity() == write.identity()) { + Some(idx) => { + if write.is_control_plane() { + existing[idx] = write.clone(); + to_append.push(write.clone()); + changed += 1; + } + // A repeated data write is ignored, not appended. + } + None => { + existing.push(write.clone()); + to_append.push(write.clone()); + changed += 1; + } + } + } + + if to_append.is_empty() { + tracing::debug!( + "[checkpoint:file] put_writes thread={} checkpoint={checkpoint_id} \ + offered={} stored=0 (no new lines appended)", + config.thread_id, + writes.len() + ); + return Ok(()); + } + + let mut buf = String::new(); + for write in to_append { + let record = WriteRecord { + namespace: config.namespace.clone(), + checkpoint_id: checkpoint_id.clone(), + write, + }; + let line = + serde_json::to_string(&record).map_err(|e| io_err("encode write", e))?; + buf.push_str(&line); + buf.push('\n'); + } + append_atomic(&path, buf.as_bytes())?; + tracing::debug!( + "[checkpoint:file] put_writes thread={} checkpoint={checkpoint_id} offered={} stored={changed}", + config.thread_id, + writes.len() + ); + Ok(()) + }) + .await + .map_err(|e| io_err("join blocking put_writes task", e))? } async fn get_writes(&self, config: &CheckpointConfig) -> Result> { let Some(checkpoint_id) = self.resolve_write_target(config).await? else { return Ok(Vec::new()); }; - let path = self.writes_path(&config.thread_id); - Ok(Self::read_write_records(&path, &config.thread_id)? - .into_iter() - .filter(|r| r.checkpoint_id == checkpoint_id && r.namespace == config.namespace) - .map(|r| r.write) - .collect()) + let this = self.clone(); + let config = config.clone(); + tokio::task::spawn_blocking(move || -> Result> { + let path = this.writes_path(&config.thread_id); + let records = Self::read_write_records(&path, &config.thread_id)?; + Ok(fold_write_records( + records, + &checkpoint_id, + &config.namespace, + )) + }) + .await + .map_err(|e| io_err("join blocking get_writes task", e))? + } + + async fn try_claim(&self, thread: &str, owner: &str, ttl: std::time::Duration) -> Result { + let base_dir = self.base_dir.clone(); + let path = self.lease_path(thread); + let owner = owner.to_string(); + let ttl_ms = ttl.as_millis() as u64; + tokio::task::spawn_blocking(move || -> Result { + fs::create_dir_all(&base_dir).map_err(|e| io_err("create base dir", e))?; + // Serialize the read/check/write sequence across processes. The + // checkpoint files already use atomic replacement for durability, + // but replacement alone cannot make this ownership decision + // atomic: two claimants could both observe an absent lease. + let lock_path = path.with_file_name(format!( + "{}.lock", + path.file_name() + .and_then(|name| name.to_str()) + .unwrap_or("lease") + )); + let mut acquired = false; + for _ in 0..1_000 { + match OpenOptions::new() + .write(true) + .create_new(true) + .open(&lock_path) + { + Ok(_) => { + acquired = true; + break; + } + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + std::thread::sleep(std::time::Duration::from_millis(1)); + } + Err(error) => return Err(io_err("create lease lock", error)), + } + } + if !acquired { + return Err(TinyAgentsError::Graph( + "timed out waiting for lease lock".to_string(), + )); + } + let now = tinyagents_harness::ids::now_ms(); + let result = if let Some(existing) = read_lease(&path)? + && existing.owner != owner + && existing.expires_at_ms > now + { + Ok(false) + } else { + let record = LeaseRecord { + owner, + expires_at_ms: now.saturating_add(ttl_ms), + }; + let bytes = serde_json::to_vec(&record).map_err(|e| io_err("encode lease", e))?; + write_atomic(&path, &bytes)?; + Ok(true) + }; + let _ = fs::remove_file(lock_path); + result + }) + .await + .map_err(|e| io_err("join blocking try_claim task", e))? + } + + async fn renew(&self, thread: &str, owner: &str, ttl: std::time::Duration) -> Result { + let path = self.lease_path(thread); + let owner = owner.to_string(); + let ttl_ms = ttl.as_millis() as u64; + tokio::task::spawn_blocking(move || -> Result { + let now = tinyagents_harness::ids::now_ms(); + match read_lease(&path)? { + Some(existing) if existing.owner == owner => { + let record = LeaseRecord { + owner, + expires_at_ms: now.saturating_add(ttl_ms), + }; + let bytes = + serde_json::to_vec(&record).map_err(|e| io_err("encode lease", e))?; + write_atomic(&path, &bytes)?; + Ok(true) + } + _ => Ok(false), + } + }) + .await + .map_err(|e| io_err("join blocking renew task", e))? + } + + async fn release(&self, thread: &str, owner: &str) -> Result<()> { + let path = self.lease_path(thread); + let owner = owner.to_string(); + tokio::task::spawn_blocking(move || -> Result<()> { + if let Some(existing) = read_lease(&path)? + && existing.owner == owner + { + let _ = fs::remove_file(&path); + } + Ok(()) + }) + .await + .map_err(|e| io_err("join blocking release task", e))? + } +} + +/// Reads a thread's execution-lease sidecar, if it exists and decodes. +/// +/// A missing file is `Ok(None)`; a corrupt/malformed file is treated the same +/// way (`Ok(None)`) rather than failing the claim — the lease is best-effort +/// advisory state layered on top of the in-process lock, not the sole source +/// of durability, so a torn write here should not strand a thread. +fn read_lease(path: &Path) -> Result> { + match fs::read(path) { + Ok(bytes) => Ok(serde_json::from_slice(&bytes).ok()), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None), + Err(e) => Err(io_err("read lease", e)), } } diff --git a/crates/tinyagents-graph/src/checkpoint/mod.rs b/crates/tinyagents-graph/src/checkpoint/mod.rs index 5fffde10..fef4ced0 100644 --- a/crates/tinyagents-graph/src/checkpoint/mod.rs +++ b/crates/tinyagents-graph/src/checkpoint/mod.rs @@ -21,10 +21,13 @@ mod types; pub use file::FileCheckpointer; #[cfg(feature = "sqlite")] pub use sqlite::SqliteCheckpointer; +#[cfg(feature = "sqlite")] +pub(crate) use sqlite::prepare_connection; pub use types::{ - BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, - CheckpointTuple, DurabilityMode, PendingActivation, PendingWrite, WRITES_IDX_ERROR, - WRITES_IDX_INTERRUPT, WRITES_IDX_RESUME, merge_writes, + BarrierArrivals, CHECKPOINT_FORMAT_VERSION, Checkpoint, CheckpointConfig, CheckpointMetadata, + CheckpointSource, CheckpointTuple, CompletedTask, DURABLE_TASK_CHANNEL_PREFIX, DurabilityMode, + INTERRUPT_AFTER_CHANNEL, PendingActivation, PendingWrite, WRITES_IDX_ERROR, + WRITES_IDX_INTERRUPT, WRITES_IDX_INTERRUPT_AFTER, WRITES_IDX_RESUME, merge_writes, }; use std::collections::{HashMap, HashSet}; @@ -126,8 +129,9 @@ where /// returned (last-write-wins, consistent with [`Checkpointer::get`]). /// /// Composed from [`Checkpointer::list`] + [`Checkpointer::get`] so every - /// backend inherits it; override for a cheaper scoped query — both durable - /// backends do, because the default costs a full thread scan per call and + /// backend inherits it; override for a cheaper scoped query — both + /// [`FileCheckpointer`] and [`SqliteCheckpointer`](crate::SqliteCheckpointer) + /// do, because the default costs a full thread scan per call and /// [`Checkpointer::state_history`] issues one per lineage hop. async fn get_scoped( &self, @@ -186,6 +190,35 @@ where Ok(()) } + /// Persists `checkpoint` and its `writes` together, at a superstep + /// boundary where both are produced at once. + /// + /// The default body is composed from [`Checkpointer::put`] followed by + /// [`Checkpointer::put_writes`] — two independent calls, so a crash + /// between them can leave the checkpoint durable with its writes lost. + /// That is no worse than calling the two methods separately (which is + /// what every caller did before this method existed), so every backend + /// keeps compiling and behaving exactly as before without overriding it. + /// + /// A backend that can share one transaction across both statements + /// should override this to do so — [`SqliteCheckpointer`] does, so a + /// crash between the two writes is impossible rather than merely + /// unlikely: either both are durable or neither is. + async fn put_with_writes( + &self, + checkpoint: Checkpoint, + writes: &[PendingWrite], + ) -> Result { + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let id = self.put(checkpoint).await?; + self.put_writes(&config, writes).await?; + Ok(id) + } + /// Reads back the writes recorded against the checkpoint addressed by /// `config`, in insertion order. /// @@ -198,6 +231,60 @@ where Ok(Vec::new()) } + // ---- Thread execution lease (C3/R4) ------------------------------------ + // + // The durable half of the per-thread execution lock. The executor + // (`compiled::executor::execute`) already holds an in-process + // `ThreadLockMap` guard for a run's whole lifetime, which is sufficient + // to serialize concurrent calls *within one process*. This lease closes + // the cross-process gap: two different processes (or two restarts of the + // same host) racing `run_with_thread`/`resume` on the same thread id + // have no shared in-process lock to serialize on. A backend that + // implements this lets a dead owner's lease be reclaimed once it expires + // instead of stranding the thread forever, while a live owner's lease + // refuses a competing claim. + // + // Every method carries a default no-op body so an out-of-tree + // `Checkpointer` (and the in-memory backend, which has no cross-process + // audience to protect against) keeps compiling and behaves exactly as it + // did before this lease existed — `try_claim` always succeeds. + + /// Attempts to claim the execution lease for `thread`, naming `owner` + /// (the run id) and expiring after `ttl`. + /// + /// Returns `Ok(true)` when the lease is unclaimed, already expired, or + /// already held by `owner` (idempotent re-claim); `Ok(false)` when a + /// different owner holds a still-live lease. + /// + /// The default body always returns `Ok(true)`. + async fn try_claim( + &self, + _thread: &str, + _owner: &str, + _ttl: std::time::Duration, + ) -> Result { + Ok(true) + } + + /// Extends `owner`'s already-held lease on `thread` by `ttl` from now. + /// + /// Returns `Ok(false)` when `owner` does not currently hold the lease + /// (it expired and was reclaimed, or was never claimed). + /// + /// The default body always returns `Ok(true)`. + async fn renew(&self, _thread: &str, _owner: &str, _ttl: std::time::Duration) -> Result { + Ok(true) + } + + /// Releases `owner`'s lease on `thread`, when it holds one. + /// + /// A no-op (not an error) when `owner` does not hold the lease. + /// + /// The default body is a no-op. + async fn release(&self, _thread: &str, _owner: &str) -> Result<()> { + Ok(()) + } + /// Resolves the checkpoint id a **read** of writes addresses. /// /// Unlike [`Checkpointer::put_writes`] (where an unaddressed id is a caller @@ -346,7 +433,7 @@ where break; }; if !visited.insert(tuple.checkpoint.checkpoint_id.clone()) { - tinyagents_tracing::warn!( + tracing::warn!( "[checkpoint] state_history: lineage cycle at checkpoint `{}` \ (thread `{thread_id}`); truncating the walk", tuple.checkpoint.checkpoint_id @@ -363,6 +450,41 @@ where Ok(out) } + /// Replays a [`crate::channel::ChannelSet::with_delta`]-tracked + /// channel's per-step write history for `config.thread_id`/ + /// `config.namespace` (I5/R3). + /// + /// The default implementation walks [`Checkpointer::state_history`] + /// (newest-first, so it is reversed to oldest-first here) and + /// concatenates each checkpoint's own + /// [`Checkpoint::channel_deltas`] entry for `channel`, in lineage + /// order — every checkpoint carries only *its own step's* writes to a + /// delta-tracked channel (not a cumulative history), which is what + /// keeps a single checkpoint's size bounded regardless of how long the + /// channel's append history grows. A checkpoint with no recorded delta + /// for `channel` (predates delta tracking, or `channel` was not + /// delta-tracked when it was written) contributes nothing. + /// + /// A backend may override this with a cheaper single-pass read; the + /// observable result must remain identical. + async fn delta_history( + &self, + config: &CheckpointConfig, + channel: &str, + ) -> Result> { + let mut tuples = self + .state_history(&config.thread_id, &config.namespace, None) + .await?; + tuples.reverse(); + let mut out = Vec::new(); + for tuple in &tuples { + if let Some(deltas) = tuple.checkpoint.channel_deltas.get(channel) { + out.extend(deltas.iter().cloned()); + } + } + Ok(out) + } + // ---- Thread operations ------------------------------------------------- // // Three storage-specific primitives (`list_threads`, `delete_thread`, @@ -621,7 +743,10 @@ where Some(id) => list.iter().rfind(|c| c.checkpoint_id == id), None => list.last(), }; - Ok(found.cloned()) + Ok(found.cloned().map(|mut c| { + c.normalize(); + c + })) } async fn list(&self, thread_id: &str) -> Result> { @@ -636,7 +761,11 @@ where // Single-pass bulk read: clone the thread's records in insertion // order, instead of the default's one `get` per listed id. let map = self.inner.lock().map_err(|_| lock_err())?; - Ok(map.get(thread_id).cloned().unwrap_or_default()) + let mut records = map.get(thread_id).cloned().unwrap_or_default(); + for record in &mut records { + record.normalize(); + } + Ok(records) } async fn list_threads(&self) -> Result> { @@ -688,7 +817,7 @@ where let mut map = self.writes.lock().map_err(|_| lock_err())?; let slot = map.entry(key).or_default(); let changed = merge_writes(slot, writes); - tinyagents_tracing::debug!( + tracing::debug!( "[checkpoint:memory] put_writes thread={} checkpoint={:?} offered={} stored={}", config.thread_id, config.checkpoint_id, diff --git a/crates/tinyagents-graph/src/checkpoint/sqlite.rs b/crates/tinyagents-graph/src/checkpoint/sqlite.rs index 519768c2..810e6a8d 100644 --- a/crates/tinyagents-graph/src/checkpoint/sqlite.rs +++ b/crates/tinyagents-graph/src/checkpoint/sqlite.rs @@ -37,7 +37,7 @@ use super::{ Checkpointer, PendingWrite, decode_json_err, merge_writes, }; use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::{CheckpointId, NodeId}; +use tinyagents_harness::ids::{CheckpointId, NodeId, TaskId}; /// A [`Checkpointer`] that persists checkpoints in a SQLite database. /// @@ -62,6 +62,32 @@ fn sqlite_err(context: &str, err: impl std::fmt::Display) -> TinyAgentsError { TinyAgentsError::Checkpoint(format!("sqlite checkpointer: {context}: {err}")) } +/// How long a statement waits for a competing writer's lock before giving up +/// with `SQLITE_BUSY`, mirroring `tinyagents-session`'s `store.rs` (see its +/// `BUSY_TIMEOUT` doc comment for why this is set explicitly rather than +/// relied on as an undocumented `rusqlite` default). +const BUSY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5); + +/// Applies the per-connection pragmas every checkpointer handle needs: +/// `journal_mode = WAL` for concurrent-reader-friendly durability, +/// `synchronous = NORMAL` (safe under WAL — only a whole-OS crash can lose a +/// commit, not a process crash) instead of the slower `FULL` default, and an +/// explicit `busy_timeout` so a writer contending with another connection +/// waits rather than failing immediately with `SQLITE_BUSY`. +/// +/// Shared with [`crate::cache::SqliteTaskCache`], which opens its own +/// connection with the same pragmas. +pub(crate) fn prepare_connection(conn: &Connection) -> Result<()> { + conn.busy_timeout(BUSY_TIMEOUT) + .map_err(|e| sqlite_err("set busy_timeout", e))?; + conn.execute_batch( + "PRAGMA journal_mode = WAL; + PRAGMA synchronous = NORMAL;", + ) + .map_err(|e| sqlite_err("apply pragmas", e))?; + Ok(()) +} + impl SqliteCheckpointer { /// Opens (creating if needed) a SQLite-backed checkpointer at `path`. /// @@ -94,8 +120,10 @@ impl SqliteCheckpointer { /// across the boundary), apply [`SqliteCheckpointer::schema_sql`] to your own /// connection instead and drive the tables directly. pub fn from_connection(conn: Connection) -> Result { + prepare_connection(&conn)?; conn.execute_batch(SCHEMA) .map_err(|e| sqlite_err("create schema", e))?; + migrate_checkpoint_format_columns(&conn)?; Ok(Self { conn: Arc::new(Mutex::new(conn)), _marker: PhantomData, @@ -113,13 +141,62 @@ impl SqliteCheckpointer { SCHEMA } - fn lock(&self) -> Result> { - self.conn.lock().map_err(|_| { - TinyAgentsError::Checkpoint("sqlite checkpointer: connection lock poisoned".to_string()) - }) + /// Test-only: reads back the live `journal_mode` pragma from this + /// checkpointer's own connection. + /// + /// Exists so the pragma regression test observes exactly what + /// [`prepare_connection`] set on `self`'s handle, rather than a second, + /// freshly opened connection (whose own pragmas default independently — + /// `synchronous` is per-connection, not persisted in the file, though + /// `journal_mode` is). + #[cfg(test)] + pub(crate) fn journal_mode(&self) -> Result { + let conn = lock_conn(&self.conn)?; + conn.query_row("PRAGMA journal_mode", [], |row| row.get(0)) + .map_err(|e| sqlite_err("read journal_mode pragma", e)) + } + + /// Test-only: reads back the live `synchronous` pragma from this + /// checkpointer's own connection. SQLite reports it as an integer + /// (`0` = OFF, `1` = NORMAL, `2` = FULL, `3` = EXTRA). + #[cfg(test)] + pub(crate) fn synchronous(&self) -> Result { + let conn = lock_conn(&self.conn)?; + conn.query_row("PRAGMA synchronous", [], |row| row.get(0)) + .map_err(|e| sqlite_err("read synchronous pragma", e)) + } + + /// Test-only: whether `checkpoints` currently has a column named `column` + /// — used to assert [`migrate_checkpoint_format_columns`] actually ran + /// against a database opened from a pre-v2 schema. + #[cfg(test)] + pub(crate) fn has_checkpoints_column(&self, column: &str) -> Result { + let conn = lock_conn(&self.conn)?; + let mut stmt = conn + .prepare("PRAGMA table_info(checkpoints)") + .map_err(|e| sqlite_err("inspect checkpoints schema", e))?; + let rows = stmt + .query_map([], |row| row.get::<_, String>(1)) + .map_err(|e| sqlite_err("query checkpoints schema", e))?; + for row in rows { + if row.map_err(|e| sqlite_err("read schema column", e))? == column { + return Ok(true); + } + } + Ok(false) } } +/// Locks a checkpointer's shared connection, mapping a poisoned mutex to a +/// [`TinyAgentsError`]. Free function (rather than a method) so it can be +/// called from inside a `spawn_blocking` closure that only holds the cloned +/// `Arc>`, not `&self`. +fn lock_conn(conn: &Arc>) -> Result> { + conn.lock().map_err(|_| { + TinyAgentsError::Checkpoint("sqlite checkpointer: connection lock poisoned".to_string()) + }) +} + /// Table + indexes. `seq` preserves insertion order; the indexes serve thread /// listing, `(thread_id, checkpoint_id)` parent-chain lookups, and — since the /// namespace-scoped overrides landed — `(thread_id, namespace, …)` scoped @@ -153,7 +230,9 @@ CREATE TABLE IF NOT EXISTS checkpoints ( source TEXT NOT NULL, step INTEGER NOT NULL, has_interrupts INTEGER NOT NULL, - record TEXT NOT NULL + record TEXT NOT NULL, + format_version INTEGER NOT NULL DEFAULT 1, + created_at INTEGER NOT NULL DEFAULT 0 ); CREATE INDEX IF NOT EXISTS idx_checkpoints_thread ON checkpoints (thread_id, seq); CREATE INDEX IF NOT EXISTS idx_checkpoints_lookup ON checkpoints (thread_id, checkpoint_id); @@ -174,8 +253,103 @@ CREATE TABLE IF NOT EXISTS checkpoint_writes ( ); CREATE INDEX IF NOT EXISTS idx_checkpoint_writes_thread ON checkpoint_writes (thread_id, checkpoint_id); + +-- C3/R4: the durable half of the per-thread execution lease. The executor +-- holds an in-process lock for the run's lifetime (see +-- `compiled::executor::execute`) AND claims this row, so a lease surviving a +-- crashed owner past its TTL is reclaimable by a different process instead of +-- stranding the thread forever. +CREATE TABLE IF NOT EXISTS thread_leases ( + thread_id TEXT PRIMARY KEY, + owner TEXT NOT NULL, + expires_at INTEGER NOT NULL +); +"; + +/// Recursive parent-chain walk for [`Checkpointer::state_history`], newest +/// first, capped by `?3` rows — a caller-supplied `limit` (already clamped to +/// the namespace's total row count) rather than a truncation applied in Rust +/// after decoding everything. +/// +/// `latest` first dedups: `put` never updates a row in place (see the module +/// doc), so a reused `checkpoint_id` — from `copy_thread`, a fork, or a +/// hand-written record — can have more than one row. Keeping only the +/// highest-`seq` row per id makes `checkpoint_id` unique within `latest`, +/// which is what makes the following recursive join well-defined: a plain +/// `JOIN` on `parent_checkpoint_id = checkpoint_id` over non-unique ids could +/// fan out. +/// +/// `chain` walks from the head (the namespace's own highest-`seq` row) along +/// `parent_checkpoint_id`, capped by `depth < ?3`. Termination is guaranteed +/// even over a corrupted lineage with a parent cycle: `latest` has at most one +/// row per distinct id, so after at most that many hops the depth cap (itself +/// bounded by the namespace's total row count, see the caller) stops the +/// recursion regardless of what the pointers do. +const STATE_HISTORY_CTE: &str = "\ +WITH RECURSIVE latest AS ( + SELECT c1.seq, c1.checkpoint_id, c1.parent_checkpoint_id, c1.record + FROM checkpoints c1 + WHERE c1.thread_id = ?1 AND c1.namespace = ?2 + AND c1.seq = ( + SELECT MAX(c2.seq) FROM checkpoints c2 + WHERE c2.thread_id = c1.thread_id AND c2.namespace = c1.namespace + AND c2.checkpoint_id = c1.checkpoint_id + ) +), +chain(seq, checkpoint_id, parent_checkpoint_id, record, depth) AS ( + SELECT seq, checkpoint_id, parent_checkpoint_id, record, 1 + FROM latest + WHERE seq = (SELECT MAX(seq) FROM latest) + UNION ALL + SELECT l.seq, l.checkpoint_id, l.parent_checkpoint_id, l.record, chain.depth + 1 + FROM latest l + JOIN chain ON l.checkpoint_id = chain.parent_checkpoint_id + WHERE chain.depth < ?3 +) +SELECT record FROM chain ORDER BY depth ASC LIMIT ?3; "; +/// Adds the checkpoint format v2 columns (`format_version`, `created_at`) to +/// an existing `checkpoints` table that predates them, guarded by +/// `PRAGMA table_info` so it is a no-op on a database that already has them +/// (a fresh database gets them for free from [`SCHEMA`] once that DDL is +/// updated to declare them directly — this migration exists for a database +/// opened by an older build, whose `checkpoints` table was created without +/// these columns). +/// +/// `format_version` defaults to `1`: an existing row predates this migration +/// by construction, so it was written by a build that only ever produced +/// checkpoint format v1 records. `created_at` defaults to `0`, the same +/// visibly-unset sentinel [`Checkpoint::created_at`] uses for a v1 record +/// decoded from JSON with no `created_at` field. +fn migrate_checkpoint_format_columns(conn: &Connection) -> Result<()> { + let mut existing: std::collections::HashSet = std::collections::HashSet::new(); + { + let mut stmt = conn + .prepare("PRAGMA table_info(checkpoints)") + .map_err(|e| sqlite_err("inspect checkpoints schema", e))?; + let rows = stmt + .query_map([], |row| row.get::<_, String>(1)) + .map_err(|e| sqlite_err("query checkpoints schema", e))?; + for row in rows { + existing.insert(row.map_err(|e| sqlite_err("read schema column", e))?); + } + } + if !existing.contains("format_version") { + conn.execute_batch( + "ALTER TABLE checkpoints ADD COLUMN format_version INTEGER NOT NULL DEFAULT 1;", + ) + .map_err(|e| sqlite_err("add format_version column", e))?; + } + if !existing.contains("created_at") { + conn.execute_batch( + "ALTER TABLE checkpoints ADD COLUMN created_at INTEGER NOT NULL DEFAULT 0;", + ) + .map_err(|e| sqlite_err("add created_at column", e))?; + } + Ok(()) +} + /// The projected listing columns read from one `checkpoints` row. struct MetaRow { thread_id: String, @@ -209,6 +383,111 @@ fn row_metadata(row: MetaRow) -> Result { }) } +/// Inserts one `checkpoints` row for `checkpoint`. +/// +/// Takes `&Connection` rather than `&SqliteCheckpointer` so it can run either +/// directly against a locked connection ([`Checkpointer::put`]) or against a +/// [`rusqlite::Transaction`] (which derefs to `Connection`) shared with a +/// `put_writes` insert in the same commit +/// ([`SqliteCheckpointer`]'s `put_with_writes` override). +fn insert_checkpoint_row( + conn: &Connection, + checkpoint: &Checkpoint, +) -> Result<()> { + let meta = checkpoint.to_metadata(); + let namespace = serde_json::to_string(&checkpoint.namespace) + .map_err(|e| sqlite_err("encode namespace", e))?; + // Projected from `to_metadata()`'s v2-or-derived-from-v1 resolution + // (`Checkpoint::effective_tasks`), not `checkpoint.next_nodes` directly — + // a v2 checkpoint (every write this crate performs) leaves that legacy + // field empty, so reading it here would silently persist an empty + // `next_nodes` listing column for every checkpoint going forward. + let next_nodes = + serde_json::to_string(&meta.next_nodes).map_err(|e| sqlite_err("encode next_nodes", e))?; + let record = serde_json::to_string(checkpoint).map_err(|e| sqlite_err("encode record", e))?; + conn.execute( + "INSERT INTO checkpoints ( + thread_id, checkpoint_id, parent_checkpoint_id, run_id, + namespace, next_nodes, source, step, has_interrupts, record, + format_version, created_at + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)", + params![ + checkpoint.thread_id, + checkpoint.checkpoint_id, + checkpoint.parent_checkpoint_id, + checkpoint.run_id, + namespace, + next_nodes, + meta.source.as_str(), + meta.step as i64, + i64::from(meta.has_interrupts), + record, + checkpoint.version as i64, + checkpoint.created_at as i64, + ], + ) + .map_err(|e| sqlite_err("insert checkpoint", e))?; + Ok(()) +} + +/// Inserts `writes` into `checkpoint_writes` for the checkpoint addressed by +/// `config`, returning how many rows were actually stored (a control-plane +/// write always stores; a data write with an already-seen `(task_id, idx)` is +/// ignored — see [`Checkpointer::put_writes`]'s doc comment for the rule). +/// +/// Takes `&Connection` for the same reason as [`insert_checkpoint_row`]: it +/// runs standalone under [`Checkpointer::put_writes`] and shares a +/// transaction with [`insert_checkpoint_row`] under `put_with_writes`. +fn insert_checkpoint_writes( + conn: &Connection, + config: &CheckpointConfig, + checkpoint_id: &str, + writes: &[PendingWrite], +) -> Result { + let namespace_json = + serde_json::to_string(&config.namespace).map_err(|e| sqlite_err("encode namespace", e))?; + let mut stored = 0usize; + for write in writes { + // The replace-vs-ignore rule pushed into SQL: a control-plane write + // (`idx < 0`) legitimately changes on a retry and upserts, while a + // data write is append-once so a retried `put_writes` is a no-op. + // Doing it with two conflict clauses rather than a read-then-write + // keeps it correct under concurrent writers. + let sql = if write.is_control_plane() { + "INSERT INTO checkpoint_writes + (thread_id, namespace, checkpoint_id, task_id, idx, node, channel, payload) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + ON CONFLICT(thread_id, namespace, checkpoint_id, task_id, idx) DO UPDATE SET + node = excluded.node, + channel = excluded.channel, + payload = excluded.payload" + } else { + "INSERT INTO checkpoint_writes + (thread_id, namespace, checkpoint_id, task_id, idx, node, channel, payload) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + ON CONFLICT(thread_id, namespace, checkpoint_id, task_id, idx) DO NOTHING" + }; + let payload = serde_json::to_string(&write.payload) + .map_err(|e| sqlite_err("encode write payload", e))?; + stored += conn + .execute( + sql, + params![ + config.thread_id, + namespace_json, + checkpoint_id, + write.task_id.as_str(), + write.idx, + write.node.as_str(), + write.channel, + payload, + ], + ) + .map_err(|e| sqlite_err("insert checkpoint write", e))?; + } + Ok(stored) +} + #[async_trait] impl Checkpointer for SqliteCheckpointer where @@ -221,42 +500,51 @@ where // never stalls a tokio worker on the step-critical path. let conn = self.conn.clone(); tokio::task::spawn_blocking(move || -> Result<()> { - let meta = checkpoint.to_metadata(); - let namespace = serde_json::to_string(&checkpoint.namespace) - .map_err(|e| sqlite_err("encode namespace", e))?; - let next_nodes = serde_json::to_string(&checkpoint.next_nodes) - .map_err(|e| sqlite_err("encode next_nodes", e))?; - let record = - serde_json::to_string(&checkpoint).map_err(|e| sqlite_err("encode record", e))?; + let conn = lock_conn(&conn)?; + insert_checkpoint_row(&conn, &checkpoint) + }) + .await + .map_err(|e| sqlite_err("join blocking put task", e))??; + Ok(id) + } - let conn = conn.lock().map_err(|_| { - TinyAgentsError::Checkpoint( - "sqlite checkpointer: connection lock poisoned".to_string(), - ) - })?; - conn.execute( - "INSERT INTO checkpoints ( - thread_id, checkpoint_id, parent_checkpoint_id, run_id, - namespace, next_nodes, source, step, has_interrupts, record - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", - params![ - checkpoint.thread_id, - checkpoint.checkpoint_id, - checkpoint.parent_checkpoint_id, - checkpoint.run_id, - namespace, - next_nodes, - meta.source.as_str(), - meta.step as i64, - i64::from(meta.has_interrupts), - record, - ], - ) - .map_err(|e| sqlite_err("insert checkpoint", e))?; + async fn put_with_writes( + &self, + checkpoint: Checkpoint, + writes: &[PendingWrite], + ) -> Result { + // One transaction covering both the checkpoint row and its writes — + // the boundary the executor commits at should never observe the + // checkpoint durable but its writes lost (or vice versa) to a crash + // between two separate autocommit statements. + let id = CheckpointId::new(checkpoint.checkpoint_id.clone()); + if writes.is_empty() { + // Nothing to share a transaction with; `put` alone is already one + // statement. + self.put(checkpoint).await?; + return Ok(id); + } + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let checkpoint_id = checkpoint.checkpoint_id.clone(); + let writes = writes.to_vec(); + let conn = self.conn.clone(); + tokio::task::spawn_blocking(move || -> Result<()> { + let mut conn = lock_conn(&conn)?; + let tx = conn + .transaction() + .map_err(|e| sqlite_err("begin put_with_writes", e))?; + insert_checkpoint_row(&tx, &checkpoint)?; + insert_checkpoint_writes(&tx, &config, &checkpoint_id, &writes)?; + tx.commit() + .map_err(|e| sqlite_err("commit put_with_writes", e))?; Ok(()) }) .await - .map_err(|e| sqlite_err("join blocking put task", e))??; + .map_err(|e| sqlite_err("join blocking put_with_writes task", e))??; Ok(id) } @@ -265,39 +553,48 @@ where thread_id: &str, checkpoint_id: Option<&str>, ) -> Result>> { - let conn = self.lock()?; - // Latest matching row (highest seq) for either the whole thread or a - // specific id, mirroring the append-only history of the other backends. - let record: Option = match checkpoint_id { - Some(id) => conn - .query_row( - "SELECT record FROM checkpoints - WHERE thread_id = ?1 AND checkpoint_id = ?2 - ORDER BY seq DESC LIMIT 1", - params![thread_id, id], - |row| row.get(0), - ) - .optional() - .map_err(|e| sqlite_err("query checkpoint", e))?, - None => conn - .query_row( - "SELECT record FROM checkpoints - WHERE thread_id = ?1 - ORDER BY seq DESC LIMIT 1", - params![thread_id], - |row| row.get(0), - ) - .optional() - .map_err(|e| sqlite_err("query latest checkpoint", e))?, - }; - match record { - Some(json) => { - Ok(Some(serde_json::from_str(&json).map_err(|e| { - decode_json_err("sqlite checkpointer", "record", e) - })?)) + let conn = self.conn.clone(); + let thread_id = thread_id.to_string(); + let checkpoint_id = checkpoint_id.map(|s| s.to_string()); + tokio::task::spawn_blocking(move || -> Result>> { + let conn = lock_conn(&conn)?; + // Latest matching row (highest seq) for either the whole thread or + // a specific id, mirroring the append-only history of the other + // backends. + let record: Option = match &checkpoint_id { + Some(id) => conn + .query_row( + "SELECT record FROM checkpoints + WHERE thread_id = ?1 AND checkpoint_id = ?2 + ORDER BY seq DESC LIMIT 1", + params![thread_id, id], + |row| row.get(0), + ) + .optional() + .map_err(|e| sqlite_err("query checkpoint", e))?, + None => conn + .query_row( + "SELECT record FROM checkpoints + WHERE thread_id = ?1 + ORDER BY seq DESC LIMIT 1", + params![thread_id], + |row| row.get(0), + ) + .optional() + .map_err(|e| sqlite_err("query latest checkpoint", e))?, + }; + match record { + Some(json) => { + let mut checkpoint: Checkpoint = serde_json::from_str(&json) + .map_err(|e| decode_json_err("sqlite checkpointer", "record", e))?; + checkpoint.normalize(); + Ok(Some(checkpoint)) + } + None => Ok(None), } - None => Ok(None), - } + }) + .await + .map_err(|e| sqlite_err("join blocking get task", e))? } async fn get_scoped( @@ -309,39 +606,47 @@ where // Pushed down to one indexed query. The trait default lists the whole // thread and then re-`get`s the winner, which costs a full thread scan // per call — and `state_history` calls it once per lineage hop. + let conn = self.conn.clone(); + let thread_id = thread_id.to_string(); + let checkpoint_id = checkpoint_id.map(|s| s.to_string()); let namespace_json = serde_json::to_string(namespace).map_err(|e| sqlite_err("encode namespace", e))?; - let conn = self.lock()?; - let record: Option = match checkpoint_id { - Some(id) => conn - .query_row( - "SELECT record FROM checkpoints - WHERE thread_id = ?1 AND namespace = ?2 AND checkpoint_id = ?3 - ORDER BY seq DESC LIMIT 1", - params![thread_id, namespace_json, id], - |row| row.get(0), - ) - .optional() - .map_err(|e| sqlite_err("query scoped checkpoint", e))?, - None => conn - .query_row( - "SELECT record FROM checkpoints - WHERE thread_id = ?1 AND namespace = ?2 - ORDER BY seq DESC LIMIT 1", - params![thread_id, namespace_json], - |row| row.get(0), - ) - .optional() - .map_err(|e| sqlite_err("query latest scoped checkpoint", e))?, - }; - match record { - Some(json) => { - Ok(Some(serde_json::from_str(&json).map_err(|e| { - decode_json_err("sqlite checkpointer", "record", e) - })?)) + tokio::task::spawn_blocking(move || -> Result>> { + let conn = lock_conn(&conn)?; + let record: Option = match &checkpoint_id { + Some(id) => conn + .query_row( + "SELECT record FROM checkpoints + WHERE thread_id = ?1 AND namespace = ?2 AND checkpoint_id = ?3 + ORDER BY seq DESC LIMIT 1", + params![thread_id, namespace_json, id], + |row| row.get(0), + ) + .optional() + .map_err(|e| sqlite_err("query scoped checkpoint", e))?, + None => conn + .query_row( + "SELECT record FROM checkpoints + WHERE thread_id = ?1 AND namespace = ?2 + ORDER BY seq DESC LIMIT 1", + params![thread_id, namespace_json], + |row| row.get(0), + ) + .optional() + .map_err(|e| sqlite_err("query latest scoped checkpoint", e))?, + }; + match record { + Some(json) => { + let mut checkpoint: Checkpoint = serde_json::from_str(&json) + .map_err(|e| decode_json_err("sqlite checkpointer", "record", e))?; + checkpoint.normalize(); + Ok(Some(checkpoint)) + } + None => Ok(None), } - None => Ok(None), - } + }) + .await + .map_err(|e| sqlite_err("join blocking get_scoped task", e))? } async fn state_history( @@ -350,204 +655,247 @@ where namespace: &[String], limit: Option, ) -> Result>> { - // One indexed range read of the namespace's rows, then the lineage walk - // in memory — instead of the default's `get_tuple` (and therefore - // `get_scoped`) per hop. + // Walks the parent chain with a recursive SQL CTE so `LIMIT` is applied + // in SQL: a `state_history(Some(1))` call decodes exactly one record + // instead of every record in the namespace. See `STATE_HISTORY_CTE`'s + // doc comment for the query shape and the cycle-termination argument. + let conn = self.conn.clone(); + let thread_id_owned = thread_id.to_string(); let namespace_json = serde_json::to_string(namespace).map_err(|e| sqlite_err("encode namespace", e))?; - let (records, writes) = { - let conn = self.lock()?; - let mut stmt = conn - .prepare( - "SELECT record FROM checkpoints - WHERE thread_id = ?1 AND namespace = ?2 ORDER BY seq ASC", + tokio::task::spawn_blocking(move || -> Result>> { + let conn = lock_conn(&conn)?; + + // Total row count in the namespace both answers "is there + // anything at all" and, more importantly, bounds the CTE's + // recursion depth: it is a safe upper bound on the number of + // *distinct* checkpoint ids reachable, so capping recursion there + // guarantees termination even over a hand-corrupted or forked + // lineage with a parent cycle (the trait's documented hazard — + // `parent_checkpoint_id` is caller-set data, not a structurally + // acyclic pointer). + let total: i64 = conn + .query_row( + "SELECT COUNT(*) FROM checkpoints WHERE thread_id = ?1 AND namespace = ?2", + params![thread_id_owned, namespace_json], + |row| row.get(0), ) + .map_err(|e| sqlite_err("count state_history rows", e))?; + if total == 0 { + return Ok(Vec::new()); + } + let cap: i64 = match limit { + Some(limit) => (limit as i64).min(total), + None => total, + }; + if cap <= 0 { + return Ok(Vec::new()); + } + + let mut stmt = conn + .prepare(STATE_HISTORY_CTE) .map_err(|e| sqlite_err("prepare state_history", e))?; let rows = stmt - .query_map(params![thread_id, namespace_json], |row| { + .query_map(params![thread_id_owned, namespace_json, cap], |row| { row.get::<_, String>(0) }) .map_err(|e| sqlite_err("query state_history", e))?; let mut records: Vec> = Vec::new(); for row in rows { let json = row.map_err(|e| sqlite_err("read record row", e))?; - records.push( - serde_json::from_str(&json) - .map_err(|e| decode_json_err("sqlite checkpointer", "record", e))?, - ); + let mut checkpoint: Checkpoint = serde_json::from_str(&json) + .map_err(|e| decode_json_err("sqlite checkpointer", "record", e))?; + checkpoint.normalize(); + records.push(checkpoint); } - let writes = read_writes_by_checkpoint(&conn, thread_id, &namespace_json)?; - (records, writes) - }; - if records.is_empty() { - return Ok(Vec::new()); - } - - // Last write wins for a re-used id, matching `get`. - let mut by_id: std::collections::HashMap> = - std::collections::HashMap::with_capacity(records.len()); - let mut cursor: Option = None; - for record in records { - cursor = Some(record.checkpoint_id.clone()); - by_id.insert(record.checkpoint_id.clone(), record); - } + if records.is_empty() { + return Ok(Vec::new()); + } + let writes = read_writes_by_checkpoint(&conn, &thread_id_owned, &namespace_json)?; - let mut out = Vec::new(); - while let Some(id) = cursor { - if let Some(limit) = limit - && out.len() >= limit - { - break; + // The CTE already returns newest-first (depth ascending from the + // head), so no further sorting or in-memory lineage walk is + // needed here. + let mut out = Vec::with_capacity(records.len()); + for checkpoint in records { + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let parent_config = + checkpoint + .parent_checkpoint_id + .as_ref() + .map(|parent| CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(parent.clone()), + namespace: checkpoint.namespace.clone(), + }); + let pending_writes = writes + .get(&checkpoint.checkpoint_id) + .cloned() + .unwrap_or_else(|| checkpoint.pending_writes.clone()); + out.push(CheckpointTuple { + config, + checkpoint, + parent_config, + pending_writes, + }); } - // `remove` doubles as the cycle guard: each id is visited once. - let Some(checkpoint) = by_id.remove(&id) else { - break; - }; - cursor = checkpoint.parent_checkpoint_id.clone(); - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; - let parent_config = - checkpoint - .parent_checkpoint_id - .as_ref() - .map(|parent| CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(parent.clone()), - namespace: checkpoint.namespace.clone(), - }); - let pending_writes = writes - .get(&checkpoint.checkpoint_id) - .cloned() - .unwrap_or_else(|| checkpoint.pending_writes.clone()); - out.push(CheckpointTuple { - config, - checkpoint, - parent_config, - pending_writes, - }); - } - Ok(out) + Ok(out) + }) + .await + .map_err(|e| sqlite_err("join blocking state_history task", e))? } async fn list(&self, thread_id: &str) -> Result> { - let conn = self.lock()?; - let mut stmt = conn - .prepare( - "SELECT thread_id, checkpoint_id, run_id, parent_checkpoint_id, - namespace, next_nodes, source, step, has_interrupts - FROM checkpoints WHERE thread_id = ?1 ORDER BY seq ASC", - ) - .map_err(|e| sqlite_err("prepare list", e))?; - let rows = stmt - .query_map(params![thread_id], |row| { - Ok(MetaRow { - thread_id: row.get(0)?, - checkpoint_id: row.get(1)?, - run_id: row.get(2)?, - parent_checkpoint_id: row.get(3)?, - namespace_json: row.get(4)?, - next_nodes_json: row.get(5)?, - source: row.get(6)?, - step: row.get(7)?, - has_interrupts: row.get(8)?, + let conn = self.conn.clone(); + let thread_id = thread_id.to_string(); + tokio::task::spawn_blocking(move || -> Result> { + let conn = lock_conn(&conn)?; + let mut stmt = conn + .prepare( + "SELECT thread_id, checkpoint_id, run_id, parent_checkpoint_id, + namespace, next_nodes, source, step, has_interrupts + FROM checkpoints WHERE thread_id = ?1 ORDER BY seq ASC", + ) + .map_err(|e| sqlite_err("prepare list", e))?; + let rows = stmt + .query_map(params![thread_id], |row| { + Ok(MetaRow { + thread_id: row.get(0)?, + checkpoint_id: row.get(1)?, + run_id: row.get(2)?, + parent_checkpoint_id: row.get(3)?, + namespace_json: row.get(4)?, + next_nodes_json: row.get(5)?, + source: row.get(6)?, + step: row.get(7)?, + has_interrupts: row.get(8)?, + }) }) - }) - .map_err(|e| sqlite_err("query list", e))?; - let mut out = Vec::new(); - for row in rows { - out.push(row_metadata( - row.map_err(|e| sqlite_err("read list row", e))?, - )?); - } - Ok(out) + .map_err(|e| sqlite_err("query list", e))?; + let mut out = Vec::new(); + for row in rows { + out.push(row_metadata( + row.map_err(|e| sqlite_err("read list row", e))?, + )?); + } + Ok(out) + }) + .await + .map_err(|e| sqlite_err("join blocking list task", e))? } async fn get_thread(&self, thread_id: &str) -> Result>> { // Single-pass bulk read: one indexed range query over the thread's // rows in insertion order, instead of the default's one point query // per listed id. - let conn = self.lock()?; - let mut stmt = conn - .prepare("SELECT record FROM checkpoints WHERE thread_id = ?1 ORDER BY seq ASC") - .map_err(|e| sqlite_err("prepare get_thread", e))?; - let rows = stmt - .query_map(params![thread_id], |row| row.get::<_, String>(0)) - .map_err(|e| sqlite_err("query get_thread", e))?; - let mut out = Vec::new(); - for row in rows { - let json = row.map_err(|e| sqlite_err("read record row", e))?; - out.push( - serde_json::from_str(&json) - .map_err(|e| decode_json_err("sqlite checkpointer", "record", e))?, - ); - } - Ok(out) + let conn = self.conn.clone(); + let thread_id = thread_id.to_string(); + tokio::task::spawn_blocking(move || -> Result>> { + let conn = lock_conn(&conn)?; + let mut stmt = conn + .prepare("SELECT record FROM checkpoints WHERE thread_id = ?1 ORDER BY seq ASC") + .map_err(|e| sqlite_err("prepare get_thread", e))?; + let rows = stmt + .query_map(params![thread_id], |row| row.get::<_, String>(0)) + .map_err(|e| sqlite_err("query get_thread", e))?; + let mut out = Vec::new(); + for row in rows { + let json = row.map_err(|e| sqlite_err("read record row", e))?; + let mut checkpoint: Checkpoint = serde_json::from_str(&json) + .map_err(|e| decode_json_err("sqlite checkpointer", "record", e))?; + checkpoint.normalize(); + out.push(checkpoint); + } + Ok(out) + }) + .await + .map_err(|e| sqlite_err("join blocking get_thread task", e))? } async fn list_threads(&self) -> Result> { - let conn = self.lock()?; - let mut stmt = conn - .prepare("SELECT DISTINCT thread_id FROM checkpoints") - .map_err(|e| sqlite_err("prepare list_threads", e))?; - let rows = stmt - .query_map([], |row| row.get::<_, String>(0)) - .map_err(|e| sqlite_err("query list_threads", e))?; - let mut out = Vec::new(); - for row in rows { - out.push(row.map_err(|e| sqlite_err("read thread row", e))?); - } - Ok(out) + let conn = self.conn.clone(); + tokio::task::spawn_blocking(move || -> Result> { + let conn = lock_conn(&conn)?; + let mut stmt = conn + .prepare("SELECT DISTINCT thread_id FROM checkpoints") + .map_err(|e| sqlite_err("prepare list_threads", e))?; + let rows = stmt + .query_map([], |row| row.get::<_, String>(0)) + .map_err(|e| sqlite_err("query list_threads", e))?; + let mut out = Vec::new(); + for row in rows { + out.push(row.map_err(|e| sqlite_err("read thread row", e))?); + } + Ok(out) + }) + .await + .map_err(|e| sqlite_err("join blocking list_threads task", e))? } async fn delete_thread(&self, thread_id: &str) -> Result<()> { - let mut conn = self.lock()?; - let tx = conn - .transaction() - .map_err(|e| sqlite_err("begin delete_thread", e))?; - tx.execute( - "DELETE FROM checkpoints WHERE thread_id = ?1", - params![thread_id], - ) - .map_err(|e| sqlite_err("delete thread", e))?; - // Writes go with the thread — and across *every* namespace, not just - // the root one, or an embedded subgraph's ledger outlives its thread. - tx.execute( - "DELETE FROM checkpoint_writes WHERE thread_id = ?1", - params![thread_id], - ) - .map_err(|e| sqlite_err("delete thread writes", e))?; - tx.commit() - .map_err(|e| sqlite_err("commit delete_thread", e))?; - Ok(()) + let conn = self.conn.clone(); + let thread_id = thread_id.to_string(); + tokio::task::spawn_blocking(move || -> Result<()> { + let mut conn = lock_conn(&conn)?; + let tx = conn + .transaction() + .map_err(|e| sqlite_err("begin delete_thread", e))?; + tx.execute( + "DELETE FROM checkpoints WHERE thread_id = ?1", + params![thread_id], + ) + .map_err(|e| sqlite_err("delete thread", e))?; + // Writes go with the thread — and across *every* namespace, not + // just the root one, or an embedded subgraph's ledger outlives + // its thread. + tx.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = ?1", + params![thread_id], + ) + .map_err(|e| sqlite_err("delete thread writes", e))?; + tx.commit() + .map_err(|e| sqlite_err("commit delete_thread", e))?; + Ok(()) + }) + .await + .map_err(|e| sqlite_err("join blocking delete_thread task", e))? } async fn delete_checkpoints(&self, thread_id: &str, ids: &[String]) -> Result { if ids.is_empty() { return Ok(0); } - let mut conn = self.lock()?; - let tx = conn - .transaction() - .map_err(|e| sqlite_err("begin transaction", e))?; - let mut removed = 0usize; - for id in ids { - removed += tx - .execute( - "DELETE FROM checkpoints WHERE thread_id = ?1 AND checkpoint_id = ?2", + let conn = self.conn.clone(); + let thread_id = thread_id.to_string(); + let ids = ids.to_vec(); + tokio::task::spawn_blocking(move || -> Result { + let mut conn = lock_conn(&conn)?; + let tx = conn + .transaction() + .map_err(|e| sqlite_err("begin transaction", e))?; + let mut removed = 0usize; + for id in &ids { + removed += tx + .execute( + "DELETE FROM checkpoints WHERE thread_id = ?1 AND checkpoint_id = ?2", + params![thread_id, id], + ) + .map_err(|e| sqlite_err("delete checkpoint", e))?; + tx.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = ?1 AND checkpoint_id = ?2", params![thread_id, id], ) - .map_err(|e| sqlite_err("delete checkpoint", e))?; - tx.execute( - "DELETE FROM checkpoint_writes WHERE thread_id = ?1 AND checkpoint_id = ?2", - params![thread_id, id], - ) - .map_err(|e| sqlite_err("delete checkpoint writes", e))?; - } - tx.commit().map_err(|e| sqlite_err("commit delete", e))?; - Ok(removed) + .map_err(|e| sqlite_err("delete checkpoint writes", e))?; + } + tx.commit().map_err(|e| sqlite_err("commit delete", e))?; + Ok(removed) + }) + .await + .map_err(|e| sqlite_err("join blocking delete_checkpoints task", e))? } async fn put_writes(&self, config: &CheckpointConfig, writes: &[PendingWrite]) -> Result<()> { @@ -555,86 +903,152 @@ where if writes.is_empty() { return Ok(()); } - let namespace_json = serde_json::to_string(&config.namespace) - .map_err(|e| sqlite_err("encode namespace", e))?; - let mut conn = self.lock()?; - let tx = conn - .transaction() - .map_err(|e| sqlite_err("begin put_writes", e))?; - let mut _stored = 0usize; - for write in writes { - // The replace-vs-ignore rule pushed into SQL: a control-plane write - // (`idx < 0`) legitimately changes on a retry and upserts, while a - // data write is append-once so a retried `put_writes` is a no-op. - // Doing it with two conflict clauses rather than a read-then-write - // keeps it correct under concurrent writers. - let sql = if write.is_control_plane() { - "INSERT INTO checkpoint_writes - (thread_id, namespace, checkpoint_id, task_id, idx, node, channel, payload) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) - ON CONFLICT(thread_id, namespace, checkpoint_id, task_id, idx) DO UPDATE SET - node = excluded.node, - channel = excluded.channel, - payload = excluded.payload" - } else { - "INSERT INTO checkpoint_writes - (thread_id, namespace, checkpoint_id, task_id, idx, node, channel, payload) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) - ON CONFLICT(thread_id, namespace, checkpoint_id, task_id, idx) DO NOTHING" - }; - let payload = serde_json::to_string(&write.payload) - .map_err(|e| sqlite_err("encode write payload", e))?; - _stored += tx - .execute( - sql, - params![ - config.thread_id, - namespace_json, - checkpoint_id, - write.task_id, - write.idx, - write.node.as_str(), - write.channel, - payload, - ], - ) - .map_err(|e| sqlite_err("insert checkpoint write", e))?; - } - tx.commit() - .map_err(|e| sqlite_err("commit put_writes", e))?; - tinyagents_tracing::debug!( - "[checkpoint:sqlite] put_writes thread={} checkpoint={checkpoint_id} offered={} stored={_stored}", - config.thread_id, - writes.len() - ); - Ok(()) + let config = config.clone(); + let writes = writes.to_vec(); + let conn = self.conn.clone(); + tokio::task::spawn_blocking(move || -> Result<()> { + let mut conn = lock_conn(&conn)?; + let tx = conn + .transaction() + .map_err(|e| sqlite_err("begin put_writes", e))?; + let stored = insert_checkpoint_writes(&tx, &config, &checkpoint_id, &writes)?; + tx.commit() + .map_err(|e| sqlite_err("commit put_writes", e))?; + tracing::debug!( + "[checkpoint:sqlite] put_writes thread={} checkpoint={checkpoint_id} offered={} stored={stored}", + config.thread_id, + writes.len() + ); + Ok(()) + }) + .await + .map_err(|e| sqlite_err("join blocking put_writes task", e))? } async fn get_writes(&self, config: &CheckpointConfig) -> Result> { let Some(checkpoint_id) = self.resolve_write_target(config).await? else { return Ok(Vec::new()); }; - let namespace_json = serde_json::to_string(&config.namespace) - .map_err(|e| sqlite_err("encode namespace", e))?; - let conn = self.lock()?; - let mut stmt = conn - .prepare( - "SELECT node, task_id, idx, channel, payload FROM checkpoint_writes - WHERE thread_id = ?1 AND namespace = ?2 AND checkpoint_id = ?3 - ORDER BY rowid ASC", + let config = config.clone(); + let conn = self.conn.clone(); + tokio::task::spawn_blocking(move || -> Result> { + let namespace_json = serde_json::to_string(&config.namespace) + .map_err(|e| sqlite_err("encode namespace", e))?; + let conn = lock_conn(&conn)?; + let mut stmt = conn + .prepare( + "SELECT node, task_id, idx, channel, payload FROM checkpoint_writes + WHERE thread_id = ?1 AND namespace = ?2 AND checkpoint_id = ?3 + ORDER BY rowid ASC", + ) + .map_err(|e| sqlite_err("prepare get_writes", e))?; + let rows = stmt + .query_map( + params![config.thread_id, namespace_json, checkpoint_id], + map_write_row, + ) + .map_err(|e| sqlite_err("query get_writes", e))?; + let mut out = Vec::new(); + for row in rows { + out.push(row.map_err(|e| sqlite_err("read write row", e))??); + } + Ok(out) + }) + .await + .map_err(|e| sqlite_err("join blocking get_writes task", e))? + } + + async fn try_claim(&self, thread: &str, owner: &str, ttl: std::time::Duration) -> Result { + let conn = self.conn.clone(); + let thread = thread.to_string(); + let owner = owner.to_string(); + let ttl_ms = ttl.as_millis() as i64; + tokio::task::spawn_blocking(move || -> Result { + let now = tinyagents_harness::ids::now_ms() as i64; + let expires_at = now.saturating_add(ttl_ms); + let conn = conn.lock().map_err(|_| { + TinyAgentsError::Checkpoint( + "sqlite checkpointer: connection lock poisoned".to_string(), + ) + })?; + let tx = conn + .unchecked_transaction() + .map_err(|e| sqlite_err("begin try_claim tx", e))?; + let existing: Option<(String, i64)> = tx + .query_row( + "SELECT owner, expires_at FROM thread_leases WHERE thread_id = ?1", + params![thread], + |row| Ok((row.get(0)?, row.get(1)?)), + ) + .optional() + .map_err(|e| sqlite_err("read thread_leases", e))?; + let claimable = match &existing { + None => true, + Some((existing_owner, _)) if existing_owner == &owner => true, + Some((_, existing_expires)) => *existing_expires <= now, + }; + if !claimable { + tx.commit().map_err(|e| sqlite_err("commit try_claim", e))?; + return Ok(false); + } + tx.execute( + "INSERT INTO thread_leases (thread_id, owner, expires_at) VALUES (?1, ?2, ?3) + ON CONFLICT(thread_id) DO UPDATE SET owner = excluded.owner, expires_at = excluded.expires_at", + params![thread, owner, expires_at], ) - .map_err(|e| sqlite_err("prepare get_writes", e))?; - let rows = stmt - .query_map( - params![config.thread_id, namespace_json, checkpoint_id], - map_write_row, + .map_err(|e| sqlite_err("upsert thread_leases", e))?; + tx.commit().map_err(|e| sqlite_err("commit try_claim", e))?; + Ok(true) + }) + .await + .map_err(|e| sqlite_err("join blocking try_claim task", e))? + } + + async fn renew(&self, thread: &str, owner: &str, ttl: std::time::Duration) -> Result { + let conn = self.conn.clone(); + let thread = thread.to_string(); + let owner = owner.to_string(); + let ttl_ms = ttl.as_millis() as i64; + tokio::task::spawn_blocking(move || -> Result { + let now = tinyagents_harness::ids::now_ms() as i64; + let expires_at = now.saturating_add(ttl_ms); + let conn = conn.lock().map_err(|_| { + TinyAgentsError::Checkpoint( + "sqlite checkpointer: connection lock poisoned".to_string(), + ) + })?; + let updated = conn + .execute( + "UPDATE thread_leases SET expires_at = ?1 + WHERE thread_id = ?2 AND owner = ?3", + params![expires_at, thread, owner], + ) + .map_err(|e| sqlite_err("renew thread_leases", e))?; + Ok(updated > 0) + }) + .await + .map_err(|e| sqlite_err("join blocking renew task", e))? + } + + async fn release(&self, thread: &str, owner: &str) -> Result<()> { + let conn = self.conn.clone(); + let thread = thread.to_string(); + let owner = owner.to_string(); + tokio::task::spawn_blocking(move || -> Result<()> { + let conn = conn.lock().map_err(|_| { + TinyAgentsError::Checkpoint( + "sqlite checkpointer: connection lock poisoned".to_string(), + ) + })?; + conn.execute( + "DELETE FROM thread_leases WHERE thread_id = ?1 AND owner = ?2", + params![thread, owner], ) - .map_err(|e| sqlite_err("query get_writes", e))?; - let mut out = Vec::new(); - for row in rows { - out.push(row.map_err(|e| sqlite_err("read write row", e))??); - } - Ok(out) + .map_err(|e| sqlite_err("release thread_leases", e))?; + Ok(()) + }) + .await + .map_err(|e| sqlite_err("join blocking release task", e))? } } @@ -653,7 +1067,7 @@ fn map_write_row(row: &rusqlite::Row<'_>) -> rusqlite::Result(&payload_json) { Ok(payload) => Ok(PendingWrite { node: NodeId::from(node), - task_id, + task_id: TaskId::from(task_id), idx, channel, payload, @@ -698,7 +1112,7 @@ fn read_writes_by_checkpoint( .map_err(|e| decode_json_err("sqlite checkpointer", "write payload", e))?; let write = PendingWrite { node: NodeId::from(node), - task_id, + task_id: TaskId::from(task_id), idx, channel, payload, diff --git a/crates/tinyagents-graph/src/checkpoint/test.rs b/crates/tinyagents-graph/src/checkpoint/test.rs index 54636caf..49c1d0e9 100644 --- a/crates/tinyagents-graph/src/checkpoint/test.rs +++ b/crates/tinyagents-graph/src/checkpoint/test.rs @@ -7,21 +7,18 @@ use serde_json::json; use tinyagents_harness::ids::NodeId; fn checkpoint(thread: &str, id: &str, parent: Option<&str>, step: usize) -> Checkpoint { - Checkpoint { - thread_id: thread.to_string(), - checkpoint_id: id.to_string(), - run_id: None, - parent_checkpoint_id: parent.map(|s| s.to_string()), - namespace: vec![], - state: step as i32, - next_nodes: vec![NodeId::from("n")], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![], - pending_activations: None, - barrier_arrivals: vec![], - metadata: json!({ "source": "loop", "step": step }), - } + Checkpoint::new( + step as i32, + vec![PendingActivation { + node: NodeId::from("n"), + send_arg: None, + task_id: tinyagents_harness::ids::TaskId::from(String::new()), + }], + ) + .with_thread_id(thread.to_string()) + .with_checkpoint_id(id.to_string()) + .with_parent_checkpoint_id(parent.map(|s| s.to_string())) + .with_metadata(json!({ "source": "loop", "step": step })) } #[tokio::test] @@ -69,41 +66,48 @@ fn legacy_checkpoint_json_without_new_fields_still_loads() { "interrupts": [], "metadata": { "source": "loop", "step": 1 } }); - let cp: Checkpoint = serde_json::from_value(legacy).unwrap(); + let mut cp: Checkpoint = serde_json::from_value(legacy).unwrap(); assert_eq!(cp.state, 7); + // Un-normalized: decodes as a v1 record with the legacy fields intact and + // `tasks`/`completed` still empty. + assert_eq!(cp.version, 1); assert_eq!(cp.next_nodes.len(), 2); + assert!(cp.tasks.is_empty()); assert!(cp.pending_activations.is_none()); assert!(cp.barrier_arrivals.is_empty()); + + // `normalize()` folds the legacy fields into the v2 shape and clears them. + cp.normalize(); + assert_eq!(cp.version, CHECKPOINT_FORMAT_VERSION); + assert_eq!(cp.tasks.len(), 2); + assert_eq!(cp.tasks[0].node, NodeId::from("a")); + assert!(cp.next_nodes.is_empty()); } #[test] fn pending_activation_send_arg_roundtrips() { - let cp = Checkpoint { - thread_id: "t".into(), - checkpoint_id: "c1".into(), - run_id: None, - parent_checkpoint_id: None, - namespace: vec![], - state: 1i32, - next_nodes: vec![NodeId::from("w")], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![], - pending_activations: Some(vec![super::PendingActivation { + let cp = Checkpoint::new( + 1i32, + vec![super::PendingActivation { node: NodeId::from("w"), - send_arg: Some(json!({ "item": 42 })), - task_id: "1:0:w".to_string(), - }]), - barrier_arrivals: vec![super::BarrierArrivals { - node: NodeId::from("join"), - arrived: vec![NodeId::from("p1")], + send_arg: Some(std::sync::Arc::new(json!({ "item": 42 }))), + task_id: tinyagents_harness::ids::TaskId::from("1:0:w"), }], - metadata: json!({ "source": "loop", "step": 1 }), - }; + ) + .with_thread_id("t") + .with_checkpoint_id("c1") + .with_barrier_arrivals(vec![super::BarrierArrivals { + node: NodeId::from("join"), + arrived: vec![NodeId::from("p1")], + }]) + .with_metadata(json!({ "source": "loop", "step": 1 })); let round: Checkpoint = serde_json::from_str(&serde_json::to_string(&cp).unwrap()).unwrap(); - let pa = round.pending_activations.unwrap(); - assert_eq!(pa[0].send_arg, Some(json!({ "item": 42 }))); + assert_eq!(round.version, super::CHECKPOINT_FORMAT_VERSION); + assert_eq!( + round.tasks[0].send_arg, + Some(std::sync::Arc::new(json!({ "item": 42 }))) + ); assert_eq!(round.barrier_arrivals[0].arrived, vec![NodeId::from("p1")]); } @@ -380,7 +384,8 @@ async fn prune_keeps_a_window_per_namespace() { mod file_backend { use super::checkpoint; - use crate::checkpoint::{CheckpointConfig, Checkpointer, FileCheckpointer}; + use crate::Checkpoint; + use crate::checkpoint::{CheckpointConfig, Checkpointer, FileCheckpointer, PendingActivation}; use std::path::PathBuf; /// A unique-per-test temp dir derived from the test name + pid (no clock). @@ -408,6 +413,127 @@ mod file_backend { } } + /// The on-disk shape a pre-v2 build wrote: no `version`/`created_at`/ + /// `tasks`/`completed` fields at all, just the v1 + /// `next_nodes`/`completed_tasks`/`completed_routes`/`pending_activations` + /// quartet. Hand-written (not produced by this build) so the test proves + /// the *wire format*, not just today's `Checkpoint::normalize` logic + /// agreeing with itself. + fn v1_fixture_line(thread: &str, id: &str, parent: Option<&str>, step: usize) -> String { + serde_json::json!({ + "thread_id": thread, + "checkpoint_id": id, + "run_id": null, + "parent_checkpoint_id": parent, + "namespace": [], + "state": step as i64, + "next_nodes": ["b"], + "completed_tasks": ["a"], + "completed_routes": [[]], + "pending_writes": [], + "interrupts": [], + "pending_activations": null, + "barrier_arrivals": [], + "metadata": { "source": "loop", "step": step }, + }) + .to_string() + } + + #[tokio::test] + async fn v1_fixture_decodes_to_normalized_v2_and_resumes() { + let tmp = TempDir::new("v1-fixture"); + std::fs::create_dir_all(tmp.path()).unwrap(); + // Write the fixture line directly — bypassing `put`, which (this + // build) only ever writes v2 — to prove the *decode* path, not just + // `Checkpoint::normalize` called directly on a value built in Rust. + std::fs::write( + tmp.path().join("v1thread.jsonl"), + format!("{}\n", v1_fixture_line("v1thread", "c1", None, 1)), + ) + .unwrap(); + + let cp = FileCheckpointer::::new(tmp.path()); + let loaded = cp.get("v1thread", None).await.unwrap().unwrap(); + assert_eq!(loaded.version, crate::checkpoint::CHECKPOINT_FORMAT_VERSION); + assert_eq!( + loaded + .tasks + .iter() + .map(|t| t.node.to_string()) + .collect::>(), + vec!["b".to_string()], + "tasks derived from the v1 next_nodes field" + ); + assert_eq!( + loaded + .completed + .iter() + .map(|c| c.node.to_string()) + .collect::>(), + vec!["a".to_string()], + "completed derived from the v1 completed_tasks/completed_routes pair" + ); + assert!( + loaded.next_nodes.is_empty(), + "legacy fields cleared by normalize" + ); + assert!(loaded.completed_tasks.is_empty()); + + // get_scoped, list, state_history, and get_thread all go through the + // same normalize-on-decode path. + let scoped = cp.get_scoped("v1thread", None, &[]).await.unwrap().unwrap(); + assert_eq!(scoped.version, crate::checkpoint::CHECKPOINT_FORMAT_VERSION); + let listed = cp.list("v1thread").await.unwrap(); + assert_eq!( + listed[0].next_nodes, + vec![tinyagents_harness::ids::NodeId::from("b")] + ); + let history = cp.state_history("v1thread", &[], None).await.unwrap(); + assert_eq!(history.len(), 1); + assert_eq!( + history[0].checkpoint.version, + crate::checkpoint::CHECKPOINT_FORMAT_VERSION + ); + let thread = cp.get_thread("v1thread").await.unwrap(); + assert_eq!( + thread[0].version, + crate::checkpoint::CHECKPOINT_FORMAT_VERSION + ); + } + + #[tokio::test] + async fn mixed_v1_and_v2_thread_lists_and_walks_state_history() { + let tmp = TempDir::new("mixed-v1-v2"); + std::fs::create_dir_all(tmp.path()).unwrap(); + // c1: hand-written v1 fixture. c2: written through `put`, which is + // always v2. Same thread, same file. + std::fs::write( + tmp.path().join("mixedthread.jsonl"), + format!("{}\n", v1_fixture_line("mixedthread", "c1", None, 1)), + ) + .unwrap(); + let cp = FileCheckpointer::::new(tmp.path()); + cp.put(checkpoint("mixedthread", "c2", Some("c1"), 2)) + .await + .unwrap(); + + let list = cp.list("mixedthread").await.unwrap(); + assert_eq!(list.len(), 2, "both the v1 and v2 record are listed"); + assert_eq!(list[0].checkpoint_id, "c1"); + assert_eq!(list[1].checkpoint_id, "c2"); + + let history = cp.state_history("mixedthread", &[], None).await.unwrap(); + assert_eq!(history.len(), 2, "the walk crosses the v1/v2 boundary"); + assert_eq!(history[0].checkpoint.checkpoint_id, "c2"); + assert_eq!(history[1].checkpoint.checkpoint_id, "c1"); + // Both normalize to v2 regardless of which format they were stored in. + assert!( + history + .iter() + .all(|t| t.checkpoint.version == crate::checkpoint::CHECKPOINT_FORMAT_VERSION) + ); + } + #[tokio::test] async fn put_get_list_roundtrip_survives_a_fresh_handle() { let tmp = TempDir::new("roundtrip"); @@ -556,6 +682,117 @@ mod file_backend { assert_eq!(records[1].state, 2); assert!(cp.get_thread("missing").await.unwrap().is_empty()); } + + #[tokio::test] + async fn concurrent_file_lease_claims_have_one_winner() { + let tmp = TempDir::new("concurrent-lease-claim"); + // Use separate handles to the same directory, matching independent + // executors rather than relying on any in-process coordination. + let first = FileCheckpointer::::new(tmp.path()); + let second = FileCheckpointer::::new(tmp.path()); + let barrier = std::sync::Arc::new(tokio::sync::Barrier::new(2)); + + let first_barrier = std::sync::Arc::clone(&barrier); + let first_claim = tokio::spawn(async move { + first_barrier.wait().await; + first + .try_claim("thread", "owner-a", std::time::Duration::from_secs(60)) + .await + }); + let second_barrier = std::sync::Arc::clone(&barrier); + let second_claim = tokio::spawn(async move { + second_barrier.wait().await; + second + .try_claim("thread", "owner-b", std::time::Duration::from_secs(60)) + .await + }); + + let first_won = first_claim.await.unwrap().unwrap(); + let second_won = second_claim.await.unwrap().unwrap(); + assert_ne!( + first_won, second_won, + "only one concurrent owner may claim a file-backed lease" + ); + } + + // ---- I9 regression: `list` must not decode full `State` ----------------- + + /// A `State` whose `Deserialize` impl counts every call it makes, so a + /// test can assert *how many times* something deserialized it rather than + /// just observing the (correct either way) return value. + #[derive(Clone, serde::Serialize)] + struct CountedState(i32); + + /// Process-wide count of `CountedState` deserializations. `CountedState` + /// is private to this test module, so nothing outside these tests can + /// bump it — safe to share across the (OS-threaded) test binary without a + /// dedicated fixture. + static STATE_DECODE_COUNT: std::sync::atomic::AtomicUsize = + std::sync::atomic::AtomicUsize::new(0); + + impl<'de> serde::Deserialize<'de> for CountedState { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + STATE_DECODE_COUNT.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + i32::deserialize(deserializer).map(CountedState) + } + } + + fn counted_checkpoint( + thread: &str, + id: &str, + parent: Option<&str>, + step: usize, + ) -> Checkpoint { + Checkpoint::new( + CountedState(step as i32), + vec![PendingActivation { + node: tinyagents_harness::ids::NodeId::from("n"), + send_arg: None, + task_id: tinyagents_harness::ids::TaskId::from(String::new()), + }], + ) + .with_thread_id(thread.to_string()) + .with_checkpoint_id(id.to_string()) + .with_parent_checkpoint_id(parent.map(|s| s.to_string())) + .with_metadata(serde_json::json!({ "source": "loop", "step": step })) + } + + #[tokio::test] + async fn list_on_a_large_thread_does_not_decode_full_state() { + let tmp = TempDir::new("list-header-only"); + let cp = FileCheckpointer::::new(tmp.path()); + + let mut parent: Option = None; + for step in 0..200usize { + let id = format!("c{step}"); + cp.put(counted_checkpoint("t", &id, parent.as_deref(), step)) + .await + .unwrap(); + parent = Some(id); + } + + // `put` only serializes, so the counter should already read 0 here; + // reset explicitly anyway so this assertion is about `list` alone, + // not an assumption about what came before it. + STATE_DECODE_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + + let list = cp.list("t").await.unwrap(); + assert_eq!( + list.len(), + 200, + "list still returns every record's metadata" + ); + assert_eq!(list[0].checkpoint_id, "c0"); + assert_eq!(list[199].checkpoint_id, "c199"); + assert_eq!( + STATE_DECODE_COUNT.load(std::sync::atomic::Ordering::SeqCst), + 0, + "list on a 200-record thread must not deserialize any record's full State" + ); + } } // ---- SQLite-backed checkpointer (feature = "sqlite") ---------------------- @@ -565,6 +802,168 @@ mod sqlite_backend { use super::checkpoint; use crate::checkpoint::{CheckpointConfig, Checkpointer, SqliteCheckpointer}; + /// Inserts a v1-shaped row directly (bypassing `insert_checkpoint_row`, + /// which — this build — only ever writes v2), with `format_version` + /// defaulting to `1` and the JSON `record` blob carrying none of the v2 + /// fields, exactly what a pre-v2 build's `INSERT` produced. Proves the + /// wire format decodes and normalizes, not just `Checkpoint::normalize` + /// agreeing with itself. + fn insert_v1_row(conn: &rusqlite::Connection, thread: &str, id: &str, parent: Option<&str>) { + let record = serde_json::json!({ + "thread_id": thread, + "checkpoint_id": id, + "run_id": null, + "parent_checkpoint_id": parent, + "namespace": [], + "state": 1, + "next_nodes": ["b"], + "completed_tasks": ["a"], + "completed_routes": [[]], + "pending_writes": [], + "interrupts": [], + "pending_activations": null, + "barrier_arrivals": [], + "metadata": { "source": "loop", "step": 1 }, + }) + .to_string(); + conn.execute( + "INSERT INTO checkpoints ( + thread_id, checkpoint_id, parent_checkpoint_id, run_id, + namespace, next_nodes, source, step, has_interrupts, record + ) VALUES (?1, ?2, ?3, NULL, '[]', '[\"b\"]', 'loop', 1, 0, ?4)", + rusqlite::params![thread, id, parent, record], + ) + .expect("insert v1 row"); + } + + #[tokio::test] + async fn v1_row_decodes_to_normalized_v2_and_resumes() { + // `insert_v1_row` needs direct SQL access, so this uses a file-backed + // database (opened once to run the schema/migration, then written to + // directly, then reopened through the checkpointer) rather than + // `in_memory`, whose connection is private to one handle. + let tmp = std::env::temp_dir().join(format!( + "tinyagents-ckpt-sqlite-v1-{}.db", + std::process::id() + )); + let _ = std::fs::remove_file(&tmp); + let cp = SqliteCheckpointer::::open(&tmp).unwrap(); + { + let raw = rusqlite::Connection::open(&tmp).unwrap(); + insert_v1_row(&raw, "v1thread", "c1", None); + } + + let loaded = cp.get("v1thread", None).await.unwrap().unwrap(); + assert_eq!(loaded.version, crate::checkpoint::CHECKPOINT_FORMAT_VERSION); + assert_eq!( + loaded + .tasks + .iter() + .map(|t| t.node.to_string()) + .collect::>(), + vec!["b".to_string()] + ); + assert_eq!( + loaded + .completed + .iter() + .map(|c| c.node.to_string()) + .collect::>(), + vec!["a".to_string()] + ); + assert!( + loaded.next_nodes.is_empty(), + "legacy fields cleared by normalize" + ); + + let history = cp.state_history("v1thread", &[], None).await.unwrap(); + assert_eq!(history.len(), 1); + assert_eq!( + history[0].checkpoint.version, + crate::checkpoint::CHECKPOINT_FORMAT_VERSION + ); + + let _ = std::fs::remove_file(&tmp); + } + + #[tokio::test] + async fn mixed_v1_and_v2_thread_lists_and_walks_state_history() { + let tmp = std::env::temp_dir().join(format!( + "tinyagents-ckpt-sqlite-mixed-{}.db", + std::process::id() + )); + let _ = std::fs::remove_file(&tmp); + let cp = SqliteCheckpointer::::open(&tmp).unwrap(); + { + let raw = rusqlite::Connection::open(&tmp).unwrap(); + insert_v1_row(&raw, "mixedthread", "c1", None); + } + cp.put(checkpoint("mixedthread", "c2", Some("c1"), 2)) + .await + .unwrap(); + + let list = cp.list("mixedthread").await.unwrap(); + assert_eq!(list.len(), 2, "both the v1 and v2 record are listed"); + + let history = cp.state_history("mixedthread", &[], None).await.unwrap(); + assert_eq!(history.len(), 2, "the walk crosses the v1/v2 boundary"); + assert!( + history + .iter() + .all(|t| t.checkpoint.version == crate::checkpoint::CHECKPOINT_FORMAT_VERSION) + ); + + let _ = std::fs::remove_file(&tmp); + } + + #[tokio::test] + async fn opening_a_pre_v2_database_migrates_the_schema_in_place() { + // A database whose `checkpoints` table predates the `format_version`/ + // `created_at` columns — the shape a build before this migration + // existed would have created. + let tmp = std::env::temp_dir().join(format!( + "tinyagents-ckpt-sqlite-migrate-{}.db", + std::process::id() + )); + let _ = std::fs::remove_file(&tmp); + { + let raw = rusqlite::Connection::open(&tmp).unwrap(); + raw.execute_batch( + "CREATE TABLE checkpoints ( + seq INTEGER PRIMARY KEY AUTOINCREMENT, + thread_id TEXT NOT NULL, + checkpoint_id TEXT NOT NULL, + parent_checkpoint_id TEXT, + run_id TEXT, + namespace TEXT NOT NULL, + next_nodes TEXT NOT NULL, + source TEXT NOT NULL, + step INTEGER NOT NULL, + has_interrupts INTEGER NOT NULL, + record TEXT NOT NULL + );", + ) + .unwrap(); + insert_v1_row(&raw, "premigrate", "c1", None); + } + + // Opening through the checkpointer must not error, and must add both + // missing columns. + let cp = SqliteCheckpointer::::open(&tmp).unwrap(); + assert!(cp.has_checkpoints_column("format_version").unwrap()); + assert!(cp.has_checkpoints_column("created_at").unwrap()); + let loaded = cp.get("premigrate", None).await.unwrap().unwrap(); + assert_eq!(loaded.version, crate::checkpoint::CHECKPOINT_FORMAT_VERSION); + + // New writes populate the migrated columns going forward. + cp.put(checkpoint("premigrate", "c2", Some("c1"), 2)) + .await + .unwrap(); + assert!(cp.get("premigrate", Some("c2")).await.unwrap().is_some()); + + let _ = std::fs::remove_file(&tmp); + } + #[tokio::test] async fn put_get_list_roundtrip_in_memory() { let cp = SqliteCheckpointer::::in_memory().unwrap(); @@ -733,4 +1132,263 @@ mod sqlite_backend { assert_eq!(records[1].state, 2); assert!(cp.get_thread("missing").await.unwrap().is_empty()); } + + // ---- C3/R4: durable per-thread execution lease ------------------------- + + #[tokio::test] + async fn a_live_lease_is_refused_to_a_different_owner() { + let cp = SqliteCheckpointer::::in_memory().unwrap(); + assert!( + cp.try_claim("t", "owner-a", std::time::Duration::from_secs(60)) + .await + .unwrap() + ); + // A different owner is refused while the lease is still live. + assert!( + !cp.try_claim("t", "owner-b", std::time::Duration::from_secs(60)) + .await + .unwrap() + ); + // The same owner re-claiming (e.g. a renew-by-reclaim) succeeds. + assert!( + cp.try_claim("t", "owner-a", std::time::Duration::from_secs(60)) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn a_stale_lease_past_its_ttl_is_reclaimable() { + let cp = SqliteCheckpointer::::in_memory().unwrap(); + // Claim with a TTL of 0 - expires immediately (simulates a dead + // owner's lease that has aged out). + assert!( + cp.try_claim("t", "dead-owner", std::time::Duration::from_millis(0)) + .await + .unwrap() + ); + // A short sleep guarantees `now` has moved past the zero-TTL expiry. + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + assert!( + cp.try_claim("t", "new-owner", std::time::Duration::from_secs(60)) + .await + .unwrap(), + "an expired lease must be reclaimable by a different owner" + ); + // The reclaim actually transferred ownership: the dead owner can no + // longer renew it. + assert!( + !cp.renew("t", "dead-owner", std::time::Duration::from_secs(60)) + .await + .unwrap() + ); + assert!( + cp.renew("t", "new-owner", std::time::Duration::from_secs(60)) + .await + .unwrap() + ); + } + + #[tokio::test] + async fn release_frees_the_lease_for_another_owner() { + let cp = SqliteCheckpointer::::in_memory().unwrap(); + assert!( + cp.try_claim("t", "owner-a", std::time::Duration::from_secs(60)) + .await + .unwrap() + ); + cp.release("t", "owner-a").await.unwrap(); + assert!( + cp.try_claim("t", "owner-b", std::time::Duration::from_secs(60)) + .await + .unwrap() + ); + } + + // ---- I8: pragmas, spawn_blocking, LIMIT-driven state_history ----------- + + /// `i32` wrapper whose [`serde::Deserialize`] impl counts every decode, so + /// tests can assert *how many* checkpoint records were actually + /// deserialized rather than just how many the call returned — the thing a + /// truncate-in-Rust `state_history` and a LIMIT-in-SQL one cannot be told + /// apart by from the returned `Vec`'s length alone. + #[derive(Clone, serde::Serialize)] + struct CountingState(i32); + + static DECODE_COUNT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + + impl<'de> serde::Deserialize<'de> for CountingState { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + let value = i32::deserialize(deserializer)?; + DECODE_COUNT.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(CountingState(value)) + } + } + + fn counting_checkpoint( + id: &str, + parent: Option<&str>, + step: usize, + ) -> crate::Checkpoint { + crate::Checkpoint::new( + CountingState(step as i32), + vec![crate::PendingActivation { + node: tinyagents_harness::ids::NodeId::from("n"), + send_arg: None, + task_id: tinyagents_harness::ids::TaskId::from(String::new()), + }], + ) + .with_thread_id("t".to_string()) + .with_checkpoint_id(id.to_string()) + .with_parent_checkpoint_id(parent.map(|s| s.to_string())) + .with_metadata(serde_json::json!({ "source": "loop", "step": step })) + } + + #[tokio::test] + async fn state_history_with_limit_decodes_only_that_many_records() { + let cp = SqliteCheckpointer::::in_memory().unwrap(); + + // A 40-checkpoint chain: if `state_history(Some(1))` decoded the whole + // namespace and truncated in Rust (the pre-fix behavior), the decode + // count below would be 40, not 1. + let mut parent: Option = None; + for step in 0..40 { + let id = format!("c{step}"); + cp.put(counting_checkpoint(&id, parent.as_deref(), step)) + .await + .unwrap(); + parent = Some(id); + } + + DECODE_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + let history = cp.state_history("t", &[], Some(1)).await.unwrap(); + assert_eq!(history.len(), 1); + assert_eq!(history[0].checkpoint.checkpoint_id, "c39"); + assert_eq!( + DECODE_COUNT.load(std::sync::atomic::Ordering::SeqCst), + 1, + "state_history(Some(1)) must decode exactly one record via a \ + LIMIT applied in SQL, not the whole namespace truncated in Rust" + ); + + // Sanity: an unlimited call still returns (and decodes) the whole + // chain, newest first. + DECODE_COUNT.store(0, std::sync::atomic::Ordering::SeqCst); + let full = cp.state_history("t", &[], None).await.unwrap(); + assert_eq!(full.len(), 40); + assert_eq!(full[0].checkpoint.checkpoint_id, "c39"); + assert_eq!(full[39].checkpoint.checkpoint_id, "c0"); + assert_eq!(DECODE_COUNT.load(std::sync::atomic::Ordering::SeqCst), 40); + } + + #[tokio::test] + async fn wal_and_synchronous_pragmas_are_set_on_open() { + // `:memory:` databases always report `journal_mode = memory` + // regardless of the pragma, so this needs a real file — WAL mode is + // stored in the database file's header. + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("checkpoints.db"); + let cp = SqliteCheckpointer::::open(&path).unwrap(); + + let journal_mode = cp.journal_mode().unwrap(); + assert_eq!(journal_mode.to_lowercase(), "wal"); + + // NORMAL == 1 (OFF = 0, FULL = 2, EXTRA = 3). + assert_eq!(cp.synchronous().unwrap(), 1); + + // The pragmas don't just read back cleanly — the checkpointer still + // works normally under them. + cp.put(checkpoint("t", "c1", None, 1)).await.unwrap(); + assert!(cp.get("t", None).await.unwrap().is_some()); + } + + #[tokio::test] + async fn put_with_writes_persists_both_in_one_call() { + use crate::checkpoint::PendingWrite; + use tinyagents_harness::ids::{NodeId, TaskId}; + + let cp = SqliteCheckpointer::::in_memory().unwrap(); + let cfg = CheckpointConfig { + thread_id: "t".to_string(), + checkpoint_id: Some("c1".to_string()), + namespace: vec![], + }; + let writes = vec![PendingWrite { + node: NodeId::from("n"), + task_id: TaskId::from("task-1"), + idx: 0, + channel: "out".to_string(), + payload: serde_json::json!("hi"), + }]; + + let id = cp + .put_with_writes(checkpoint("t", "c1", None, 1), &writes) + .await + .unwrap(); + assert_eq!(id.as_str(), "c1"); + + assert!(cp.get("t", Some("c1")).await.unwrap().is_some()); + let stored = cp.get_writes(&cfg).await.unwrap(); + assert_eq!(stored.len(), 1); + assert_eq!(stored[0].channel, "out"); + } +} + +#[test] +fn replay_memo_writes_are_distinguished_from_completion_markers() { + use tinyagents_harness::ids::TaskId; + + let marker = PendingWrite::completion_marker("n", "task-1"); + assert!(!marker.is_task_replay()); + assert!(!marker.is_durable_task()); + assert!(!marker.is_interrupt_after()); + + let memo = PendingWrite::durable_task("n", "task-1", 1, "call-api", json!({ "id": 7 })); + assert!(memo.is_task_replay()); + assert!(memo.is_durable_task()); + assert!( + !memo.is_control_plane(), + "memos are append-once data writes" + ); + assert_eq!(memo.durable_task_key(), Some("call-api")); + assert_eq!( + memo.channel, + format!("{DURABLE_TASK_CHANNEL_PREFIX}call-api") + ); + assert_eq!(memo.task_id, TaskId::from("task-1")); + + let deferred = PendingWrite::interrupt_after("n", "task-1", json!({ "update": 1, "goto": [] })); + assert!(deferred.is_task_replay()); + assert!(deferred.is_interrupt_after()); + assert!( + deferred.is_control_plane(), + "one deferred result per task, upserted" + ); + assert_eq!(deferred.idx, WRITES_IDX_INTERRUPT_AFTER); + assert_eq!(deferred.channel, INTERRUPT_AFTER_CHANNEL); + + // Merge semantics follow from the idx classes: a second memo under a + // fresh idx appends, a re-put deferred result replaces. + let mut stored = vec![marker.clone(), memo.clone(), deferred.clone()]; + let second_memo = PendingWrite::durable_task("n", "task-1", 2, "other", json!(2)); + let replaced = PendingWrite::interrupt_after("n", "task-1", json!({ "update": 9, "goto": [] })); + let changed = merge_writes(&mut stored, &[memo.clone(), second_memo, replaced.clone()]); + assert_eq!(changed, 2); + assert_eq!(stored.len(), 4); + assert_eq!( + stored + .iter() + .find(|w| w.is_interrupt_after()) + .unwrap() + .payload, + replaced.payload + ); + + // Round-trips through JSON with the reserved channel/idx intact. + let decoded: PendingWrite = + serde_json::from_value(serde_json::to_value(&deferred).unwrap()).unwrap(); + assert_eq!(decoded, deferred); } diff --git a/crates/tinyagents-graph/src/checkpoint/types.rs b/crates/tinyagents-graph/src/checkpoint/types.rs index 9760064e..d3850377 100644 --- a/crates/tinyagents-graph/src/checkpoint/types.rs +++ b/crates/tinyagents-graph/src/checkpoint/types.rs @@ -12,10 +12,43 @@ //! scopes nested subgraph checkpoints so a parent run and the child graphs it //! embeds never overwrite each other. +use std::collections::BTreeMap; use std::fmt; +use std::sync::Arc; -use crate::command::Interrupt; -use tinyagents_harness::ids::NodeId; +use crate::command::{Interrupt, RouteTarget}; +use tinyagents_harness::ids::{NodeId, TaskId}; + +/// Default value for a `TaskId` field carrying `#[serde(default = "..")]`: +/// `TaskId` is a foreign newtype (from `tinyagents_harness`), so it cannot +/// implement `Default` here (orphan rule) — this free function stands in for +/// it. An empty task id is exactly what a checkpoint written before task +/// identities existed decodes to. +fn empty_task_id() -> TaskId { + TaskId::from(String::new()) +} + +/// `#[serde(skip_serializing_if = "..")]` predicate pairing with +/// [`empty_task_id`]. +fn task_id_is_empty(id: &TaskId) -> bool { + id.as_str().is_empty() +} + +/// The current on-disk checkpoint record shape (checkpoint format v2): a +/// single `tasks`/`completed` pair replaces the four overlapping v1 +/// projections of pending work (`next_nodes`, `completed_tasks` + +/// `completed_routes`, `pending_activations`). See the module docs on +/// [`Checkpoint`] and `docs/modules/graph/checkpointing.md` for the full +/// decode story. +pub const CHECKPOINT_FORMAT_VERSION: u32 = 2; + +/// `#[serde(default = "..")]` for [`Checkpoint::version`]: a record with no +/// `version` field on disk predates the field entirely, which is exactly +/// what checkpoint format v1 (the shape before this constant existed) looked +/// like. +fn checkpoint_version_v1() -> u32 { + 1 +} /// Why a checkpoint was written. /// @@ -149,6 +182,21 @@ pub struct CheckpointTuple { /// through JSON. The in-memory path never needs it. #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] pub struct Checkpoint { + /// The on-disk record shape. `1` (the implicit shape before this field + /// existed — `#[serde(default = "checkpoint_version_v1")]`) or + /// [`CHECKPOINT_FORMAT_VERSION`] (`2`). Every writer in this crate stamps + /// `2`; a `1` is only ever seen decoding a record written by an older + /// build. See [`Checkpoint::normalize`]. + #[serde(default = "checkpoint_version_v1")] + pub version: u32, + /// Wall-clock time this checkpoint was written, in milliseconds since the + /// Unix epoch (see [`tinyagents_harness::ids::now_ms`]). + /// + /// `#[serde(default)]` (`0`) for a v1 record, which never carried a + /// timestamp at all — `0` is a visibly-unset sentinel, not a plausible + /// wall-clock value. + #[serde(default)] + pub created_at: u64, /// Checkpoint lineage key for a conversation/workflow/tenant run series. pub thread_id: String, /// This checkpoint's id within the thread. @@ -165,24 +213,35 @@ pub struct Checkpoint { pub namespace: Vec, /// Committed graph state at this boundary. pub state: State, - /// Nodes that should run when resuming from this checkpoint. - pub next_nodes: Vec, - /// Nodes that completed in the step that produced this checkpoint. - pub completed_tasks: Vec, + /// The single source of truth for what runs when this checkpoint is + /// resumed: every pending activation, preserving each one's + /// per-invocation [`Send`](crate::Send) argument and task identity. + /// + /// Checkpoint format v2 (see [`Checkpoint::version`]). Replaces the v1 + /// pair of `next_nodes` (a node-id-only projection) and + /// `pending_activations` (an `Option`-wrapped superset that was the same + /// information, just optional) with exactly one field that is never + /// ambiguous with anything else on the record. A v1 record decodes with + /// this empty; call [`Checkpoint::normalize`] (every bundled backend's + /// decode path does) to populate it from the legacy fields. + #[serde(default)] + pub tasks: Vec, + /// The single source of truth for what completed in the step that + /// produced this checkpoint, and how each task explicitly routed (if it + /// returned a `Command::goto`). + /// + /// Checkpoint format v2. Replaces the v1 pair of parallel vectors + /// `completed_tasks: Vec` and + /// `completed_routes: Vec>`, which had to stay + /// positionally aligned by convention rather than by type. A v1 record + /// decodes with this empty; [`Checkpoint::normalize`] zips the legacy + /// pair back into this shape. + #[serde(default)] + pub completed: Vec, /// Per-task partial writes preserved when a step partially completes. pub pending_writes: Vec, /// Interrupts that paused the run at this boundary. pub interrupts: Vec, - /// Pending activations to schedule on resume, preserving each pending - /// node's per-invocation [`Send`](crate::Send) argument. - /// - /// A richer superset of [`next_nodes`](Self::next_nodes) (which stays the - /// node-id projection used for listing and status). `#[serde(default)]` - /// keeps checkpoints written before this field loadable: they deserialize - /// to `None`, and resume falls back to `next_nodes` (node-only, no send - /// arg) — exactly the pre-field behavior. - #[serde(default)] - pub pending_activations: Option>, /// Barrier (waiting-edge) arrivals accumulated across supersteps, persisted /// so a join node's precondition survives an interrupt/failure + resume. /// @@ -190,8 +249,74 @@ pub struct Checkpoint { /// empty set (the pre-field behavior, where arrivals were run-local). #[serde(default)] pub barrier_arrivals: Vec, + /// Cumulative per-channel version counters as of this boundary (I5/R3): + /// name -> a monotonically-increasing count of writes to that channel. + /// + /// For a [`crate::channel::ChannelState`] graph this is + /// [`crate::channel::ChannelState::channel_versions`] verbatim. For a + /// plain whole-`State` graph (any other `State` type) it is the single + /// entry `{"state": }`, bumped once per checkpoint — + /// see `compiled::channel_bookkeeping`, the one function every + /// checkpoint-construction call site (a normal superstep boundary, + /// `update_state`, `fork_state`) uses to fill this field, so replay and + /// a manual write cannot disagree about it. + /// + /// `#[serde(default)]` for back-compat: a checkpoint written before this + /// field existed decodes with it empty. + #[serde(default)] + pub channel_versions: BTreeMap, + /// Per-node snapshot of [`Checkpoint::channel_versions`] as of the last + /// time each node ran, keyed by node id string (`NodeId` itself has no + /// `Ord` impl to key a `BTreeMap` on). Backs + /// [`crate::builder::NodeContext::changed_since_last_run`): a node + /// compares its own entry here (what it last observed) against the + /// checkpoint's live `channel_versions` (what is current) to tell + /// whether a channel changed since it last ran. + /// + /// `#[serde(default)]` for back-compat. + #[serde(default)] + pub versions_seen: BTreeMap>, + /// Per-step delta-channel write history (I5/R3): for every channel + /// registered with [`crate::channel::ChannelSet::with_delta`], the raw + /// values written to it *in the step that produced this checkpoint* + /// (not cumulative — each checkpoint carries only its own step's + /// writes, which is what keeps per-checkpoint size bounded for a + /// long-running append channel). Replayed across a thread's lineage by + /// [`crate::Checkpointer::delta_history`]. + /// + /// `#[serde(default)]` for back-compat and for every checkpoint of a + /// graph that declares no delta-tracked channels (always empty there). + #[serde(default)] + pub channel_deltas: BTreeMap>, /// Free-form metadata (source, step, etc.). pub metadata: serde_json::Value, + + // ---- Checkpoint format v1 fields (decode-only) ------------------------- + // + // Every writer in this crate leaves these at their empty default, so a + // freshly-written record serializes with none of them present + // (`skip_serializing_if`) — only [`Checkpoint::tasks`]/ + // [`Checkpoint::completed`] above carry pending/completed work going + // forward. They exist purely so a record written by a build that + // predates checkpoint format v2 still deserializes; [`Checkpoint::normalize`] + // is the single place that reads them and folds them into the v2 shape. + // Every reader elsewhere in this crate (`compiled::{resume,boundary, + // state_api,mod}`) reads `tasks`/`completed` only. + /// v1: nodes that should run when resuming from this checkpoint. Decode-only. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub next_nodes: Vec, + /// v1: nodes that completed in the step that produced this checkpoint. + /// Decode-only. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub completed_tasks: Vec, + /// v1: the explicit `Command::goto` routing for each entry of + /// [`completed_tasks`](Self::completed_tasks), positionally aligned. + /// Decode-only. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub completed_routes: Vec>, + /// v1: pending activations superset of `next_nodes`. Decode-only. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub pending_activations: Option>, } /// One pending node activation persisted in a checkpoint: the node to run on @@ -207,15 +332,70 @@ pub struct PendingActivation { pub node: NodeId, /// The per-invocation `Send` argument, when the activation was a `Send` /// packet (plain edge/goto activations carry `None`). + /// + /// `Arc`-wrapped (M2 in `docs/runtime-comparison/code-review-graph.md`) + /// so a repeated `Send` fan-out of the same node shares one allocation + /// in memory; serde's blanket `Arc` impl serializes/deserializes it + /// exactly as a bare `serde_json::Value`, so on-disk checkpoint records + /// are unaffected by this type change. #[serde(default, skip_serializing_if = "Option::is_none")] - pub send_arg: Option, + pub send_arg: Option>, /// Stable identity of this scheduled task within its superstep. /// /// Unlike `node`, this distinguishes repeated `Send` fan-out activations /// targeting the same node. Empty on checkpoints written before task - /// identities were persisted. - #[serde(default, skip_serializing_if = "String::is_empty")] - pub task_id: String, + /// identities were persisted. Serializes transparently as the underlying + /// string, so on-disk records are unaffected by the `String` -> `TaskId` + /// type change (R5). + #[serde(default = "empty_task_id", skip_serializing_if = "task_id_is_empty")] + pub task_id: TaskId, +} + +/// One task that completed in the step a checkpoint's boundary closes, +/// checkpoint format v2's replacement for the v1 +/// `completed_tasks: Vec` / `completed_routes: Vec>` +/// pair (see [`Checkpoint::completed`]). +#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct CompletedTask { + /// The task that completed, unique within the superstep that produced + /// it. Empty (`TaskId::from(String::new())`) for a task carried forward + /// from a checkpoint written before task identities existed, or derived + /// from a v1 record's `completed_tasks` (which carried no task id at + /// all). + #[serde(default = "empty_task_id", skip_serializing_if = "task_id_is_empty")] + pub task_id: TaskId, + /// The node that completed. + pub node: NodeId, + /// The explicit `Command::goto` routing this task returned, or empty when + /// it returned none (route via static/conditional edges instead). + #[serde(default)] + pub routes: Vec, +} + +impl CompletedTask { + /// Builds a completed-task record with no explicit `Command::goto` + /// routing (route via static/conditional edges). + pub fn new(task_id: impl Into, node: impl Into) -> Self { + Self { + task_id: task_id.into(), + node: node.into(), + routes: Vec::new(), + } + } + + /// Builds a completed-task record carrying an explicit `Command::goto` + /// routing. + pub fn with_routes( + task_id: impl Into, + node: impl Into, + routes: Vec, + ) -> Self { + Self { + task_id: task_id.into(), + node: node.into(), + routes, + } + } } /// The persisted arrivals recorded against one barrier (waiting-edge) join node: @@ -230,14 +410,232 @@ pub struct BarrierArrivals { } impl Checkpoint { + /// Builds a fresh checkpoint format v2 record. + /// + /// Sensible defaults for everything except `state` and `tasks`: a + /// freshly-minted [`checkpoint_id`](Self::checkpoint_id) (collision-free + /// across process restarts, matching what every executor-driven write + /// already used — see + /// [`tinyagents_harness::ids::new_checkpoint_id`]), the current + /// [`created_at`](Self::created_at), [`version`](Self::version) == + /// [`CHECKPOINT_FORMAT_VERSION`], and empty `thread_id`/`completed`/ + /// `pending_writes`/`interrupts`/`barrier_arrivals`/`namespace`, with + /// `metadata` left `null`. Chain the `with_*` setters below to fill in + /// the rest; every field is also directly `pub` for call sites that + /// prefer plain field assignment. + pub fn new(state: State, tasks: Vec) -> Self { + Self { + version: CHECKPOINT_FORMAT_VERSION, + created_at: tinyagents_harness::ids::now_ms(), + thread_id: String::new(), + checkpoint_id: tinyagents_harness::ids::new_checkpoint_id() + .as_str() + .to_string(), + run_id: None, + parent_checkpoint_id: None, + namespace: Vec::new(), + state, + tasks, + completed: Vec::new(), + pending_writes: Vec::new(), + interrupts: Vec::new(), + barrier_arrivals: Vec::new(), + channel_versions: BTreeMap::new(), + versions_seen: BTreeMap::new(), + channel_deltas: BTreeMap::new(), + metadata: serde_json::Value::Null, + next_nodes: Vec::new(), + completed_tasks: Vec::new(), + completed_routes: Vec::new(), + pending_activations: None, + } + } + + /// Alias for [`Checkpoint::new`] with no pending tasks yet — the start of + /// a fluent build, e.g. `Checkpoint::builder(state).with_tasks(pending)`. + pub fn builder(state: State) -> Self { + Self::new(state, Vec::new()) + } + + /// Sets [`Checkpoint::thread_id`]. + pub fn with_thread_id(mut self, thread_id: impl Into) -> Self { + self.thread_id = thread_id.into(); + self + } + + /// Sets [`Checkpoint::checkpoint_id`], overriding the freshly-minted + /// default from [`Checkpoint::new`]. + pub fn with_checkpoint_id(mut self, checkpoint_id: impl Into) -> Self { + self.checkpoint_id = checkpoint_id.into(); + self + } + + /// Sets [`Checkpoint::run_id`]. + pub fn with_run_id(mut self, run_id: impl Into) -> Self { + self.run_id = Some(run_id.into()); + self + } + + /// Sets [`Checkpoint::parent_checkpoint_id`]. + pub fn with_parent_checkpoint_id(mut self, parent: Option) -> Self { + self.parent_checkpoint_id = parent; + self + } + + /// Sets [`Checkpoint::namespace`]. + pub fn with_namespace(mut self, namespace: Vec) -> Self { + self.namespace = namespace; + self + } + + /// Sets [`Checkpoint::tasks`]. + pub fn with_tasks(mut self, tasks: Vec) -> Self { + self.tasks = tasks; + self + } + + /// Sets [`Checkpoint::completed`]. + pub fn with_completed(mut self, completed: Vec) -> Self { + self.completed = completed; + self + } + + /// Sets [`Checkpoint::pending_writes`]. + pub fn with_pending_writes(mut self, writes: Vec) -> Self { + self.pending_writes = writes; + self + } + + /// Sets [`Checkpoint::interrupts`]. + pub fn with_interrupts(mut self, interrupts: Vec) -> Self { + self.interrupts = interrupts; + self + } + + /// Sets [`Checkpoint::barrier_arrivals`]. + pub fn with_barrier_arrivals(mut self, barrier_arrivals: Vec) -> Self { + self.barrier_arrivals = barrier_arrivals; + self + } + + /// Sets [`Checkpoint::metadata`]. + pub fn with_metadata(mut self, metadata: serde_json::Value) -> Self { + self.metadata = metadata; + self + } + + /// Sets [`Checkpoint::channel_versions`]. + pub fn with_channel_versions(mut self, channel_versions: BTreeMap) -> Self { + self.channel_versions = channel_versions; + self + } + + /// Sets [`Checkpoint::versions_seen`]. + pub fn with_versions_seen( + mut self, + versions_seen: BTreeMap>, + ) -> Self { + self.versions_seen = versions_seen; + self + } + + /// Sets [`Checkpoint::channel_deltas`]. + pub fn with_channel_deltas( + mut self, + channel_deltas: BTreeMap>, + ) -> Self { + self.channel_deltas = channel_deltas; + self + } + + /// The effective pending-task set: [`Checkpoint::tasks`] directly on a + /// v2 record (`version >= 2`), or derived from the v1 fields + /// (preferring `pending_activations`, falling back to `next_nodes`) on a + /// v1 record. Non-mutating — shared by [`Checkpoint::normalize`] (which + /// writes the result back) and [`Checkpoint::to_metadata`] (which only + /// needs to read it). + fn effective_tasks(&self) -> Vec { + if self.version >= CHECKPOINT_FORMAT_VERSION { + return self.tasks.clone(); + } + match &self.pending_activations { + Some(pending) if !pending.is_empty() => pending.clone(), + _ => self + .next_nodes + .iter() + .cloned() + .map(|node| PendingActivation { + node, + send_arg: None, + task_id: empty_task_id(), + }) + .collect(), + } + } + + /// The effective completed-task set: [`Checkpoint::completed`] directly + /// on a v2 record, or zipped from the v1 `completed_tasks`/ + /// `completed_routes` pair (padding a shorter/missing `completed_routes` + /// with empty routing — the pre-`completed_routes` behavior) on a v1 + /// record. Non-mutating, mirroring [`Checkpoint::effective_tasks`]. + fn effective_completed(&self) -> Vec { + if self.version >= CHECKPOINT_FORMAT_VERSION { + return self.completed.clone(); + } + self.completed_tasks + .iter() + .cloned() + .zip( + self.completed_routes + .iter() + .cloned() + .chain(std::iter::repeat(Vec::new())), + ) + .map(|(node, routes)| CompletedTask { + task_id: empty_task_id(), + node, + routes, + }) + .collect() + } + + /// Folds a checkpoint format v1 record into the current (v2) shape, + /// in place: populates [`Checkpoint::tasks`]/[`Checkpoint::completed`] + /// from whichever legacy fields the record carries (see + /// [`Checkpoint::effective_tasks`]/[`Checkpoint::effective_completed`]), + /// clears the legacy fields (so a subsequent `put` of the same value + /// re-serializes as clean v2), and stamps [`Checkpoint::version`] to + /// [`CHECKPOINT_FORMAT_VERSION`]. + /// + /// A no-op on an already-v2 record. Every bundled [`Checkpointer`] + /// backend calls this on every decode path (`get`/`get_scoped`/`list`/ + /// `state_history`/`get_thread`), so callers outside this module never + /// observe a v1 record — see `docs/modules/graph/checkpointing.md`. + pub fn normalize(&mut self) { + if self.version >= CHECKPOINT_FORMAT_VERSION { + return; + } + self.tasks = self.effective_tasks(); + self.completed = self.effective_completed(); + self.next_nodes = Vec::new(); + self.completed_tasks = Vec::new(); + self.completed_routes = Vec::new(); + self.pending_activations = None; + self.version = CHECKPOINT_FORMAT_VERSION; + } + /// Builds the lightweight [`CheckpointMetadata`] summary for this checkpoint. /// /// The single source of truth for projecting a stored checkpoint onto its /// listing record: it parses the `source`/`step` out of the free-form - /// `metadata` (falling back to [`CheckpointSource::Loop`]/`0`) and copies the - /// lineage fields. Both `Checkpointer::list` and the state-inspection API + /// `metadata` (falling back to [`CheckpointSource::Loop`]/`0`), projects + /// [`Checkpoint::effective_tasks`] onto its node ids for + /// [`CheckpointMetadata::next_nodes`], and copies the lineage fields. Both + /// `Checkpointer::list` and the state-inspection API /// (`get_state`/`get_state_history`) use it so a snapshot's metadata always - /// matches what listing reports. + /// matches what listing reports. Correct on an un-normalized v1 record too + /// (it never mutates `self`), which is what lets a header-only listing + /// path (no full-record decode) project it without first normalizing. pub fn to_metadata(&self) -> CheckpointMetadata { let source = self .metadata @@ -250,13 +648,14 @@ impl Checkpoint { .get("step") .and_then(|v| v.as_u64()) .unwrap_or(0) as usize; + let next_nodes = self.effective_tasks().into_iter().map(|t| t.node).collect(); CheckpointMetadata { thread_id: self.thread_id.clone(), checkpoint_id: self.checkpoint_id.clone(), run_id: self.run_id.clone(), parent_checkpoint_id: self.parent_checkpoint_id.clone(), namespace: self.namespace.clone(), - next_nodes: self.next_nodes.clone(), + next_nodes, has_interrupts: !self.interrupts.is_empty(), source, step, @@ -277,6 +676,22 @@ pub const WRITES_IDX_ERROR: i64 = -2; /// The `idx` reserved for a task's **interrupt** control-plane write. pub const WRITES_IDX_INTERRUPT: i64 = -3; +/// The `idx` reserved for a task's **deferred result** control-plane write: +/// the `Update`/`Command::goto` a node produced but that an +/// [`interrupt_after`](crate::GraphBuilder::interrupt_after) pause held back +/// from committed state. At most one per task (a task completes once), and +/// re-put on a repeated pause replaces it — the control-plane upsert rule. +/// Replayed by the executor on resume instead of re-running the handler. +pub const WRITES_IDX_INTERRUPT_AFTER: i64 = -4; + +/// Channel prefix for a [`crate::NodeContext::durable_task`] memo write: +/// the full channel is this prefix followed by the caller's `key`. +pub const DURABLE_TASK_CHANNEL_PREFIX: &str = "__durable_task__:"; + +/// Channel of a task's deferred `interrupt_after` result write (see +/// [`WRITES_IDX_INTERRUPT_AFTER`]). +pub const INTERRUPT_AFTER_CHANNEL: &str = "__interrupt_after__"; + /// A partial write produced by a completed task, preserved across reruns. /// /// # Why writes are recorded separately from the checkpoint @@ -323,8 +738,8 @@ pub struct PendingWrite { /// A plain node id is not enough on its own: a fan-out step runs the same /// node several times with different [`Send`](crate::Send) args, and /// each of those is a separately resumable task. - #[serde(default)] - pub task_id: String, + #[serde(default = "empty_task_id")] + pub task_id: TaskId, /// Position of this write within its task's emission order, or one of the /// `WRITES_IDX_*` constants for a control-plane write. #[serde(default)] @@ -351,7 +766,7 @@ impl PendingWrite { /// Builds an ordinary data write for `task_id` at position `idx`. pub fn data( node: impl Into, - task_id: impl Into, + task_id: impl Into, idx: i64, channel: impl Into, payload: serde_json::Value, @@ -367,7 +782,7 @@ impl PendingWrite { /// Builds a completion marker: a data write at index `0` whose payload is /// `null`, recording only that `task_id` ran to completion. - pub fn completion_marker(node: impl Into, task_id: impl Into) -> Self { + pub fn completion_marker(node: impl Into, task_id: impl Into) -> Self { let node = node.into(); let channel = node.as_str().to_string(); Self { @@ -379,12 +794,83 @@ impl PendingWrite { } } + /// Builds a [`crate::NodeContext::durable_task`] memo write: the + /// serialized output of the task's durable sub-step `key`, stored as an + /// ordinary data write (`idx >= 1`, append-once) on the + /// [`DURABLE_TASK_CHANNEL_PREFIX`]`key` channel. `idx` must be unique + /// among the task's writes (the executor allocates it past every write + /// the task already holds); `0` is reserved for the completion marker. + pub fn durable_task( + node: impl Into, + task_id: impl Into, + idx: i64, + key: &str, + payload: serde_json::Value, + ) -> Self { + debug_assert!(idx >= 1, "durable-task writes use idx >= 1"); + Self { + node: node.into(), + task_id: task_id.into(), + idx, + channel: format!("{DURABLE_TASK_CHANNEL_PREFIX}{key}"), + payload, + } + } + + /// Builds a task's deferred `interrupt_after` result write (see + /// [`WRITES_IDX_INTERRUPT_AFTER`]): `payload` is + /// `{"update": , "goto": [...]}`. + pub fn interrupt_after( + node: impl Into, + task_id: impl Into, + payload: serde_json::Value, + ) -> Self { + Self { + node: node.into(), + task_id: task_id.into(), + idx: WRITES_IDX_INTERRUPT_AFTER, + channel: INTERRUPT_AFTER_CHANNEL.to_string(), + payload, + } + } + /// Whether this is a control-plane write (`idx < 0`), which upserts rather /// than appends. See the type docs. pub fn is_control_plane(&self) -> bool { self.idx < 0 } + /// Whether this is a [`crate::NodeContext::durable_task`] memo write + /// (see [`Self::durable_task`]). + pub fn is_durable_task(&self) -> bool { + self.channel.starts_with(DURABLE_TASK_CHANNEL_PREFIX) + } + + /// The caller's `key` of a durable-task memo write, or `None` for any + /// other write. + pub fn durable_task_key(&self) -> Option<&str> { + self.channel.strip_prefix(DURABLE_TASK_CHANNEL_PREFIX) + } + + /// Whether this is a task's deferred `interrupt_after` result write (see + /// [`Self::interrupt_after`]). + pub fn is_interrupt_after(&self) -> bool { + self.idx == WRITES_IDX_INTERRUPT_AFTER && self.channel == INTERRUPT_AFTER_CHANNEL + } + + /// Whether this write is a per-task *replay memo* — a durable-task memo + /// or a deferred `interrupt_after` result — that a re-run of the same + /// (still pending) task consumes, as opposed to a completion marker or + /// any other write that records the task as already done. + /// + /// Resume keys "which pending tasks already ran" off the writes that are + /// *not* replay memos: a replay memo belongs to a task that has *not* + /// completed yet (that is the whole point of memoising it), so counting + /// it as a completion marker would wrongly skip the task. + pub fn is_task_replay(&self) -> bool { + self.is_durable_task() || self.is_interrupt_after() + } + /// The `(task_id, idx)` identity pair this write is deduplicated on within /// a checkpoint. pub fn identity(&self) -> (&str, i64) { diff --git a/crates/tinyagents-graph/src/command/mod.rs b/crates/tinyagents-graph/src/command/mod.rs index a3db9abc..7702ed5d 100644 --- a/crates/tinyagents-graph/src/command/mod.rs +++ b/crates/tinyagents-graph/src/command/mod.rs @@ -14,11 +14,7 @@ mod types; pub use types::{Command, Interrupt, NodeResult, RouteTarget, Send}; -use std::sync::atomic::{AtomicU64, Ordering}; - -use tinyagents_harness::ids::NodeId; - -static INTERRUPT_SEQ: AtomicU64 = AtomicU64::new(0); +use tinyagents_harness::ids::{NodeId, TaskId}; impl Command { /// Creates an empty command (no update, no routing, no resume). @@ -27,6 +23,7 @@ impl Command { update: None, goto: Vec::new(), resume: None, + resume_by_task: std::collections::HashMap::new(), } } @@ -39,6 +36,7 @@ impl Command { .map(|t| RouteTarget::Node(t.into())) .collect(), resume: None, + resume_by_task: std::collections::HashMap::new(), } } @@ -50,6 +48,7 @@ impl Command { update: None, goto: sends.into_iter().map(RouteTarget::Send).collect(), resume: None, + resume_by_task: std::collections::HashMap::new(), } } @@ -59,6 +58,7 @@ impl Command { update: Some(update), goto: Vec::new(), resume: None, + resume_by_task: std::collections::HashMap::new(), } } @@ -68,6 +68,21 @@ impl Command { update: None, goto: Vec::new(), resume: Some(value), + resume_by_task: std::collections::HashMap::new(), + } + } + + /// Creates a resume command carrying a distinct value per interrupted + /// task (I1): a `Send` fan-out of the same node produces several + /// concurrently-interrupted tasks, and this is how a caller delivers each + /// its own resume value in one call, keyed by + /// [`crate::builder::NodeContext::task_id`]. + pub fn resume_tasks(values: impl IntoIterator) -> Self { + Self { + update: None, + goto: Vec::new(), + resume: None, + resume_by_task: values.into_iter().collect(), } } @@ -105,13 +120,28 @@ impl Default for Command { impl Interrupt { /// Creates an interrupt with an auto-generated unique id. + /// + /// I7 (`docs/runtime-comparison/code-review-graph.md`): built from + /// [`tinyagents_harness::ids::process_nonce`] + + /// [`tinyagents_harness::ids::next_seq`] — the same restart-safe scheme + /// [`tinyagents_harness::ids::new_checkpoint_id`] uses — rather than a + /// bare process-local counter. A bare counter restarts at `0` in every + /// new process, so two pauses minted in different process lifetimes + /// could collide on `(node, seq)` and conflate two distinct interrupts + /// in `GraphRunStatus::pending_interrupts` or a UI keyed on interrupt id. pub fn new(node: impl Into, payload: serde_json::Value) -> Self { let node = node.into(); - let seq = INTERRUPT_SEQ.fetch_add(1, Ordering::Relaxed); + let id = format!( + "interrupt-{node}-{}-{}", + tinyagents_harness::ids::process_nonce(), + tinyagents_harness::ids::next_seq() + ); Self { - id: format!("interrupt-{node}-{seq}"), + id, node, payload, + task_id: None, + response_schema: None, } } @@ -125,8 +155,28 @@ impl Interrupt { id: id.into(), node: node.into(), payload, + task_id: None, + response_schema: None, } } + + /// Attaches the JSON schema a resume value for this interrupt must + /// satisfy (see [`Interrupt::response_schema`]). + pub fn with_response_schema(mut self, schema: serde_json::Value) -> Self { + self.response_schema = Some(schema); + self + } + + /// Returns this interrupt with its scheduled task id set (R5/I1). + /// + /// The interrupt boundary calls this on the emitted interrupt before + /// persisting/returning it, so a `Send` fan-out of the same node (or a + /// re-emitted subgraph interrupt) is resumable by its own task rather + /// than sharing the node's identity with its siblings. + pub fn with_task_id(mut self, task_id: tinyagents_harness::ids::TaskId) -> Self { + self.task_id = Some(task_id); + self + } } #[cfg(test)] diff --git a/crates/tinyagents-graph/src/command/test.rs b/crates/tinyagents-graph/src/command/test.rs index 52019b41..1347a9d0 100644 --- a/crates/tinyagents-graph/src/command/test.rs +++ b/crates/tinyagents-graph/src/command/test.rs @@ -47,3 +47,48 @@ fn interrupt_ids_are_unique() { let fixed = Interrupt::with_id("fixed", "n", json!(null)); assert_eq!(fixed.id, "fixed"); } + +/// I7 regression: interrupt ids are minted with the same restart-safe +/// process-nonce scheme `tinyagents_harness::ids::new_checkpoint_id` uses, +/// not a bare process-local counter that restarts at `0` every process (see +/// `docs/runtime-comparison/code-review-graph.md`). Asserts the id embeds +/// the process nonce and that a large batch of mints never collides. +#[test] +fn interrupt_ids_embed_the_process_nonce_and_never_collide() { + let nonce = tinyagents_harness::ids::process_nonce().to_string(); + let mut seen = std::collections::HashSet::new(); + for _ in 0..1000 { + let interrupt = Interrupt::new("n", json!(null)); + assert!( + interrupt.id.contains(&nonce), + "interrupt id `{}` must embed the process nonce `{nonce}`", + interrupt.id + ); + assert!( + seen.insert(interrupt.id.clone()), + "interrupt id `{}` was minted twice", + interrupt.id + ); + } +} + +#[test] +fn interrupt_response_schema_round_trips_and_defaults_to_none() { + let schema = json!({ "type": "object", "required": ["approved"] }); + let interrupt = + Interrupt::new("approve", json!({ "ask": "ok?" })).with_response_schema(schema.clone()); + assert_eq!(interrupt.response_schema, Some(schema.clone())); + let encoded = serde_json::to_value(&interrupt).unwrap(); + assert_eq!(encoded["response_schema"], schema); + let decoded: Interrupt = serde_json::from_value(encoded).unwrap(); + assert_eq!(decoded, interrupt); + + // Legacy checkpoint JSON without the field decodes with `None`, and a + // schema-less interrupt omits the key entirely on the wire. + let bare = Interrupt::new("approve", json!(null)); + let encoded = serde_json::to_value(&bare).unwrap(); + assert!(encoded.get("response_schema").is_none()); + let legacy: Interrupt = + serde_json::from_value(json!({ "id": "i", "node": "approve", "payload": {} })).unwrap(); + assert_eq!(legacy.response_schema, None); +} diff --git a/crates/tinyagents-graph/src/command/types.rs b/crates/tinyagents-graph/src/command/types.rs index 390202a4..ad2fb28e 100644 --- a/crates/tinyagents-graph/src/command/types.rs +++ b/crates/tinyagents-graph/src/command/types.rs @@ -10,7 +10,7 @@ //! - [`NodeResult::Interrupt`]: an [`Interrupt`] that pauses the run for //! human-in-the-loop input. -use tinyagents_harness::ids::NodeId; +use tinyagents_harness::ids::{NodeId, TaskId}; /// The outcome of running a durable graph node. #[derive(Clone, Debug)] @@ -32,7 +32,7 @@ pub enum NodeResult { /// pointing at the *same* target node — and each scheduled invocation receives /// its own `arg`. Distinct from a plain `goto`, which simply activates a node /// against the shared state with no per-activation input. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] pub struct Send { /// The node to schedule. pub node: NodeId, @@ -53,7 +53,13 @@ impl Send { /// A single routing target produced by a [`Command`]: either a plain node /// activation ([`RouteTarget::Node`]) or a [`Send`] packet carrying /// per-invocation input ([`RouteTarget::Send`]). -#[derive(Clone, Debug)] +/// +/// Serializable (R1 in `docs/runtime-comparison/code-review-graph.md`): a +/// completed sibling's explicit `Command::goto` is persisted alongside +/// `Checkpoint::completed_tasks` (see [`crate::Checkpoint::completed_routes`]) +/// so it survives a resume instead of being re-resolved via +/// static/conditional edges only. +#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] pub enum RouteTarget { /// Activate the node against the shared committed state. Node(NodeId), @@ -95,7 +101,17 @@ pub struct Command { /// plain node activation or a [`Send`] packet (see [`RouteTarget`]). pub goto: Vec, /// Resume value for an interrupted node (used by `CompiledGraph::resume`). + /// + /// Applies to every interrupted task when `resume_by_task` is empty. When + /// a `Send` fan-out of the same node produced several concurrent + /// interrupted tasks (I1), prefer `resume_by_task` so each gets its own + /// value; this field alone cannot distinguish them. pub resume: Option, + /// Per-task resume values (R5/I1), keyed by the interrupted task's + /// [`TaskId`] (see [`crate::builder::NodeContext::task_id`]). Consulted + /// before `resume`: a task named here gets its own value; every other + /// pending task falls back to `resume` (if set). + pub resume_by_task: std::collections::HashMap, } /// A human-in-the-loop pause point. @@ -111,4 +127,29 @@ pub struct Interrupt { pub node: NodeId, /// Arbitrary payload presented to the human/approver. pub payload: serde_json::Value, + /// The scheduled task this interrupt paused, when known (R5/I1). + /// + /// Stamped by the interrupt boundary from the pausing branch's + /// [`crate::compiled` activation task id — distinct fan-out activations + /// of the same node (a `Send` `[node_id, task_id]`-scoped subgraph, for + /// example) each get their own interrupt/resume identity instead of + /// sharing the node's. `None` for a hand-built interrupt or one recorded + /// before task identity was tracked; `#[serde(default)]` keeps legacy + /// checkpoint JSON without this field decoding. + #[serde(default)] + pub task_id: Option, + /// Optional JSON schema the resume value answering this interrupt must + /// satisfy (a subset: `type`, `required`, `properties`, + /// `additionalProperties`, `items`, `enum` — see + /// [`tinyagents_harness::tool::validate_against_schema`]). + /// + /// Enforced *fail-closed* by `CompiledGraph::resume`/`resume_from`: + /// the value a `Command::resume`/`Command::resume_tasks` would deliver + /// to this interrupt's task is validated before the resumed run starts, + /// and a mismatch returns [`crate::TinyAgentsError::Validation`] with + /// the thread's checkpoint left untouched. `None` (the default, and what + /// legacy checkpoint JSON without this field decodes to) accepts any + /// value. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub response_schema: Option, } diff --git a/crates/tinyagents-graph/src/compiled/boundary.rs b/crates/tinyagents-graph/src/compiled/boundary.rs new file mode 100644 index 00000000..b7ebd719 --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/boundary.rs @@ -0,0 +1,1086 @@ +//! The step boundary: applying reducer updates, routing a completed step's +//! active set into the next superstep, persisting checkpoints, and the +//! failure/interrupt boundaries that pause or abort a run. +//! +//! Routing itself (`route_completed`, static/conditional/`Send` resolution, +//! barrier gating) stays in `routing.rs`; this module is what calls it at +//! each of the three boundary shapes a superstep can end at — the normal +//! boundary ([`CompiledGraph::advance`]), a node-handler failure +//! ([`CompiledGraph::handle_failure_boundary`]), and an interrupt +//! ([`CompiledGraph::handle_interrupt_boundary`]) — plus the run-abort +//! bookkeeping ([`CompiledGraph::fail_run`], [`CompiledGraph::fail_and_return`]) +//! shared by every early-exit path in `execute_run`. + +use super::*; + +use crate::compiled::run_ctx::RunCtx; + +/// The channel-bookkeeping trio a checkpoint persists (I5/R3): current +/// per-channel versions, this boundary's delta-channel writes, and the +/// per-node `versions_seen` snapshot. See +/// [`CompiledGraph::channel_checkpoint_fields`]. +type ChannelCheckpointFields = ( + BTreeMap, + BTreeMap>, + BTreeMap>, +); + +/// The step data a boundary persist needs beyond the (possibly narrowed) +/// pending/completed activation slices: the committed state snapshot and +/// this step's child-run metadata. Bundled so the persist helpers below stay +/// under the arity that would otherwise need `#[allow(too_many_arguments)]`. +pub(super) struct BoundaryCheckpoint<'a, State> { + pub(super) state: &'a State, + pub(super) pending: &'a [Activation], + /// The step's completed tasks, each carrying its own explicit + /// `Command::goto` routing (R1: see + /// [`crate::checkpoint::CompletedTask`]). + pub(super) completed: Vec, + pub(super) child_runs: &'a serde_json::Value, + /// Per-task replay memos of this step's *stalled* branches — their + /// [`NodeContext::durable_task`] writes and any deferred + /// `interrupt_after` result — persisted alongside the completion + /// markers so a resumed re-run of those tasks can consult them. Empty at + /// a normal boundary (nothing stalled, so nothing to replay). + pub(super) task_writes: Vec, +} + +/// The transient data one superstep's boundary handling needs: the step's +/// active set, its fold outcome (`completed`/`stalled`, per [`StepRun`]), +/// its folded routing (`goto_map`), the step's child-run metadata, and the +/// step number. Shared by the normal, failure, and interrupt boundaries so +/// none of them re-take these as separate parameters. +pub(super) struct StepBoundary<'a> { + pub(super) active: &'a [Activation], + /// Every branch of this step that completed, original-index-paired (see + /// [`crate::compiled::step::StepRun::completed`]) — used both to build + /// the persisted `completed_tasks` and, at the normal boundary, as the + /// routing input. + pub(super) completed: &'a [(usize, Activation)], + /// Every branch of this step that errored or interrupted — the + /// `pending` set at a failure/interrupt boundary. + pub(super) stalled: &'a [(usize, Activation)], + pub(super) goto_map: &'a HashMap>, + pub(super) child_runs_meta: &'a serde_json::Value, + /// Replay memos of the stalled branches (see + /// [`crate::compiled::step::StepRun::task_writes`]). + pub(super) task_writes: &'a [crate::checkpoint::PendingWrite], + pub(super) step: usize, +} + +impl CompiledGraph +where + State: Clone + Send + Sync + 'static, + Update: Send + 'static, +{ + /// Applies each collected update through the reducer, in order, at the + /// step boundary. A reducer error here must still fail the run (not just + /// unwind leaving it `Running`) — surfaced to the caller as `Err`. + pub(super) fn apply_updates(&self, mut state: State, updates: Vec) -> Result { + for update in updates { + state = self.reducer.apply(state, update)?; + } + Ok(state) + } + + /// The channel-bookkeeping trio (`channel_versions`, `channel_deltas`, + /// `versions_seen`) to embed in a checkpoint built at this boundary — + /// I5/R3, via the single [`crate::channel::channel_bookkeeping`] + /// dispatch point shared with `compiled::state_api`. `state` is this + /// boundary's freshly-committed state (already folded through + /// [`Self::apply_updates`]); `ctx.steps` is the fallback per-checkpoint + /// version for a plain whole-state graph. + fn channel_checkpoint_fields( + &self, + ctx: &RunCtx<'_, State, Update>, + state: &State, + ) -> ChannelCheckpointFields { + let (channel_versions, channel_deltas) = + crate::channel::channel_bookkeeping(state, ctx.steps as u64); + let versions_seen = ctx.versions_seen.clone().into_iter().collect(); + (channel_versions, channel_deltas, versions_seen) + } + + /// Real `defer` scheduling (A1/D2): holds back every activation in + /// `next` whose effective [`NodePolicy::defer`] is set, so long as + /// `next` also contains at least one non-deferred activation — + /// accumulating the held-back set in `ctx.deferred_pending` across + /// however many supersteps that takes. The *first* time this would + /// otherwise route to an empty frontier (nothing non-deferred left + /// anywhere), every pending deferred activation is released at once. + /// + /// This makes a deferred node behave as a "run once everything else is + /// done" synthesis/join, without needing an explicit barrier naming + /// every other node in the graph. [`crate::GraphBuilder::mark_deferred`] + /// is a thin alias over the same per-node [`NodePolicy::defer`] flag + /// this reads via [`Self::effective_policy`]. + pub(super) fn apply_defer( + &self, + ctx: &mut RunCtx<'_, State, Update>, + next: Vec, + ) -> Vec { + if next.is_empty() && ctx.deferred_pending.is_empty() { + return next; + } + let (deferred, immediate): (Vec, Vec) = next + .into_iter() + .partition(|activation| self.effective_policy(&activation.node).defer); + ctx.deferred_pending.extend(deferred); + if !immediate.is_empty() { + return immediate; + } + std::mem::take(&mut ctx.deferred_pending) + } + + /// The normal (non-interrupt/non-failure) step boundary: routes the + /// completed active set into the next superstep's activations and + /// persists a boundary checkpoint per the configured + /// [`DurabilityMode`], updating `ctx.last_checkpoint`/`parent_checkpoint` + /// when one is written. Returns the next active set. + /// + /// When this run was resumed from a mid-step checkpoint (an + /// interrupt/failure boundary whose completed siblings were never + /// routed — see [`Self::handle_interrupt_boundary`] / + /// [`Self::handle_failure_boundary`]), `ctx.carried_completed` carries + /// those siblings' node ids forward. The *first* `advance` call of the + /// resumed run consumes it (`take`) and routes it together with this + /// step's own `sb.completed`, so every branch of the original step is + /// routed in one pass against one committed state — matching what an + /// uninterrupted run would have done (the C2 fix). Those carried + /// branches have no persisted `goto_map` entry (a `Command`'s explicit + /// `goto` is not durable across the boundary), so they route via + /// static/conditional edges only; see the module and `RunCtx` docs. + pub(super) async fn advance( + &self, + ctx: &mut RunCtx<'_, State, Update>, + sb: StepBoundary<'_>, + state: &State, + ) -> Result> { + // Select the next active set from commands or static/conditional + // edges, evaluated against the freshly-committed state. Barrier + // arrivals accumulate into `ctx.barrier_arrivals` (persisted below). + let carried = ctx.carried_completed.take(); + let completed_tasks: Vec; + let next = match &carried { + Some(carried_completions) => { + // Reserve an index range that cannot collide with `sb`'s own + // (0-based) active-set indices, so `goto_map.get(&index)` + // correctly misses for every carried entry instead of + // aliasing onto this step's own routing. + let offset = sb.active.len().max(sb.completed.len()) + 1; + let mut pairs: Vec<(usize, Activation)> = carried_completions + .iter() + .enumerate() + .map(|(i, (node, _))| (offset + i, Activation::node(node.clone()))) + .collect(); + pairs.extend(sb.completed.iter().cloned()); + // Merge in each carried completion's persisted `goto` (R1): + // without this, a completed sibling's explicit + // `Command::goto` is lost across the boundary and it + // re-resolves via static/conditional edges only. + let mut merged_goto_map = sb.goto_map.clone(); + for (i, (_, goto)) in carried_completions.iter().enumerate() { + if !goto.is_empty() { + merged_goto_map.insert(offset + i, goto.clone()); + } + } + let next = self.route_completed( + &ctx.run_id, + &pairs, + &merged_goto_map, + state, + &mut ctx.barrier_arrivals, + )?; + completed_tasks = pairs.into_iter().map(|(_, a)| a).collect(); + next + } + None => { + completed_tasks = sb.completed.iter().map(|(_, a)| a.clone()).collect(); + self.route_completed( + &ctx.run_id, + sb.completed, + sb.goto_map, + state, + &mut ctx.barrier_arrivals, + )? + } + }; + let next = self.apply_defer(ctx, next); + + // Persist a boundary checkpoint. Under `Exit` durability only the + // terminal boundary (the step that empties the active set) is + // written; `Sync`/`Async` persist every boundary. `Async` hands + // non-terminal writes to background tasks instead of awaiting them + // inline. + let persist_now = match self.durability { + DurabilityMode::Exit => next.is_empty(), + DurabilityMode::Sync | DurabilityMode::Async => true, + }; + // Async durability: surface any background write failure recorded + // since the previous boundary. The run fails at the first + // durability boundary that observes the loss rather than silently + // continuing with a hole in its lineage. + if let Some(err) = ctx.async_writes.take_failure().await { + return Err(err); + } + let terminal = next.is_empty(); + let checkpoint_id = if persist_now { + // Fully routed at this normal boundary, so nothing is left to + // carry forward: every completed task's routing is empty. + let completed = completed_tasks + .iter() + .map(|a| crate::checkpoint::CompletedTask::new(a.task_id.clone(), a.node.clone())) + .collect(); + let boundary = BoundaryCheckpoint { + state, + pending: &next, + completed, + child_runs: sb.child_runs_meta, + task_writes: Vec::new(), + }; + if matches!(self.durability, DurabilityMode::Async) && !terminal { + self.persist_checkpoint_nonblocking(ctx, boundary, sb.step) + .await? + } else { + // Terminal boundary: drain every in-flight background write + // first (the "final await at run end"), so a lost Async + // checkpoint fails the run instead of being swallowed. The + // final checkpoint itself is then written synchronously in + // every mode. + if terminal { + ctx.async_writes.drain().await?; + } + self.persist_checkpoint(ctx, boundary, sb.step, Vec::new(), &[]) + .await? + } + } else { + None + }; + if let Some(id) = &checkpoint_id { + ctx.last_checkpoint = Some(id.clone()); + ctx.parent_checkpoint = Some(id.to_string()); + } + + ctx.emit(GraphEvent::StepCompleted { step: sb.step }); + Ok(next) + } + + /// The failure boundary: a node-handler failure that survived the + /// node-retry policy. + /// + /// The updates of every branch that completed this step — regardless of + /// its index relative to the failed one — are already folded into + /// `state` (by [`Self::apply_updates`] before this is called, from + /// [`crate::compiled::step::StepRun::updates`]). This boundary does + /// *not* route those completed branches yet (see [`Self::advance`]'s + /// `carried_completed` doc): routing them now, before the failed/pending + /// branches are known, would let their successors observe a state that + /// omits whatever those pending branches eventually write — the exact + /// same-superstep ordering bug C2 describes for the interrupt boundary. + /// Instead `pending` is exactly `sb.stalled` (the failed node plus any + /// other branch that also errored/interrupted this step — the + /// not-yet-run set, not `active[failed_index..]`), and the completed + /// branches' node ids are stamped into the checkpoint's + /// `completed_tasks` (merged with any already-carried-forward ones from + /// an earlier resume of this same logical step) so a resuming + /// `retry`/`resume` can route the whole step together once the pending + /// branches finish. Persists a resumable failure-boundary checkpoint, + /// records a `Failed` status carrying the error and that checkpoint, and + /// returns the error. Without a checkpointer/thread the checkpoint is a + /// no-op and the run aborts exactly as before. + pub(super) async fn handle_failure_boundary( + &self, + ctx: &mut RunCtx<'_, State, Update>, + sb: StepBoundary<'_>, + state: &State, + fail: StepFailure, + ) -> Result> { + ctx.disarm_drop_guard(); + let StepFailure { + failed_index, + error, + } = fail; + let failed_node = sb.active[failed_index].node.clone(); + let pending: Vec = sb.stalled.iter().map(|(_, a)| a.clone()).collect(); + let completed = self.merged_completed(ctx, sb.completed, sb.goto_map); + // Settle any in-flight Async background writes before the + // failure-boundary persist so earlier boundaries are durable when + // the run aborts. Like the persist error below, a background write + // error must not replace the original node error, so it is + // intentionally dropped here. + let _ = ctx.async_writes.drain().await; + // A failure-boundary persist error must not replace the original + // node error: keep reporting the node error and just drop the + // resumable checkpoint reference. + let checkpoint_id = self + .persist_failure_checkpoint( + ctx, + BoundaryCheckpoint { + state, + pending: &pending, + completed, + child_runs: sb.child_runs_meta, + task_writes: sb.task_writes.to_vec(), + }, + sb.step, + &failed_node, + &error, + ) + .await + .unwrap_or(None); + self.fail_run( + &ctx.run_id, + &ctx.thread_id, + ctx.started_at, + sb.step, + &error, + checkpoint_id, + ) + .await; + Err(error) + } + + /// The interrupt boundary: persists a checkpoint whose pending + /// activations are the successors of the branches that completed before + /// the interrupt (their routing must survive) followed by the + /// not-yet-completed members of this step (interrupted node first). + /// Each pending branch keeps its `Send` arg; accumulated barrier + /// arrivals are persisted too. Returns control to the caller. + /// + /// `interrupted` is every branch of this step whose result was an + /// interrupt (I1), in ascending active-set index order — a `Send` + /// fan-out of one subgraph node interrupting on every one of its + /// concurrent activations, for example, surfaces all of them here rather + /// than only the (arbitrarily chosen) lowest-index one. Each is stamped + /// with its own branch's task id before being persisted/returned, so the + /// caller's subsequent `resume` can address each individually (see + /// `resume_from_inner`'s `resume_map` / `Command::resume_tasks`). + pub(super) async fn handle_interrupt_boundary( + &self, + ctx: &mut RunCtx<'_, State, Update>, + sb: StepBoundary<'_>, + state: State, + interrupted: Vec<(usize, Interrupt)>, + ) -> Result> { + ctx.disarm_drop_guard(); + if let Err(err) = self.require_interrupt_durability(&ctx.thread_id) { + return self.fail_and_return(ctx, err).await; + } + let stamped: Vec = interrupted + .into_iter() + .map(|(index, interrupt)| interrupt.with_task_id(sb.active[index].task_id.clone())) + .collect(); + // Deferred routing, same as the failure boundary above: the + // completed siblings (whichever side of the interrupted branches + // they fall on) are not routed here. `pending` is exactly + // `sb.stalled` (every interrupted branch — no error can be mixed in + // here, since the executor dispatches a step with any failure to + // `handle_failure_boundary` first), and `completed_tasks` carries + // every completed node id forward (merged with anything already + // carried from an earlier resume of this step) for `advance` to + // route once the pending set finishes. + let pending: Vec = sb.stalled.iter().map(|(_, a)| a.clone()).collect(); + let completed = self.merged_completed(ctx, sb.completed, sb.goto_map); + let pending_nodes = activation_nodes(&pending); + let interrupt_ids: Vec = stamped + .iter() + .map(|i| InterruptId::new(i.id.clone())) + .collect(); + // An interrupt hands control back to the caller expecting a fully + // durable pause point: settle any in-flight Async background writes + // first, failing the run if one was lost (a broken lineage cannot + // be safely resumed from). + if let Err(err) = ctx.async_writes.drain().await { + return self.fail_and_return(ctx, err).await; + } + let checkpoint_id = match self + .persist_checkpoint( + ctx, + BoundaryCheckpoint { + state: &state, + pending: &pending, + completed, + child_runs: sb.child_runs_meta, + task_writes: sb.task_writes.to_vec(), + }, + sb.step, + stamped.clone(), + &pending_nodes, + ) + .await + { + Ok(id) => id, + Err(persist_err) => return self.fail_and_return(ctx, persist_err).await, + }; + + let mut status = ctx.base_status(); + status.status = ExecutionStatus::Interrupted; + status.current_step = sb.step; + status.active_nodes = pending_nodes; + status.pending_interrupts = interrupt_ids; + status.checkpoint_id = checkpoint_id.clone(); + ctx.save_status(status.clone()).await; + + Ok(GraphExecution { + state, + run_id: ctx.run_id.clone(), + graph_id: self.graph_id.clone(), + root_run_id: ctx.root_run_id.clone(), + parent_run_id: ctx.parent_run_id.clone(), + child_runs: std::mem::take(&mut ctx.all_child_runs), + visited: std::mem::take(&mut ctx.visited), + steps: sb.step, + interrupts: stamped, + status, + checkpoint_id, + drained: false, + }) + } + + /// The cancellation boundary (I4 part 2): the run's cooperative + /// cancellation token was observed cancelled, either between supersteps + /// or while this step's node handlers were still in flight and had to be + /// abandoned (raced against the token via + /// [`super::executor::CompiledGraph::run_step_with_cancel`]). + /// + /// Unlike the failure/interrupt boundaries, nothing from this step is + /// trusted to have completed — a mid-step cancellation abandons the + /// step's future rather than awaiting it to a folded result — so `active` + /// (exactly what the next superstep would have run) is persisted whole + /// as the resumable checkpoint's pending set, mirroring how the failure + /// boundary reuses the checkpoint machinery. Persists a resumable + /// checkpoint (on a checkpointed thread), records a `Cancelled` status, + /// and returns `Ok` (cancellation is a normal, requested outcome, not an + /// error) carrying no interrupts. + pub(super) async fn handle_cancel_boundary( + &self, + ctx: &mut RunCtx<'_, State, Update>, + active: &[Activation], + state: &State, + ) -> Result> { + ctx.disarm_drop_guard(); + // Settle in-flight Async background writes before persisting the + // cancellation checkpoint, same as the failure boundary — best + // effort, since a lost background write here must not turn a + // successfully-requested cancellation into a hard error. + let _ = ctx.async_writes.drain().await; + let checkpoint_id = self + .persist_stop_checkpoint(ctx, state, active, "cancelled") + .await + .unwrap_or(None); + + let mut status = ctx.base_status(); + status.status = ExecutionStatus::Cancelled; + status.current_step = ctx.steps; + status.active_nodes = activation_nodes(active); + status.checkpoint_id = checkpoint_id.clone(); + status.ended_at = Some(SystemTime::now()); + ctx.save_status(status.clone()).await; + ctx.emit(GraphEvent::RunCancelled { + run_id: ctx.run_id.clone(), + }); + + Ok(GraphExecution { + state: state.clone(), + run_id: ctx.run_id.clone(), + graph_id: self.graph_id.clone(), + root_run_id: ctx.root_run_id.clone(), + parent_run_id: ctx.parent_run_id.clone(), + child_runs: std::mem::take(&mut ctx.all_child_runs), + visited: std::mem::take(&mut ctx.visited), + steps: ctx.steps, + interrupts: Vec::new(), + status, + checkpoint_id, + drained: false, + }) + } + + /// The graceful-drain boundary: the run's [`super::DrainSignal`] was + /// observed raised between supersteps (`execute_run` polls it only + /// there), so the previous step finished and committed normally and + /// `active` — the next step's whole set, none of which has started — is + /// persisted as the resumable checkpoint's pending set, exactly as the + /// cancellation boundary does. Records a `Drained` status, emits + /// [`GraphEvent::RunDrained`], and returns `Ok` with + /// [`GraphExecution::drained`] set. A later `resume`/`retry` continues + /// from that checkpoint like any other pending-tasks checkpoint. + pub(super) async fn handle_drain_boundary( + &self, + ctx: &mut RunCtx<'_, State, Update>, + active: &[Activation], + state: &State, + ) -> Result> { + ctx.disarm_drop_guard(); + // A drain is a requested stop, so like cancellation a lost Async + // background write must not turn it into a hard error. + let _ = ctx.async_writes.drain().await; + let checkpoint_id = self + .persist_stop_checkpoint(ctx, state, active, "drained") + .await + .unwrap_or(None); + + let mut status = ctx.base_status(); + status.status = ExecutionStatus::Drained; + status.current_step = ctx.steps; + status.active_nodes = activation_nodes(active); + status.checkpoint_id = checkpoint_id.clone(); + status.ended_at = Some(SystemTime::now()); + ctx.save_status(status.clone()).await; + ctx.emit(GraphEvent::RunDrained { + run_id: ctx.run_id.clone(), + steps: ctx.steps, + }); + + Ok(GraphExecution { + state: state.clone(), + run_id: ctx.run_id.clone(), + graph_id: self.graph_id.clone(), + root_run_id: ctx.root_run_id.clone(), + parent_run_id: ctx.parent_run_id.clone(), + child_runs: std::mem::take(&mut ctx.all_child_runs), + visited: std::mem::take(&mut ctx.visited), + steps: ctx.steps, + interrupts: Vec::new(), + status, + checkpoint_id, + drained: true, + }) + } + + /// Persists a resumable stop-boundary checkpoint for a cancellation or + /// drain, mirroring [`Self::persist_failure_checkpoint`]: `next_nodes` + /// schedules exactly the activations that were still pending when the + /// stop was observed, so `resume`/`retry` re-runs exactly what did not + /// complete. `marker` (`"cancelled"`/`"drained"`) is stamped `true` into + /// the checkpoint metadata. A no-op returning `None` without a + /// checkpointer/thread, exactly like the failure boundary. + async fn persist_stop_checkpoint( + &self, + ctx: &RunCtx<'_, State, Update>, + state: &State, + pending: &[Activation], + marker: &str, + ) -> Result> { + let (Some(checkpointer), Some(thread)) = (&self.checkpointer, &ctx.thread_id) else { + return Ok(None); + }; + let (channel_versions, channel_deltas, versions_seen) = + self.channel_checkpoint_fields(ctx, state); + // Nothing in `pending` ran this boundary, so the only replay memos + // (and interrupt acks) it can own are the ones this run was seeded + // with on resume; re-persist them so a stop straight after a resume + // does not strip a task of its `durable_task` memos or deferred + // `interrupt_after` result. + let pending_writes: Vec = pending + .iter() + .filter_map(|a| ctx.task_writes.get(a.task_id.as_str())) + .flat_map(|writes| writes.iter().cloned()) + .collect(); + let metadata = Self::with_carried_acks( + serde_json::json!({ + "source": "loop", + "step": ctx.steps, + "recursion": ctx.recursion_meta, + marker: true, + "node_visits": node_visits_to_json(&ctx.node_visits), + }), + ctx, + pending, + ); + let checkpoint = Checkpoint::new( + state.clone(), + pending.iter().map(PendingActivation::from).collect(), + ) + .with_thread_id(thread.to_string()) + .with_checkpoint_id(next_checkpoint_id()) + .with_run_id(ctx.run_id.to_string()) + .with_parent_checkpoint_id(ctx.parent_checkpoint.clone()) + .with_namespace(self.namespace.clone()) + .with_pending_writes(pending_writes) + .with_barrier_arrivals(barriers_to_persisted(&ctx.barrier_arrivals)) + .with_channel_versions(channel_versions) + .with_channel_deltas(channel_deltas) + .with_versions_seen(versions_seen) + .with_metadata(metadata); + let writes = checkpoint.pending_writes.clone(); + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let id = checkpointer.put(checkpoint).await?; + if !writes.is_empty() { + checkpointer.put_writes(&config, &writes).await?; + } + self.emit( + &ctx.run_id, + GraphEvent::CheckpointSaved { + checkpoint_id: id.clone(), + step: Some(ctx.steps), + }, + ); + Ok(Some(id)) + } + + /// Emits a [`GraphEvent::RunFailed`] and records a terminal `Failed` + /// status for a run that aborted with `err`. + /// + /// `checkpoint_id` is the resumable failure-boundary checkpoint when the + /// run left one (a node-handler failure on a checkpointed thread), or + /// `None` for a structural/non-resumable abort. When present it is + /// recorded on the status so an observer can locate the checkpoint to + /// `resume`/`retry` from. + pub(super) async fn fail_run( + &self, + run_id: &RunId, + thread_id: &Option, + started_at: SystemTime, + steps: usize, + err: &TinyAgentsError, + checkpoint_id: Option, + ) { + self.emit( + run_id, + GraphEvent::RunFailed { + run_id: run_id.clone(), + error: err.to_string(), + }, + ); + let mut status = self.base_status(run_id, thread_id, started_at); + status.status = ExecutionStatus::Failed; + status.current_step = steps; + status.ended_at = Some(SystemTime::now()); + status.error = Some(err.to_string()); + status.checkpoint_id = checkpoint_id; + self.save_status(status).await; + } + + /// Records a terminal `Failed` status for `err` (via [`Self::fail_run`], + /// reading identity/timing off `ctx`) and returns it as `Err`. + /// + /// Used at every early-exit path in `execute_run` — a guard trip, a + /// node-runner error, a reducer merge, a routing resolution, or a + /// checkpoint persist — so the run transitions 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. + pub(super) async fn fail_and_return( + &self, + ctx: &mut RunCtx<'_, State, Update>, + err: TinyAgentsError, + ) -> Result { + ctx.disarm_drop_guard(); + let _ = ctx.async_writes.drain().await; + self.fail_run( + &ctx.run_id, + &ctx.thread_id, + ctx.started_at, + ctx.steps, + &err, + None, + ) + .await; + Err(err) + } + + /// Persists a resumable failure-boundary checkpoint for a node-handler + /// failure that survived the node-retry policy. + /// + /// Mirrors the interrupt boundary: `next_nodes` schedules the failed + /// node (and any not-yet-run members of the step) so `resume`/`retry` + /// re-runs exactly what did not complete, while `completed_tasks` + /// records the branches that already succeeded (their updates are + /// folded into `state` before this is called). The rendered error and + /// failed node id are stamped into the checkpoint metadata for + /// diagnosis. A no-op returning `None` when no checkpointer/thread is + /// configured — the run then aborts without a resumable checkpoint, + /// exactly as before this policy existed. + async fn persist_failure_checkpoint( + &self, + ctx: &RunCtx<'_, State, Update>, + boundary: BoundaryCheckpoint<'_, State>, + step: usize, + failed_node: &NodeId, + error: &TinyAgentsError, + ) -> Result> { + let (Some(checkpointer), Some(thread)) = (&self.checkpointer, &ctx.thread_id) else { + return Ok(None); + }; + let pending_writes = Self::completion_writes(&boundary.completed, &boundary.task_writes); + let (channel_versions, channel_deltas, versions_seen) = + self.channel_checkpoint_fields(ctx, boundary.state); + let checkpoint = Checkpoint::new( + boundary.state.clone(), + boundary + .pending + .iter() + .map(PendingActivation::from) + .collect(), + ) + .with_thread_id(thread.to_string()) + .with_checkpoint_id(next_checkpoint_id()) + .with_run_id(ctx.run_id.to_string()) + .with_parent_checkpoint_id(ctx.parent_checkpoint.clone()) + .with_namespace(self.namespace.clone()) + .with_completed(boundary.completed) + .with_pending_writes(pending_writes) + .with_barrier_arrivals(barriers_to_persisted(&ctx.barrier_arrivals)) + .with_channel_versions(channel_versions) + .with_channel_deltas(channel_deltas) + .with_versions_seen(versions_seen) + .with_metadata(Self::with_carried_acks( + serde_json::json!({ + "source": "loop", + "step": step, + "recursion": ctx.recursion_meta, + "child_runs": boundary.child_runs, + "failed_node": failed_node.as_str(), + "error": error.to_string(), + "node_visits": node_visits_to_json(&ctx.node_visits), + }), + ctx, + boundary.pending, + )); + let writes = checkpoint.pending_writes.clone(); + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let id = checkpointer.put(checkpoint).await?; + // Also record the ledger through the write protocol, so backends + // that implement it can answer "did this task run?" without loading + // the whole state payload. + checkpointer.put_writes(&config, &writes).await?; + self.emit( + &ctx.run_id, + GraphEvent::CheckpointSaved { + checkpoint_id: id.clone(), + step: Some(step), + }, + ); + Ok(Some(id)) + } + + /// Persists a loop-boundary checkpoint (the normal step boundary, or an + /// interrupt boundary when `interrupts`/`interrupted` are non-empty). + async fn persist_checkpoint( + &self, + ctx: &RunCtx<'_, State, Update>, + boundary: BoundaryCheckpoint<'_, State>, + step: usize, + interrupts: Vec, + interrupted: &[NodeId], + ) -> Result> { + let (Some(checkpointer), Some(thread)) = (&self.checkpointer, &ctx.thread_id) else { + return Ok(None); + }; + let checkpoint = + self.build_loop_checkpoint(ctx, thread, boundary, step, interrupts, interrupted); + let writes = checkpoint.pending_writes.clone(); + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let id = checkpointer.put(checkpoint).await?; + checkpointer.put_writes(&config, &writes).await?; + self.emit( + &ctx.run_id, + GraphEvent::CheckpointSaved { + checkpoint_id: id.clone(), + step: Some(step), + }, + ); + Ok(Some(id)) + } + + /// Persists a boundary checkpoint without blocking the superstep loop + /// ([`DurabilityMode::Async`]). + /// + /// The checkpoint id is minted up front and returned immediately so the + /// loop keeps chaining lineage onto it, while the actual `put` (and the + /// [`GraphEvent::CheckpointSaved`] emitted on its success) runs on a + /// spawned background task tracked in `ctx.async_writes`. + /// + /// # Failure semantics + /// + /// A background write error is never dropped: it is recorded in + /// `ctx.async_writes` and surfaced by the executor at the next + /// durability boundary, or at the latest when the run drains all + /// in-flight writes at its terminal/interrupt boundary — so the run + /// result reflects persistence failures. Because the `CheckpointSaved` + /// event is emitted from the background task, its ordering relative to + /// subsequent step events is not deterministic under `Async` durability. + /// + /// Outside a tokio runtime there is nothing to spawn onto, so the write + /// happens inline — degrading to [`DurabilityMode::Sync`] behavior. + async fn persist_checkpoint_nonblocking( + &self, + ctx: &mut RunCtx<'_, State, Update>, + boundary: BoundaryCheckpoint<'_, State>, + step: usize, + ) -> Result> { + let (Some(checkpointer), Some(thread)) = (&self.checkpointer, &ctx.thread_id) else { + return Ok(None); + }; + let thread = thread.clone(); + let checkpoint = self.build_loop_checkpoint(ctx, &thread, boundary, step, Vec::new(), &[]); + let id = CheckpointId::new(checkpoint.checkpoint_id.clone()); + // M5: mirror the synchronous path, which persists both the state + // record (`put`) and the write ledger (`put_writes`). Without this + // the ledger tooling sees no completion markers for any checkpoint + // written under `DurabilityMode::Async`. + let writes = checkpoint.pending_writes.clone(); + let write_config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + + match tokio::runtime::Handle::try_current() { + Ok(handle) => { + let checkpointer = Arc::clone(checkpointer); + let sink = self.event_sink.clone(); + // Build the envelope (and so claim its `seq` value) on the + // calling thread, before spawning: the background task's + // completion order relative to other work is not + // deterministic, but the sequence number it carries still + // reflects when this write was *requested*. + let envelope = sink.as_ref().map(|_| { + self.envelope( + &ctx.run_id, + GraphEvent::CheckpointSaved { + checkpoint_id: id.clone(), + step: Some(step), + }, + ) + }); + ctx.async_writes.spawn_ordered(&handle, async move { + let id = checkpointer.put(checkpoint).await?; + checkpointer.put_writes(&write_config, &writes).await?; + if let (Some(sink), Some(envelope)) = (sink, envelope) { + sink.emit(envelope); + } + Ok(id) + }); + Ok(Some(id)) + } + Err(_) => { + let id = checkpointer.put(checkpoint).await?; + checkpointer.put_writes(&write_config, &writes).await?; + self.emit( + &ctx.run_id, + GraphEvent::CheckpointSaved { + checkpoint_id: id.clone(), + step: Some(step), + }, + ); + Ok(Some(id)) + } + } + } + + /// Builds the `completed_tasks`/`completed_routes` pair for a + /// failure/interrupt boundary checkpoint: this step's own completed + /// branches (`sb.completed`, with their `goto_map` routing captured + /// positionally), prefixed by any node ids (and their persisted + /// `goto`) already carried forward from an earlier interrupt/failure of + /// this *same* logical step (`ctx.carried_completed` — set once, at + /// resume, from the loaded checkpoint's `completed_tasks`/ + /// `completed_routes`, and left untouched here; only [`Self::advance`] + /// consumes it, once the step finally finishes routing). This is what + /// lets a step interrupt or fail more than once across repeated resumes + /// without losing track of which of its branches have already + /// completed, or what they explicitly routed to (R1). + fn merged_completed( + &self, + ctx: &RunCtx<'_, State, Update>, + completed: &[(usize, Activation)], + goto_map: &HashMap>, + ) -> Vec { + let mut out: Vec = Vec::new(); + if let Some(carried) = &ctx.carried_completed { + for (node, goto) in carried { + // No task id is carried across a resume: `RunCtx::carried_completed` + // stores only the node id and persisted routing. + out.push(crate::checkpoint::CompletedTask::with_routes( + TaskId::from(String::new()), + node.clone(), + goto.clone(), + )); + } + } + for (index, activation) in completed { + out.push(crate::checkpoint::CompletedTask::with_routes( + activation.task_id.clone(), + activation.node.clone(), + goto_map.get(index).cloned().unwrap_or_default(), + )); + } + out + } + + /// Stamps `metadata.acknowledged_interrupts` with the executor-injected + /// interrupt phases (`":"`, see + /// [`RunCtx::acknowledged_interrupts`]) already acknowledged for any + /// task still in `pending`, so they survive this boundary. Without this + /// a task that acknowledged its `interrupt_before` pause and then paused + /// again (its `interrupt_after`, or an interrupt it emitted itself) + /// would be paused *before* a second time on the next resume, since + /// resume derives acknowledgements from the latest checkpoint's own + /// interrupts. Acks of tasks no longer pending are dropped; the key is + /// omitted entirely when nothing carries over. + fn with_carried_acks( + mut metadata: serde_json::Value, + ctx: &RunCtx<'_, State, Update>, + pending: &[Activation], + ) -> serde_json::Value { + let carried: Vec<&String> = ctx + .acknowledged_interrupts + .iter() + .filter(|key| { + pending.iter().any(|a| { + key.split_once(':') + .is_some_and(|(_, task)| task == a.task_id.as_str()) + }) + }) + .collect(); + if !carried.is_empty() { + let mut carried: Vec<&String> = carried; + carried.sort(); + metadata["acknowledged_interrupts"] = serde_json::json!(carried); + } + metadata + } + + /// Records completion markers for the tasks that finished in the step a + /// boundary checkpoint closes. + /// + /// A graph's `Update` carries no `Serialize` bound, so the executor + /// cannot persist *what* a task wrote — but it does not need to: the + /// applied value is already durable in the checkpoint's `state`. What + /// was missing was the other half, the per-task record of *that* it + /// ran, which is what lets a resume distinguish "already done" from + /// "not yet started". See [`PendingWrite`](crate::checkpoint::PendingWrite)'s + /// docs for why that distinction is the whole point of the ledger. + /// + /// The task id is persisted on the activation itself, so a resume can + /// match a marker to one fan-out task rather than every task with its + /// node. + /// + /// `task_writes` — the stalled branches' replay memos (durable-task + /// writes, deferred `interrupt_after` results) — are appended verbatim. + /// They belong to tasks that are still *pending*, and resume tells them + /// apart from completion markers by + /// [`PendingWrite::is_task_replay`](crate::checkpoint::PendingWrite::is_task_replay). + fn completion_writes( + completed: &[crate::checkpoint::CompletedTask], + task_writes: &[crate::checkpoint::PendingWrite], + ) -> Vec { + completed + .iter() + .map(|task| { + crate::checkpoint::PendingWrite::completion_marker( + task.node.clone(), + task.task_id.clone(), + ) + }) + .chain(task_writes.iter().cloned()) + .collect() + } + + /// Builds the loop-boundary [`Checkpoint`] record shared by the sync and + /// async persist paths, minting a fresh checkpoint id. + fn build_loop_checkpoint( + &self, + ctx: &RunCtx<'_, State, Update>, + thread: &ThreadId, + boundary: BoundaryCheckpoint<'_, State>, + step: usize, + interrupts: Vec, + interrupted: &[NodeId], + ) -> Checkpoint { + let mut metadata = serde_json::json!({ + "source": "loop", + "step": step, + "recursion": ctx.recursion_meta, + "child_runs": boundary.child_runs, + "node_visits": node_visits_to_json(&ctx.node_visits), + }); + // 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::>() + ); + } + let metadata = Self::with_carried_acks(metadata, ctx, boundary.pending); + let pending_writes = Self::completion_writes(&boundary.completed, &boundary.task_writes); + let (channel_versions, channel_deltas, versions_seen) = + self.channel_checkpoint_fields(ctx, boundary.state); + Checkpoint::new( + boundary.state.clone(), + boundary + .pending + .iter() + .map(PendingActivation::from) + .collect(), + ) + .with_thread_id(thread.to_string()) + .with_checkpoint_id(next_checkpoint_id()) + .with_run_id(ctx.run_id.to_string()) + .with_parent_checkpoint_id(ctx.parent_checkpoint.clone()) + .with_namespace(self.namespace.clone()) + .with_completed(boundary.completed) + .with_pending_writes(pending_writes) + .with_barrier_arrivals(barriers_to_persisted(&ctx.barrier_arrivals)) + .with_channel_versions(channel_versions) + .with_channel_deltas(channel_deltas) + .with_versions_seen(versions_seen) + .with_interrupts(interrupts) + .with_metadata(metadata) + } + + pub(super) fn base_status( + &self, + run_id: &RunId, + thread_id: &Option, + started_at: SystemTime, + ) -> GraphRunStatus { + let mut status = GraphRunStatus::new( + run_id.clone(), + self.graph_id.clone(), + ExecutionStatus::Running, + ); + status.thread_id = thread_id.clone(); + status.checkpoint_namespace = self.namespace.clone(); + status.started_at = started_at; + status.updated_at = SystemTime::now(); + status + } + + /// Best-effort status write; never aborts the run on a status-store + /// error, but logs it (M4) so a dead status backend is at least visible + /// rather than silently discarded. + pub(super) async fn save_status(&self, status: GraphRunStatus) { + if let Some(store) = &self.status_store { + let run_id = status.run_id.clone(); + if let Err(err) = store.put_status(status).await { + tracing::warn!( + "[graph:status] failed to persist run status for run `{run_id}`: {err}" + ); + } + } + } +} diff --git a/crates/tinyagents-graph/src/compiled/drain_test.rs b/crates/tinyagents-graph/src/compiled/drain_test.rs new file mode 100644 index 00000000..322c4527 --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/drain_test.rs @@ -0,0 +1,267 @@ +//! Graceful drain (`RunOptions::drain`, `DrainSignal`/`DrainHandle`): the +//! superstep in flight finishes and commits, the next step's activations are +//! checkpointed instead of run, and the run reports `Drained` — resumable +//! to the same final state an undrained run reaches. + +use super::*; +use crate::builder::{GraphBuilder, NodeContext}; +use crate::checkpoint::InMemoryCheckpointer; +use crate::command::NodeResult; +use crate::stream::{CollectingSink, GraphEvent}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use tinyagents_harness::ids::ExecutionStatus; + +/// Three sequential supersteps `a (+1) -> b (+10) -> c (+100)`, each node +/// counted; `a` raises `drain` (when given) from inside its handler, i.e. +/// while step 1 is in flight. +fn chain(counts: [Arc; 3], drain: Option) -> CompiledGraph { + let [a, b, c] = counts; + GraphBuilder::::overwrite() + .add_node("a", move |s, _c: NodeContext| { + let a = a.clone(); + let drain = drain.clone(); + async move { + a.fetch_add(1, AtomicOrdering::SeqCst); + if let Some(drain) = drain { + drain.drain(); + } + Ok(NodeResult::Update(s + 1)) + } + }) + .add_node("b", move |s, _c: NodeContext| { + let b = b.clone(); + async move { + b.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(s + 10)) + } + }) + .add_node("c", move |s, _c: NodeContext| { + let c = c.clone(); + async move { + c.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(s + 100)) + } + }) + .add_sequence(["a", "b", "c"]) + .set_entry("a") + .set_finish("c") + .compile() + .unwrap() +} + +fn counters() -> [Arc; 3] { + std::array::from_fn(|_| Arc::new(AtomicUsize::new(0))) +} + +#[tokio::test] +async fn drain_finishes_the_in_flight_step_and_stops_before_the_next() { + let counts = counters(); + let (handle, signal) = DrainSignal::new(); + let sink = Arc::new(CollectingSink::new()); + let graph = chain(counts.clone(), Some(handle.clone())) + .with_checkpointer(Arc::new(InMemoryCheckpointer::::new())) + .with_event_sink(sink.clone()); + + let run = graph + .run_with_thread_options("drain", 0, RunOptions::with_drain(signal)) + .await + .unwrap(); + + assert!(run.drained); + assert_eq!(run.status.status, ExecutionStatus::Drained); + assert!(run.status.is_terminal()); + assert!(!run.is_interrupted()); + assert!(handle.is_requested()); + // Step 1 (`a`) completed and committed; `b`/`c` never started. + assert_eq!(run.state, 1); + assert_eq!(run.steps, 1); + assert_eq!(counts[0].load(AtomicOrdering::SeqCst), 1); + assert_eq!(counts[1].load(AtomicOrdering::SeqCst), 0); + assert_eq!(counts[2].load(AtomicOrdering::SeqCst), 0); + assert_eq!( + run.status + .active_nodes + .iter() + .map(|n| n.as_str()) + .collect::>(), + vec!["b"] + ); + assert!(run.checkpoint_id.is_some()); + + // The terminal event is `RunDrained` (and only that). + let events = sink.events(); + assert!(events.iter().any(|e| matches!( + e, + GraphEvent::RunDrained { run_id, steps: 1 } if *run_id == run.run_id + ))); + assert!(!events.iter().any(|e| matches!( + e, + GraphEvent::RunCompleted { .. } + | GraphEvent::RunCancelled { .. } + | GraphEvent::RunFailed { .. } + ))); + + // The checkpoint names `b` as the pending work and is marked drained. + let snapshot = graph.get_state("drain", None).await.unwrap().unwrap(); + assert_eq!(snapshot.values, 1); + assert_eq!( + snapshot + .next_nodes + .iter() + .map(|n| n.as_str()) + .collect::>(), + vec!["b"] + ); + assert!(!snapshot.metadata.has_interrupts); + let tuple = graph + .checkpointer + .as_ref() + .unwrap() + .get_tuple(snapshot.config) + .await + .unwrap() + .unwrap(); + assert_eq!( + tuple.checkpoint.metadata["drained"], + serde_json::json!(true) + ); +} + +#[tokio::test] +async fn resume_after_drain_reaches_the_same_state_as_an_undrained_run() { + // Reference: the same graph, no drain. + let reference = chain(counters(), None).run(0).await.unwrap(); + assert_eq!(reference.status.status, ExecutionStatus::Completed); + + let counts = counters(); + let (handle, signal) = DrainSignal::new(); + let graph = chain(counts.clone(), Some(handle)) + .with_checkpointer(Arc::new(InMemoryCheckpointer::::new())); + let drained = graph + .run_with_thread_options("resume", 0, RunOptions::with_drain(signal)) + .await + .unwrap(); + assert!(drained.drained); + + // `retry`/`resume` continue from the drained checkpoint like any other + // pending-tasks checkpoint: `b` then `c`, each once. + let resumed = graph.retry("resume").await.unwrap(); + assert!(!resumed.drained); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(resumed.state, reference.state); + assert_eq!(resumed.state, 111); + assert_eq!( + resumed + .visited + .iter() + .map(|n| n.as_str()) + .collect::>(), + vec!["b", "c"] + ); + assert_eq!(counts[0].load(AtomicOrdering::SeqCst), 1); + assert_eq!(counts[1].load(AtomicOrdering::SeqCst), 1); + assert_eq!(counts[2].load(AtomicOrdering::SeqCst), 1); +} + +#[tokio::test] +async fn drain_requested_before_the_run_starts_runs_nothing() { + let counts = counters(); + let (handle, signal) = DrainSignal::new(); + handle.drain(); + let graph = + chain(counts.clone(), None).with_checkpointer(Arc::new(InMemoryCheckpointer::::new())); + let run = graph + .run_with_thread_options("early", 5, RunOptions::with_drain(signal)) + .await + .unwrap(); + assert!(run.drained); + assert_eq!(run.steps, 0); + assert_eq!(run.state, 5); + assert!(counts.iter().all(|c| c.load(AtomicOrdering::SeqCst) == 0)); + // Still resumable from the entry node. + let resumed = graph.retry("early").await.unwrap(); + assert_eq!(resumed.state, 116); +} + +#[tokio::test] +async fn drain_without_a_thread_reports_drained_without_a_checkpoint() { + let counts = counters(); + let (handle, signal) = DrainSignal::new(); + let graph = chain(counts.clone(), Some(handle)); + let run = graph + .run_with_options(0, RunOptions::with_drain(signal)) + .await + .unwrap(); + assert!(run.drained); + assert_eq!(run.status.status, ExecutionStatus::Drained); + assert_eq!(run.state, 1); + assert!(run.checkpoint_id.is_none()); + assert_eq!(counts[1].load(AtomicOrdering::SeqCst), 0); +} + +#[tokio::test] +async fn undrained_runs_report_drained_false_on_every_outcome() { + let completed = chain(counters(), None).run(0).await.unwrap(); + assert!(!completed.drained); + let (_handle, signal) = DrainSignal::new(); + let unsignalled = chain(counters(), None) + .run_with_options(0, RunOptions::with_drain(signal)) + .await + .unwrap(); + assert!(!unsignalled.drained); + assert_eq!(unsignalled.status.status, ExecutionStatus::Completed); +} + +#[tokio::test] +async fn drain_straight_after_resume_keeps_a_deferred_interrupt_after_result() { + // interrupt_after pause -> resume with drain already raised: the drain + // boundary must re-persist `b`'s deferred result (and its ack), so the + // eventual retry replays it instead of running `b` a second time. + let b_runs = Arc::new(AtomicUsize::new(0)); + let runs = b_runs.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", move |s, _c: NodeContext| { + let runs = runs.clone(); + async move { + runs.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(s + 10)) + } + }) + .add_sequence(["a", "b"]) + .set_entry("a") + .set_finish("b") + .interrupt_after(["b"]) + .compile() + .unwrap() + .with_checkpointer(Arc::new(InMemoryCheckpointer::::new())); + + let paused = graph.run_with_thread("drain-after", 0).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); + + let (handle, signal) = DrainSignal::new(); + handle.drain(); + let drained = graph + .resume_with_options( + "drain-after", + crate::command::Command::new(), + RunOptions::with_drain(signal), + ) + .await + .unwrap(); + assert!(drained.drained); + assert_eq!(drained.state, 1); + + let done = graph.retry("drain-after").await.unwrap(); + assert_eq!(done.status.status, ExecutionStatus::Completed); + assert_eq!(done.state, 11); + assert_eq!( + b_runs.load(AtomicOrdering::SeqCst), + 1, + "replayed, not re-run" + ); +} diff --git a/crates/tinyagents-graph/src/compiled/durable_task_test.rs b/crates/tinyagents-graph/src/compiled/durable_task_test.rs new file mode 100644 index 00000000..45bfce5e --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/durable_task_test.rs @@ -0,0 +1,325 @@ +//! `NodeContext::durable_task`: per-task memoisation of a side-effecting +//! sub-step, keyed by `(task_id, key)` in the checkpoint write ledger. +//! +//! The contract under test: a handler is re-run from its start after an +//! interrupt/resume, a failure/retry, or an in-process node retry, but a +//! side effect wrapped in `durable_task` happens **once** — the re-run gets +//! the stored output back without polling the future — and distinct keys +//! are memoised independently of each other. + +use super::*; +use crate::builder::{GraphBuilder, NodeContext}; +use crate::checkpoint::{Checkpointer, FileCheckpointer, InMemoryCheckpointer}; +use crate::command::{Command, Interrupt, NodeResult}; +use serde_json::json; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use tinyagents_harness::ids::ExecutionStatus; +use tinyagents_harness::retry::RetryPolicy; + +fn memory() -> Arc> { + Arc::new(InMemoryCheckpointer::::new()) +} + +/// A single node that performs a counted side effect inside +/// `durable_task("side-effect")`, then either pauses (first pass, no resume +/// value) or commits `state + `. +fn interrupting_graph(effects: Arc) -> GraphBuilder { + GraphBuilder::::overwrite() + .add_node("work", move |s, ctx: NodeContext| { + let effects = effects.clone(); + async move { + let n: usize = ctx + .durable_task("side-effect", async { + effects.fetch_add(1, AtomicOrdering::SeqCst); + Ok(effects.load(AtomicOrdering::SeqCst)) + }) + .await?; + if ctx.resume.is_none() { + return Ok(NodeResult::Interrupt(Interrupt::new( + "work", + json!({ "ask": "continue?" }), + ))); + } + Ok(NodeResult::Update(s + n as i32)) + } + }) + .set_entry("work") + .set_finish("work") +} + +#[tokio::test] +async fn durable_task_is_memoised_across_an_interrupt_and_resume() { + let effects = Arc::new(AtomicUsize::new(0)); + let graph = interrupting_graph(effects.clone()) + .compile() + .unwrap() + .with_checkpointer(memory()); + + let paused = graph.run_with_thread("memo", 100).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); + + // The memo travelled with the pending task into the checkpoint ledger. + let snapshot = graph.get_state("memo", None).await.unwrap().unwrap(); + let tuple = graph + .checkpointer + .as_ref() + .unwrap() + .get_tuple(snapshot.config) + .await + .unwrap() + .unwrap(); + let memos: Vec<_> = tuple + .pending_writes + .iter() + .filter(|w| w.is_durable_task()) + .collect(); + assert_eq!(memos.len(), 1); + assert_eq!(memos[0].durable_task_key(), Some("side-effect")); + assert_eq!(memos[0].payload, json!(1)); + assert!(memos[0].idx >= 1, "durable-task writes never reuse idx 0"); + + // The handler re-runs from the top on resume, but the side effect does + // not: the memoised `1` is what it sees. + let resumed = graph + .resume("memo", Command::resume(json!({ "go": true }))) + .await + .unwrap(); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(resumed.state, 101); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); +} + +/// `work` performs a counted side effect, then fails hard on its first +/// attempt (per `fails`) and succeeds afterwards. +fn failing_graph(effects: Arc, fails: Arc) -> GraphBuilder { + GraphBuilder::::overwrite() + .add_node("work", move |s, ctx: NodeContext| { + let effects = effects.clone(); + let fails = fails.clone(); + async move { + let n: usize = ctx + .durable_task("side-effect", async { + effects.fetch_add(1, AtomicOrdering::SeqCst); + Ok(effects.load(AtomicOrdering::SeqCst)) + }) + .await?; + if fails.fetch_add(1, AtomicOrdering::SeqCst) == 0 { + return Err(TinyAgentsError::Graph("crash after the side effect".into())); + } + Ok(NodeResult::Update(s + n as i32)) + } + }) + .set_entry("work") + .set_finish("work") +} + +#[tokio::test] +async fn durable_task_is_memoised_across_a_failure_and_retry() { + let effects = Arc::new(AtomicUsize::new(0)); + let fails = Arc::new(AtomicUsize::new(0)); + let graph = failing_graph(effects.clone(), fails.clone()) + .compile() + .unwrap() + .with_checkpointer(memory()); + + let err = graph.run_with_thread("retry", 100).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Graph(_)), "got {err:?}"); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); + + let resumed = graph.retry("retry").await.unwrap(); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(resumed.state, 101); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); + assert_eq!( + fails.load(AtomicOrdering::SeqCst), + 2, + "handler itself ran twice" + ); +} + +#[tokio::test] +async fn durable_task_keys_are_memoised_independently() { + // `a` runs, the handler crashes, then on retry `a` is a hit and `b` a + // fresh miss — proving memoisation is per key, not per node. + let a_effects = Arc::new(AtomicUsize::new(0)); + let b_effects = Arc::new(AtomicUsize::new(0)); + let attempts = Arc::new(AtomicUsize::new(0)); + let (a, b, att) = (a_effects.clone(), b_effects.clone(), attempts.clone()); + let graph = GraphBuilder::::overwrite() + .add_node("work", move |s, ctx: NodeContext| { + let (a, b, att) = (a.clone(), b.clone(), att.clone()); + async move { + let x: i32 = ctx + .durable_task("a", async { + a.fetch_add(1, AtomicOrdering::SeqCst); + Ok(10) + }) + .await?; + if att.fetch_add(1, AtomicOrdering::SeqCst) == 0 { + return Err(TinyAgentsError::Graph("crash between a and b".into())); + } + let y: i32 = ctx + .durable_task("b", async { + b.fetch_add(1, AtomicOrdering::SeqCst); + Ok(20) + }) + .await?; + Ok(NodeResult::Update(s + x + y)) + } + }) + .set_entry("work") + .set_finish("work") + .compile() + .unwrap() + .with_checkpointer(memory()); + + graph.run_with_thread("keys", 0).await.unwrap_err(); + assert_eq!(a_effects.load(AtomicOrdering::SeqCst), 1); + assert_eq!( + b_effects.load(AtomicOrdering::SeqCst), + 0, + "`b` never reached" + ); + + let resumed = graph.retry("keys").await.unwrap(); + assert_eq!(resumed.state, 30); + assert_eq!( + a_effects.load(AtomicOrdering::SeqCst), + 1, + "`a` replayed from memo" + ); + assert_eq!( + b_effects.load(AtomicOrdering::SeqCst), + 1, + "`b` ran fresh once" + ); +} + +#[tokio::test] +async fn durable_task_memo_is_shared_by_in_process_node_retries() { + // Under a node retry policy the handler is re-invoked in-process with a + // cloned context; the clone shares the memo buffer, so the retry hits. + let effects = Arc::new(AtomicUsize::new(0)); + let attempts = Arc::new(AtomicUsize::new(0)); + let (eff, att) = (effects.clone(), attempts.clone()); + let graph = GraphBuilder::::overwrite() + .add_node("work", move |s, ctx: NodeContext| { + let (eff, att) = (eff.clone(), att.clone()); + async move { + let n: i32 = ctx + .durable_task("side-effect", async { + eff.fetch_add(1, AtomicOrdering::SeqCst); + Ok(7) + }) + .await?; + if att.fetch_add(1, AtomicOrdering::SeqCst) < 2 { + return Err(TinyAgentsError::Model("transient".into())); + } + Ok(NodeResult::Update(s + n)) + } + }) + .set_entry("work") + .set_finish("work") + .compile() + .unwrap() + .with_node_retry(RetryPolicy::default().with_max_attempts(4)); + + let run = graph.run(0).await.unwrap(); + assert_eq!(run.state, 7); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 3); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); +} + +#[tokio::test] +async fn durable_task_memo_survives_a_file_checkpointer_restart() { + let dir = tempfile::tempdir().unwrap(); + let effects = Arc::new(AtomicUsize::new(0)); + { + let cp: Arc> = Arc::new(FileCheckpointer::::new(dir.path())); + let graph = interrupting_graph(effects.clone()) + .compile() + .unwrap() + .with_checkpointer(cp); + let paused = graph.run_with_thread("disk", 0).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); + } + let cp: Arc> = Arc::new(FileCheckpointer::::new(dir.path())); + let graph = interrupting_graph(effects.clone()) + .compile() + .unwrap() + .with_checkpointer(cp); + let resumed = graph + .resume("disk", Command::resume(json!({}))) + .await + .unwrap(); + assert_eq!(resumed.state, 1); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); +} + +#[tokio::test] +async fn durable_task_does_not_memoise_a_failed_future() { + // An `Err` from the wrapped future is returned as-is and leaves no memo, + // so the next attempt re-runs it. + let calls = Arc::new(AtomicUsize::new(0)); + let c = calls.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("work", move |s, ctx: NodeContext| { + let c = c.clone(); + async move { + let n: i32 = ctx + .durable_task("flaky", async { + if c.fetch_add(1, AtomicOrdering::SeqCst) == 0 { + Err(TinyAgentsError::Model("first call fails".into())) + } else { + Ok(3) + } + }) + .await?; + Ok(NodeResult::Update(s + n)) + } + }) + .set_entry("work") + .set_finish("work") + .compile() + .unwrap() + .with_node_retry(RetryPolicy::default().with_max_attempts(3)); + + let run = graph.run(0).await.unwrap(); + assert_eq!(run.state, 3); + assert_eq!(calls.load(AtomicOrdering::SeqCst), 2); +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn durable_task_memo_survives_a_sqlite_checkpointer_restart() { + use crate::checkpoint::SqliteCheckpointer; + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("checkpoints.db"); + let effects = Arc::new(AtomicUsize::new(0)); + { + let cp: Arc> = + Arc::new(SqliteCheckpointer::::open(&db_path).unwrap()); + let graph = interrupting_graph(effects.clone()) + .compile() + .unwrap() + .with_checkpointer(cp); + let paused = graph.run_with_thread("sqlite", 0).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); + } + let cp: Arc> = + Arc::new(SqliteCheckpointer::::open(&db_path).unwrap()); + let graph = interrupting_graph(effects.clone()) + .compile() + .unwrap() + .with_checkpointer(cp); + let resumed = graph + .resume("sqlite", Command::resume(json!({}))) + .await + .unwrap(); + assert_eq!(resumed.state, 1); + assert_eq!(effects.load(AtomicOrdering::SeqCst), 1); +} diff --git a/crates/tinyagents-graph/src/compiled/durable_test.rs b/crates/tinyagents-graph/src/compiled/durable_test.rs new file mode 100644 index 00000000..cf9b8ffb --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/durable_test.rs @@ -0,0 +1,198 @@ +//! Durable executor tests: the interrupt→resume and failure→retry scenarios +//! from `test.rs` (`interrupt_then_resume_reruns_node`, +//! `exhausted_retries_leave_a_resumable_failure_checkpoint`), replayed against +//! real on-disk checkpointers instead of [`InMemoryCheckpointer`]. +//! +//! Each scenario simulates a process restart between its write half and its +//! read/resume half: the checkpointer used for the first half is dropped +//! entirely, and a **fresh** checkpointer instance is opened against the same +//! on-disk location (the same directory for [`FileCheckpointer`], the same +//! database file for [`SqliteCheckpointer`]) for the second half. This +//! exercises the "close it, come back, resume from disk" path rather than +//! merely calling `.resume()`/`.retry()` on an already-warm in-process +//! checkpointer. + +use super::*; +use crate::builder::{GraphBuilder, NodeContext}; +#[cfg(feature = "sqlite")] +use crate::checkpoint::SqliteCheckpointer; +use crate::checkpoint::{Checkpointer, FileCheckpointer}; +use crate::command::{Command, Interrupt, NodeResult}; +use serde_json::json; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use tinyagents_harness::ids::ExecutionStatus; + +/// Same "human approval" graph shape as `interrupt_then_resume_reruns_node`: +/// pauses on first run, applies a resume-supplied bump on the second. +fn approve_graph() -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node("approve", |s, ctx: NodeContext| async move { + match ctx.resume { + Some(value) => { + let bump = value.get("bump").and_then(|v| v.as_i64()).unwrap_or(0) as i32; + Ok(NodeResult::Update(s + bump)) + } + None => Ok(NodeResult::Interrupt(Interrupt::new( + "approve", + json!({ "ask": "approve?" }), + ))), + } + }) + .add_node("done", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .set_entry("approve") + .add_edge("approve", "done") + .set_finish("done") + .compile() + .unwrap() +} + +/// Same "flaky node" shape as `flaky_graph` in `test.rs`: fails the first +/// `fail_times` invocations, then succeeds with `+1`. +fn flaky_graph(fail_times: usize, attempts: Arc) -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node("flaky", move |s, _c: NodeContext| { + let attempts = attempts.clone(); + async move { + let n = attempts.fetch_add(1, AtomicOrdering::SeqCst); + if n < fail_times { + Err(TinyAgentsError::Model(format!("transient blip {n}"))) + } else { + Ok(NodeResult::Update(s + 1)) + } + } + }) + .set_entry("flaky") + .set_finish("flaky") + .compile() + .unwrap() +} + +// ── FileCheckpointer ───────────────────────────────────────────────────── + +#[tokio::test] +async fn file_backend_interrupt_then_restart_then_resume() { + let dir = tempfile::tempdir().unwrap(); + + // Write phase: a fresh `FileCheckpointer` opened on the temp directory + // runs the graph to its interrupt point. + { + let cp: Arc> = Arc::new(FileCheckpointer::::new(dir.path())); + let graph = approve_graph().with_checkpointer(cp); + let paused = graph.run_with_thread("hitl", 10).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(paused.status.status, ExecutionStatus::Interrupted); + assert_eq!(paused.interrupts.len(), 1); + // `graph` (and the checkpointer it owns) is dropped at the end of + // this block, simulating the process exiting. + } + + // Read/resume phase: a brand-new `FileCheckpointer` instance, opened + // fresh against the same directory, resumes the run from disk. + let cp2: Arc> = Arc::new(FileCheckpointer::::new(dir.path())); + let graph2 = approve_graph().with_checkpointer(cp2); + let resumed = graph2 + .resume("hitl", Command::resume(json!({ "bump": 5 }))) + .await + .unwrap(); + assert!(!resumed.is_interrupted()); + assert_eq!(resumed.state, 15); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); +} + +#[tokio::test] +async fn file_backend_failure_then_restart_then_retry() { + let dir = tempfile::tempdir().unwrap(); + let attempts = Arc::new(AtomicUsize::new(0)); + + // Write phase: the node fails once and the run aborts, leaving a + // resumable failure checkpoint on disk. + { + let cp: Arc> = Arc::new(FileCheckpointer::::new(dir.path())); + let graph = flaky_graph(1, attempts.clone()).with_checkpointer(cp); + let err = graph.run_with_thread("net", 100).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 1); + // `graph` (and the checkpointer it owns) is dropped here, simulating + // the process exiting after the failure. + } + + // Read/retry phase: a fresh `FileCheckpointer` instance, opened against + // the same directory, retries the failed node from disk. The transient + // condition has cleared by the time this attempt runs. + let cp2: Arc> = Arc::new(FileCheckpointer::::new(dir.path())); + let graph2 = flaky_graph(1, attempts.clone()).with_checkpointer(cp2); + let resumed = graph2.retry("net").await.unwrap(); + assert_eq!(resumed.state, 101); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 2); +} + +// ── SqliteCheckpointer (feature = "sqlite") ────────────────────────────── + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn sqlite_backend_interrupt_then_restart_then_resume() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("checkpoints.db"); + + // Write phase: a fresh `SqliteCheckpointer` opened on the db file runs + // the graph to its interrupt point. + { + let cp: Arc> = + Arc::new(SqliteCheckpointer::::open(&db_path).unwrap()); + let graph = approve_graph().with_checkpointer(cp); + let paused = graph.run_with_thread("hitl", 10).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(paused.status.status, ExecutionStatus::Interrupted); + assert_eq!(paused.interrupts.len(), 1); + // `graph` (and the `Connection` it owns) is dropped at the end of + // this block, simulating the process exiting. + } + + // Read/resume phase: a brand-new `Connection`/`SqliteCheckpointer`, + // opened fresh against the same database file, resumes from disk. + let cp2: Arc> = + Arc::new(SqliteCheckpointer::::open(&db_path).unwrap()); + let graph2 = approve_graph().with_checkpointer(cp2); + let resumed = graph2 + .resume("hitl", Command::resume(json!({ "bump": 5 }))) + .await + .unwrap(); + assert!(!resumed.is_interrupted()); + assert_eq!(resumed.state, 15); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn sqlite_backend_failure_then_restart_then_retry() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("checkpoints.db"); + let attempts = Arc::new(AtomicUsize::new(0)); + + // Write phase: the node fails once and the run aborts, leaving a + // resumable failure checkpoint on disk. + { + let cp: Arc> = + Arc::new(SqliteCheckpointer::::open(&db_path).unwrap()); + let graph = flaky_graph(1, attempts.clone()).with_checkpointer(cp); + let err = graph.run_with_thread("net", 100).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 1); + // `graph` (and the `Connection` it owns) is dropped here, simulating + // the process exiting after the failure. + } + + // Read/retry phase: a fresh `Connection`/`SqliteCheckpointer`, opened + // against the same database file, retries the failed node from disk. + let cp2: Arc> = + Arc::new(SqliteCheckpointer::::open(&db_path).unwrap()); + let graph2 = flaky_graph(1, attempts.clone()).with_checkpointer(cp2); + let resumed = graph2.retry("net").await.unwrap(); + assert_eq!(resumed.state, 101); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 2); +} diff --git a/crates/tinyagents-graph/src/compiled/executor.rs b/crates/tinyagents-graph/src/compiled/executor.rs index 5c047850..e8560387 100644 --- a/crates/tinyagents-graph/src/compiled/executor.rs +++ b/crates/tinyagents-graph/src/compiled/executor.rs @@ -2,10 +2,113 @@ //! //! Split out of `compiled/mod.rs`; see that module's doc comment for the //! full executor design (superstep loop, concurrency, and resumable-failure -//! semantics). +//! semantics). The superstep loop itself is a thin wire-up over three +//! sibling modules: [`step`] runs a superstep's active node set and folds +//! its results ([`step::StepRunner`]), [`boundary`] applies the reducer, +//! routes, and persists a checkpoint at each of the three boundary shapes a +//! step can end at ([`boundary::StepBoundary`]), and [`resume`] loads a +//! checkpoint back into a fresh run. [`run_ctx::RunCtx`] carries the run +//! identity and bookkeeping all three share. use super::*; +use crate::compiled::boundary::StepBoundary; +use crate::compiled::run_ctx::{ResumeSeed, RunCtx}; +use crate::compiled::step::StepRunner; +use crate::thread_locks::ThreadLockMap; +use std::sync::OnceLock; + +/// Default TTL for the durable execution lease claimed in [`CompiledGraph::execute`] +/// (C3/R4). Renewal is not wired up (a single run is expected to complete, or at +/// least reach its next boundary, well inside this window); a lease that +/// outlives its owning process by more than this is reclaimable by the next +/// claimant. +const THREAD_LEASE_TTL: Duration = Duration::from_secs(300); + +/// Process-wide map of per-`(thread, namespace)` in-process execution locks +/// (C3/R4). Distinct from `delegation::run::thread_lock`'s map: that one +/// serializes delegation's own pre-`execute` checkpoint classification, this +/// one serializes the executor's run/resume/retry entry points themselves — +/// the gap the review's C3 finding describes (`executor.rs` took no lock of +/// its own). Keyed on `thread_id` *and* namespace so a parent run and a +/// subgraph run sharing a thread id never contend on each other's lock. +fn execution_lock_map() -> &'static ThreadLockMap { + static LOCKS: OnceLock = OnceLock::new(); + LOCKS.get_or_init(|| ThreadLockMap::new("graph executor per-thread run lock")) +} + +/// Builds the in-process lock map key for `thread_id` scoped to `namespace`. +/// `\u{1}` is not a legal thread-id or namespace-segment character in +/// practice and is used only as an internal separator, never persisted. +fn execution_lock_key(thread_id: &str, namespace: &[String]) -> String { + let mut key = thread_id.to_string(); + for segment in namespace { + key.push('\u{1}'); + key.push_str(segment); + } + key +} + +/// Everything a fresh or resumed run is seeded with, bundled so +/// [`CompiledGraph::execute`]/[`CompiledGraph::execute_run`] take one +/// parameter instead of positional state/thread/resume/barrier/binding +/// arguments. +pub(super) struct RunSeed { + pub(super) state: State, + pub(super) active: Vec, + pub(super) thread_id: Option, + /// Keyed by task id, falling back to node id (I1/R5), so a `Send` + /// fan-out of the same node can deliver each interrupted activation its + /// own resume value. + pub(super) resume_map: HashMap, + pub(super) barriers: HashMap>, + pub(super) parent: Option, + pub(super) binding: Option, + /// Resume-only seeding (step/node-visit continuation, carried-forward + /// mid-step completions) — see [`ResumeSeed`]. Left at its `Default` + /// (empty/zero) for a fresh run. + pub(super) resume_seed: ResumeSeed, + /// Optional per-run options (I4 part 2) — the cooperative cancellation + /// token and/or graceful-drain signal, if the caller opted in via + /// [`CompiledGraph::run_with_options`]/[`CompiledGraph::resume_with_options`]. + pub(super) options: RunOptions, + pub(super) _update: std::marker::PhantomData, +} + +impl RunSeed { + pub(super) fn fresh( + state: State, + active: Vec, + thread_id: Option, + ) -> Self { + Self { + state, + active, + thread_id, + resume_map: HashMap::new(), + barriers: HashMap::new(), + parent: None, + binding: None, + resume_seed: ResumeSeed::default(), + options: RunOptions::default(), + _update: std::marker::PhantomData, + } + } + + pub(super) fn with_binding( + mut self, + binding: crate::subagent_node::AgentInvocationBinding, + ) -> Self { + self.binding = Some(binding); + self + } + + pub(super) fn with_options(mut self, options: RunOptions) -> Self { + self.options = options; + self + } +} + impl CompiledGraph where State: Clone + Send + Sync + 'static, @@ -16,14 +119,74 @@ where /// Without a thread id no checkpoints are persisted even if a checkpointer /// is configured, since checkpoints are keyed by thread. pub async fn run(&self, state: State) -> Result> { - self.execute( + self.execute(RunSeed::fresh( state, vec![Activation::node(self.entry.clone())], None, - HashMap::new(), - HashMap::new(), - None, + )) + .await + } + + /// Runs the graph to completion (or to an interrupt/cancellation/drain) + /// without a thread, honoring `options` (I4 part 2): a cooperative + /// [`RunOptions::cancellation`] token checked at every superstep + /// boundary and raced against that step's in-flight node handlers, and + /// a graceful [`RunOptions::drain`] signal checked between supersteps + /// only (the step in flight always completes). + /// + /// Without a thread id, cancellation/drain still stop the run and record + /// a `Cancelled`/`Drained` status, but there is nothing to persist a + /// resumable checkpoint against (checkpoints are keyed by thread), + /// exactly like [`Self::run`]. + pub async fn run_with_options( + &self, + state: State, + options: RunOptions, + ) -> Result> { + self.execute( + RunSeed::fresh(state, vec![Activation::node(self.entry.clone())], None) + .with_options(options), + ) + .await + } + + /// Runs the graph under a thread id, honoring `options` (I4 part 2). + /// + /// This is the checkpointed counterpart to [`Self::run_with_options`]: a + /// cancellation or drain observed mid-run persists a resumable checkpoint + /// naming the still-pending activations, so the run can be continued + /// later with [`Self::resume`]/[`Self::retry`]. + pub async fn run_with_thread_options( + &self, + thread_id: impl Into, + state: State, + options: RunOptions, + ) -> Result> { + self.execute( + RunSeed::fresh( + state, + vec![Activation::node(self.entry.clone())], + Some(thread_id.into()), + ) + .with_options(options), + ) + .await + } + + /// Resumes a run from its latest checkpoint, honoring `options` (I4 part + /// 2) — see [`Self::run_with_thread_options`]. + pub async fn resume_with_options( + &self, + thread_id: impl Into, + command: Command, + options: RunOptions, + ) -> Result> { + self.resume_from_inner( + thread_id.into(), + ResumeTarget::Latest, + command, None, + options, ) .await } @@ -38,13 +201,8 @@ where binding: crate::subagent_node::AgentInvocationBinding, ) -> Result> { self.execute( - state, - vec![Activation::node(self.entry.clone())], - None, - HashMap::new(), - HashMap::new(), - None, - Some(binding), + RunSeed::fresh(state, vec![Activation::node(self.entry.clone())], None) + .with_binding(binding), ) .await } @@ -63,16 +221,7 @@ where inputs: impl IntoIterator, ) -> Result> { let active = self.initial_inputs(inputs)?; - self.execute( - state, - active, - None, - HashMap::new(), - HashMap::new(), - None, - None, - ) - .await + self.execute(RunSeed::fresh(state, active, None)).await } /// Runs the graph under a thread id, persisting checkpoints at every @@ -82,15 +231,11 @@ where thread_id: impl Into, state: State, ) -> Result> { - self.execute( + self.execute(RunSeed::fresh( state, vec![Activation::node(self.entry.clone())], Some(thread_id.into()), - HashMap::new(), - HashMap::new(), - None, - None, - ) + )) .await } @@ -102,13 +247,12 @@ where binding: crate::subagent_node::AgentInvocationBinding, ) -> Result> { self.execute( - state, - vec![Activation::node(self.entry.clone())], - Some(thread_id.into()), - HashMap::new(), - HashMap::new(), - None, - Some(binding), + RunSeed::fresh( + state, + vec![Activation::node(self.entry.clone())], + Some(thread_id.into()), + ) + .with_binding(binding), ) .await } @@ -123,16 +267,8 @@ where inputs: impl IntoIterator, ) -> Result> { let active = self.initial_inputs(inputs)?; - self.execute( - state, - active, - Some(thread_id.into()), - HashMap::new(), - HashMap::new(), - None, - None, - ) - .await + self.execute(RunSeed::fresh(state, active, Some(thread_id.into()))) + .await } /// Resumes an interrupted run from its latest checkpoint, re-running the @@ -221,8 +357,14 @@ where target: ResumeTarget, command: Command, ) -> Result> { - self.resume_from_inner(thread_id.into(), target, command, None) - .await + self.resume_from_inner( + thread_id.into(), + target, + command, + None, + RunOptions::default(), + ) + .await } /// Resumes a run from `target` with a host-bound recursive-agent binding. @@ -239,154 +381,12 @@ where command: Command, binding: crate::subagent_node::AgentInvocationBinding, ) -> Result> { - self.resume_from_inner(thread_id.into(), target, command, Some(binding)) - .await - } - - async fn resume_from_inner( - &self, - thread_id: ThreadId, - target: ResumeTarget, - command: Command, - binding: Option, - ) -> Result> { - let checkpointer = self - .checkpointer - .as_ref() - .ok_or_else(|| TinyAgentsError::Resume("no checkpointer configured".to_string()))?; - - let checkpoint_id = match &target { - ResumeTarget::Latest => None, - ResumeTarget::Checkpoint(id) => Some(id.as_str()), - }; - let checkpoint = checkpointer - .get_scoped(thread_id.as_str(), checkpoint_id, &self.namespace) - .await? - .ok_or_else(|| match &target { - ResumeTarget::Latest => { - TinyAgentsError::Resume(format!("no checkpoint found for thread `{thread_id}`")) - } - ResumeTarget::Checkpoint(id) => TinyAgentsError::Resume(format!( - "no checkpoint `{id}` found for thread `{thread_id}`" - )), - })?; - // Resume *loads* this checkpoint — it is a read, not a write — so emit a - // restore event, not `CheckpointSaved` (which would falsely inflate - // persisted-checkpoint counts and mislead durability observers). - self.emit(GraphEvent::CheckpointRestored { - checkpoint_id: CheckpointId::new(checkpoint.checkpoint_id.clone()), - }); - - // Prefer the persisted pending activations (which preserve each pending - // node's `Send` arg); fall back to the node-id projection for - // checkpoints written before that field existed. - let active: Vec = match &checkpoint.pending_activations { - Some(pending) if !pending.is_empty() => pending.iter().map(Activation::from).collect(), - _ => checkpoint - .next_nodes - .iter() - .cloned() - .map(Activation::node) - .collect(), - }; - if active.is_empty() { - return Err(TinyAgentsError::Resume( - "checkpoint has no pending nodes to resume".to_string(), - )); - } - - // Partial-failure guard. The boundary that produced this checkpoint - // recorded a completion marker per task that had already finished; a - // node named by *both* the pending set and that ledger has therefore - // already run, and re-running it would repeat its side effects. On a - // checkpoint the executor itself wrote the two sets are disjoint, so - // this is a no-op — it earns its keep on a checkpoint that was - // hand-built, time-travelled to, or edited through `update_state`, - // where `next_nodes` can legitimately disagree with what ran. - let completed_config = CheckpointConfig { - thread_id: thread_id.to_string(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: self.namespace.clone(), - }; - let recorded = checkpointer.get_writes(&completed_config).await?; - let done: HashSet = if recorded.is_empty() { - checkpoint - .pending_writes - .iter() - .map(|w| w.task_id.clone()) - .collect() - } else { - recorded.iter().map(|w| w.task_id.clone()).collect() - }; - let active: Vec = if done.is_empty() { - active - } else { - let filtered: Vec = active - .iter() - // A node name is not a task identity: a Send fan-out can have - // several live activations of one node. Legacy checkpoints - // have no persisted task id, so leave them runnable. - .filter(|a| a.task_id.is_empty() || !done.contains(&a.task_id)) - .cloned() - .collect(); - if filtered.is_empty() { - // Every pending node claims to have run. Trust the pending set - // rather than turning a resumable checkpoint into a hard error: - // a wrong re-run is recoverable, a stuck thread is not. - tinyagents_tracing::warn!( - "[graph:resume] every pending node of checkpoint `{}` has a completion \ - marker; resuming them anyway rather than stranding the thread", - checkpoint.checkpoint_id - ); - active - } else { - if filtered.len() != active.len() { - tinyagents_tracing::debug!( - "[graph:resume] checkpoint `{}`: skipping {} already-completed task(s)", - checkpoint.checkpoint_id, - active.len() - filtered.len() - ); - } - filtered - } - }; - - // 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 { - 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()); - } - } - } - - // Restore accumulated barrier arrivals so a join's precondition survives - // the interrupt/failure boundary this checkpoint recorded. - let initial_barriers = barriers_from_persisted(&checkpoint.barrier_arrivals); - // Chain the first post-resume boundary onto the checkpoint we loaded so - // the lineage spine stays connected across the resume. - let initial_parent = Some(checkpoint.checkpoint_id.clone()); - - self.execute( - checkpoint.state, - active, - Some(thread_id), - resume_map, - initial_barriers, - initial_parent, - binding, + self.resume_from_inner( + thread_id.into(), + target, + command, + Some(binding), + RunOptions::default(), ) .await } @@ -411,8 +411,8 @@ where }; active.push(Activation { node, - send_arg: input.payload, - task_id: String::new(), + send_arg: input.payload.map(Arc::new), + task_id: TaskId::from(String::new()), }); } if active.is_empty() { @@ -427,49 +427,64 @@ where /// Returns the configured checkpointer or a [`TinyAgentsError::Checkpoint`] /// when inspection is attempted on a graph without durability. - #[allow(clippy::too_many_arguments)] - async fn execute( + pub(super) async fn execute( &self, - state: State, - initial_active: Vec, - thread_id: Option, - resume_map: HashMap, - initial_barriers: HashMap>, - initial_parent: Option, - binding: Option, + seed: RunSeed, ) -> Result> { let run_id = tinyagents_harness::ids::new_run_id(); + + // C3/R4: hold this thread's execution lock for the run's whole + // lifetime, in-process first (cheap, always available) then a + // durable lease when a checkpointer is configured (cross-process). + // Held across the entire `execute_run` below — including checkpoint + // reads that would otherwise race a concurrent caller's — not just + // around the write, which is what closes the interleaving the C3 + // finding describes (two concurrent `run_with_thread`/`resume` on + // one thread id previously had nothing serializing them at this + // layer at all). + let _in_process_guard = if let Some(thread) = &seed.thread_id { + let key = execution_lock_key(thread.as_str(), &self.namespace); + Some(execution_lock_map().lock_for(&key).lock_owned().await) + } else { + None + }; + let lease_owner = + if let (Some(checkpointer), Some(thread)) = (&self.checkpointer, &seed.thread_id) { + match checkpointer + .try_claim(thread.as_str(), run_id.as_str(), THREAD_LEASE_TTL) + .await + { + Ok(true) => Some((checkpointer.clone(), thread.clone())), + Ok(false) => { + return Err(TinyAgentsError::Validation(format!( + "thread `{thread}` is leased by another run" + ))); + } + // A lease-claim I/O error must not silently degrade to + // running unprotected: propagate it rather than proceeding + // as if the claim had succeeded. + Err(err) => return Err(err), + } + } else { + None + }; + // When a durable journal is configured, run against a clone whose event // sink wraps every emitted event into a `GraphObservation` and appends // it (while still forwarding to any pre-existing live sink). The journal // sink carries this graph's checkpoint namespace so subgraph runs record // their nested path. Default (no journal) leaves `self` untouched. - if self.journal.is_some() { - let this = self.clone_with_journal_sink(&run_id, &thread_id); - this.execute_run( - run_id, - state, - initial_active, - thread_id, - resume_map, - initial_barriers, - initial_parent, - binding, - ) - .await + let result = if self.journal.is_some() { + let this = self.clone_with_journal_sink(&run_id, &seed.thread_id); + this.execute_run(run_id.clone(), seed).await } else { - self.execute_run( - run_id, - state, - initial_active, - thread_id, - resume_map, - initial_barriers, - initial_parent, - binding, - ) - .await + self.execute_run(run_id.clone(), seed).await + }; + + if let Some((checkpointer, thread)) = lease_owner { + let _ = checkpointer.release(thread.as_str(), run_id.as_str()).await; } + result } /// Builds a clone whose `event_sink` is a [`JournalGraphSink`] for `run_id`, @@ -494,1390 +509,257 @@ where this } - /// Best-effort status write; never aborts the run on a status-store error. - async fn save_status(&self, status: GraphRunStatus) { - if let Some(store) = &self.status_store { - let _ = store.put_status(status).await; - } - } - - #[allow(clippy::too_many_arguments)] + /// Drives one run's superstep loop to completion, an interrupt, or a + /// failure. + /// + /// Builds this run's [`RunCtx`] (identity, recursion stack, and the + /// accumulators the loop carries forward) and its [`StepRunner`], then + /// loops: check the recursion/deadline/visit-count guards, run the + /// active set's node handlers ([`StepRunner::run_sequential`] or + /// [`StepRunner::run_parallel`]), fold the results + /// ([`StepRunner::fold_step`]), apply updates through the reducer + /// ([`CompiledGraph::apply_updates`]), and dispatch to whichever + /// boundary the step ended at — failure + /// ([`CompiledGraph::handle_failure_boundary`]), interrupt + /// ([`CompiledGraph::handle_interrupt_boundary`]), or the normal boundary + /// ([`CompiledGraph::advance`], which returns the next active set). async fn execute_run( &self, run_id: RunId, - mut state: State, - initial_active: Vec, - thread_id: Option, - mut resume_map: HashMap, - initial_barriers: HashMap>, - initial_parent: Option, - binding: Option, + seed: RunSeed, ) -> Result> { - let started_at = SystemTime::now(); - let mut visited: Vec = Vec::new(); - let mut steps = 0usize; - let mut last_checkpoint: Option = None; - // On resume this is the loaded checkpoint's id, so the first boundary - // checkpoint after a resume chains onto pre-interrupt history rather - // than orphaning the lineage (which would stop `get_state_history` at - // the resume point and let `prune` delete the ancestors). - let mut parent_checkpoint: Option = initial_parent; + let RunSeed { + mut state, + active: initial_active, + thread_id, + resume_map, + barriers: initial_barriers, + parent: initial_parent, + binding, + resume_seed, + options, + .. + } = seed; - // Build this run's recursion stack from the inherited parent frames and - // push the frame for this graph call. A push that would exceed - // `max_depth` fails the run with a clear recursion error before any - // node executes. Graph-call depth (the stack) is tracked separately - // from node-loop visits (`node_visits`, below). - let mut recursion = - RecursionStack::with_frames(self.recursion_frames.clone(), self.recursion_policy); - // Run lineage: the root is the first inherited frame's run (the top of - // the recursion tree) or this run when top-level; the parent is the - // enclosing run, if any. - let root_run_id = self - .recursion_frames - .first() - .map(|f| f.run_id.clone()) - .unwrap_or_else(|| run_id.clone()); - let parent_run_id = self.recursion_frames.last().map(|f| f.run_id.clone()); - let this_frame = RecursionFrame { - graph_id: self.graph_id.clone(), - node_id: self.recursion_node.clone(), - run_id: run_id.clone(), - task_id: None, - namespace: self.namespace.clone(), - depth: recursion.depth(), - parent: parent_run_id.clone(), - }; - if let Err(err) = recursion.push(this_frame) { - self.emit(GraphEvent::RunStarted { - run_id: run_id.clone(), - }); - self.fail_run(&run_id, &thread_id, started_at, steps, &err, None) - .await; - return Err(err); - } - // Serialized once per run for embedding in every checkpoint's metadata. - let recursion_meta = - serde_json::to_value(recursion.frames()).unwrap_or(serde_json::Value::Null); - // The live frame stack handed to node contexts so a subgraph node can - // seed an embedded child with this run's recursion path, plus the - // per-run sink the node reports its spawned child run into. - let live_frames = recursion.frames().to_vec(); - let child_sink = ChildRunSink::new(); - // Accumulates every child run spawned across all supersteps for the - // final `GraphExecution::child_runs`. - let mut all_child_runs: Vec = Vec::new(); - // Per-node activation counts for `max_visits_per_node` enforcement. - let mut node_visits: HashMap = HashMap::new(); - let mut active = initial_active; - // Barrier/waiting-edge arrivals accumulate across supersteps: a waiting - // node only activates once every required predecessor has arrived. - // Seeded from the resumed checkpoint so a join's precondition survives - // an interrupt/failure boundary. - let mut barrier_arrivals: HashMap> = initial_barriers; - // Under `DurabilityMode::Async`, boundary checkpoint writes run on - // spawned background tasks tracked here. Failures are surfaced at the - // next durability boundary; every terminal path drains the tracker so - // the run result reflects persistence failures (see - // `AsyncCheckpointWrites`). - let mut async_writes = AsyncCheckpointWrites::default(); + let mut ctx = RunCtx::start( + self, + run_id, + thread_id, + resume_map, + initial_barriers, + initial_parent, + binding, + resume_seed, + options, + ) + .await?; + let runner = StepRunner { graph: self }; - self.emit(GraphEvent::RunStarted { - run_id: run_id.clone(), - }); - // Surface this run's recursion depth so observers can attribute nested - // runs without reconstructing the tree from logs. - self.emit(GraphEvent::RecursionDepthChanged { - depth: recursion.depth(), - }); // Record the run as live before the first superstep is scheduled. - let mut running = self.base_status(&run_id, &thread_id, started_at); - running.active_nodes = activation_nodes(&active); - self.save_status(running).await; + let mut running = ctx.base_status(); + running.active_nodes = activation_nodes(&initial_active); + ctx.save_status(running).await; + let mut active = initial_active; while !active.is_empty() { - // The effective step cap is the smaller of the builder's recursion - // limit and the policy's `max_total_steps`, so a policy never - // loosens an existing limit. Both surface a `RecursionLimit`. - let step_limit = self - .recursion_limit - .min(self.recursion_policy.max_total_steps); - if steps >= step_limit { - let err = TinyAgentsError::RecursionLimit(step_limit); - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - err, - ) - .await; + // I4 part 2: check cooperative cancellation at every superstep + // boundary, before starting a new step. `active` at this point is + // exactly what the next step would run, so a cancellation here + // schedules the whole set as pending (nothing of this step has + // executed yet). + if ctx.is_cancelled() { + return self.handle_cancel_boundary(&mut ctx, &active, &state).await; } - // Whole-run wall-clock deadline: stop *between* super-steps once the - // elapsed run time reaches it, leaving the last committed boundary - // checkpoint intact (unlike an external `tokio::time::timeout`, which - // aborts mid-super-step and cannot). The already-completed super-steps - // and their checkpoints are preserved; the run fails with `Timeout`. - if let Some(deadline) = self.run_deadline { - let elapsed = started_at.elapsed().unwrap_or_default(); - if elapsed >= deadline { - let err = TinyAgentsError::Timeout(format!( - "graph run exceeded its {deadline:?} deadline after {steps} super-step(s) \ - ({elapsed:?} elapsed)" - )); - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - err, - ) - .await; - } + // Graceful drain: checked *only* here, between supersteps, so a + // drain raised while a step was in flight lets that step finish + // and commit its boundary (above/`advance`), then stops before + // `active` — the next step's whole set — runs anything. + if ctx.is_drain_requested() { + return self.handle_drain_boundary(&mut ctx, &active, &state).await; } - // 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) { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - err, - ) - .await; - } - } - steps += 1; - // Assign identities before any branch runs. A failure checkpoint - // carries these identities with its pending activations, letting a - // later resume skip only the completed fan-out task. - for (index, activation) in active.iter_mut().enumerate() { - if activation.task_id.is_empty() { - activation.task_id = format!("{steps}:{index}:{}", activation.node); - } - } - self.emit(GraphEvent::StepStarted { - step: steps, - active: activation_nodes(&active), - }); - let run_result = if self.parallel && active.len() > 1 { - self.run_active_parallel( - &active, - &state, - &run_id, - &thread_id, - steps, - &mut resume_map, - &mut visited, - &root_run_id, - &live_frames, - &child_sink, - &binding, - ) - .await - } else { - self.run_active_sequential( - &active, - &state, - &run_id, - &thread_id, - steps, - &mut resume_map, - &mut visited, - &root_run_id, - &live_frames, - &child_sink, - &binding, - ) - .await + let step = match self.begin_step(&mut ctx, &mut active).await { + Ok(step) => step, + Err(err) => return self.fail_and_return(&mut ctx, err).await, }; - let StepRun { - updates, - goto_map, - interrupt, - failure, - } = match run_result { - Ok(step_run) => step_run, - Err(err) => { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - err, - ) - .await; - } + + let step_run = match self + .run_step_with_cancel(&runner, &mut ctx, &active, &state, step) + .await + { + Ok(Some(step_run)) => step_run, + // Cancelled while this step's handlers were in flight: none + // of them are trusted to have applied (the step's own future + // was raced and abandoned, not awaited to completion), so the + // whole active set is still pending. + Ok(None) => return self.handle_cancel_boundary(&mut ctx, &active, &state).await, + Err(err) => return self.fail_and_return(&mut ctx, err).await, }; // Apply collected updates through the reducer at the boundary. A // reducer error here must still fail the run (not just unwind // leaving it `Running`). - for update in updates { - state = match self.reducer.apply(state, update) { - Ok(state) => state, - Err(err) => { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - err, - ) - .await; - } - }; - } - - // Collect any child runs spawned by subgraph nodes this step. They - // are embedded into this boundary's checkpoint metadata (keyed by - // node) and accumulated onto the final `GraphExecution`. - let step_child_runs = child_sink.drain(); - all_child_runs.extend(step_child_runs.iter().cloned()); - let child_runs_meta = - serde_json::to_value(&step_child_runs).unwrap_or(serde_json::Value::Null); - - // Node-handler failure (survived any node-retry policy): the updates - // of the branches that completed before it are already folded into - // `state` above, so persist a resumable failure-boundary checkpoint - // scheduling the failed node (and the not-yet-run tail) for a later - // `resume`/`retry`, record a `Failed` status carrying the error and - // that checkpoint, and abort. Without a checkpointer/thread the - // checkpoint is a no-op and the run aborts exactly as before. - if let Some(fail) = failure { - let StepFailure { - failed_index, - error, - } = fail; - let failed_node = active[failed_index].node.clone(); - // Schedule the successors of the branches that completed before - // the failure (they succeeded; their routing must not be lost) - // followed by the failed branch and the not-yet-run tail, which - // re-run on resume with their `Send` args preserved. - let successors = match self.route_completed( - &active[..failed_index], - &goto_map, - &state, - &mut barrier_arrivals, - ) { - Ok(successors) => successors, - Err(route_err) => { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - route_err, - ) - .await; - } - }; - let mut pending = successors; - pending.extend(active[failed_index..].iter().cloned()); - // Settle any in-flight Async background writes before the - // failure-boundary persist so earlier boundaries are durable - // when the run aborts. Like the persist error below, a - // background write error must not replace the original node - // error, so it is intentionally dropped here. - let _ = async_writes.drain().await; - // A failure-boundary persist error must not replace the original - // node error: keep reporting the node error and just drop the - // resumable checkpoint reference. - let checkpoint_id = self - .persist_failure_checkpoint( - &thread_id, - &run_id, - &state, - &pending, - &active[..failed_index], - &barrier_arrivals, - parent_checkpoint.clone(), - steps, - &failed_node, - &error, - &recursion_meta, - &child_runs_meta, - ) - .await - .unwrap_or(None); - self.fail_run( - &run_id, - &thread_id, - started_at, - steps, - &error, - checkpoint_id, - ) - .await; - return Err(error); - } - - // Interrupt: persist a checkpoint whose pending activations are the - // successors of the branches that completed before the interrupt - // (their routing must survive) followed by the not-yet-completed - // members of this step (interrupted node first). Each pending branch - // keeps its `Send` arg; accumulated barrier arrivals are persisted - // too. Then return control to the caller. - if let Some((index, emitted)) = interrupt { - if let Err(err) = self.require_interrupt_durability(&thread_id) { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - err, - ) - .await; - } - let successors = match self.route_completed( - &active[..index], - &goto_map, - &state, - &mut barrier_arrivals, - ) { - Ok(successors) => successors, - Err(route_err) => { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - route_err, - ) - .await; - } - }; - let mut pending = successors; - pending.extend(active[index..].iter().cloned()); - let pending_nodes = activation_nodes(&pending); - let interrupt_id = InterruptId::new(emitted.id.clone()); - // An interrupt hands control back to the caller expecting a - // fully durable pause point: settle any in-flight Async - // background writes first, failing the run if one was lost - // (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, - &mut async_writes, - err, - ) - .await; - } - let checkpoint_id = match self - .persist_checkpoint( - &thread_id, - &run_id, - &state, - &pending, - &active[..index], - vec![emitted.clone()], - std::slice::from_ref(&active[index].node), - &barrier_arrivals, - parent_checkpoint.clone(), - steps, - "loop", - &recursion_meta, - &child_runs_meta, - ) - .await - { - Ok(id) => id, - Err(persist_err) => { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - persist_err, - ) - .await; - } - }; - - let mut status = self.base_status(&run_id, &thread_id, started_at); - status.status = ExecutionStatus::Interrupted; - status.current_step = steps; - status.active_nodes = pending_nodes; - status.pending_interrupts = vec![interrupt_id]; - status.checkpoint_id = checkpoint_id.clone(); - self.save_status(status.clone()).await; - - return Ok(GraphExecution { - state, - run_id: run_id.clone(), - graph_id: self.graph_id.clone(), - root_run_id: root_run_id.clone(), - parent_run_id: parent_run_id.clone(), - child_runs: all_child_runs, - visited, - steps, - interrupts: vec![emitted], - status, - checkpoint_id, - }); - } - - // Select the next active set from commands or static/conditional - // edges, evaluated against the freshly-committed state. Barrier - // arrivals accumulate into `barrier_arrivals` (persisted below). - let next = match self.route_completed(&active, &goto_map, &state, &mut barrier_arrivals) - { - Ok(next) => next, - Err(route_err) => { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - route_err, - ) - .await; - } + state = match self.apply_updates(state, step_run.updates) { + Ok(state) => state, + Err(err) => return self.fail_and_return(&mut ctx, err).await, }; - // Persist a boundary checkpoint. Under `Exit` durability only the - // terminal boundary (the step that empties the active set) is - // written; `Sync`/`Async` persist every boundary. `Async` hands - // non-terminal writes to background tasks instead of awaiting them - // inline. - let persist_now = match self.durability { - DurabilityMode::Exit => next.is_empty(), - DurabilityMode::Sync | DurabilityMode::Async => true, + // Child runs spawned by subgraph nodes this step are embedded + // into this boundary's checkpoint metadata (keyed by node) and + // accumulated onto the final `GraphExecution`. + let child_runs_meta = ctx.take_step_child_runs(); + let sb = StepBoundary { + active: &active, + completed: &step_run.completed, + stalled: &step_run.stalled, + goto_map: &step_run.goto_map, + child_runs_meta: &child_runs_meta, + task_writes: &step_run.task_writes, + step, }; - // Async durability: surface any background write failure recorded - // since the previous boundary. The run fails at the first - // durability boundary that observes the loss rather than silently - // continuing with a hole in its lineage. - if let Some(err) = async_writes.take_failure().await { + + // Node-handler failure (survived any node-retry policy) or an + // interrupt: both are terminal for this run, persisting a + // resumable boundary checkpoint before returning. + if let Some(fail) = step_run.failure { return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - err, - ) + .handle_failure_boundary(&mut ctx, sb, &state, fail) .await; } - let terminal = next.is_empty(); - let checkpoint_id = if persist_now { - let persisted = if matches!(self.durability, DurabilityMode::Async) && !terminal { - self.persist_checkpoint_nonblocking( - &mut async_writes, - &thread_id, - &run_id, - &state, - &next, - &active, - &barrier_arrivals, - parent_checkpoint.clone(), - steps, - &recursion_meta, - &child_runs_meta, - ) - .await - } else { - // Terminal boundary: drain every in-flight background - // write first (the "final await at run end"), so a lost - // Async checkpoint fails the run instead of being - // swallowed. The final checkpoint itself is then written - // 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, - &mut async_writes, - err, - ) - .await; - } - self.persist_checkpoint( - &thread_id, - &run_id, - &state, - &next, - &active, - Vec::new(), - &[], - &barrier_arrivals, - parent_checkpoint.clone(), - steps, - "loop", - &recursion_meta, - &child_runs_meta, - ) - .await - }; - match persisted { - Ok(id) => id, - Err(persist_err) => { - return self - .fail_and_return( - &run_id, - &thread_id, - started_at, - steps, - &mut async_writes, - persist_err, - ) - .await; - } - } - } else { - None - }; - if let Some(id) = &checkpoint_id { - last_checkpoint = Some(id.clone()); - parent_checkpoint = Some(id.to_string()); + if !step_run.interrupted.is_empty() { + return self + .handle_interrupt_boundary(&mut ctx, sb, state, step_run.interrupted) + .await; } - self.emit(GraphEvent::StepCompleted { step: steps }); - active = next; + active = match self.advance(&mut ctx, sb, &state).await { + Ok(next) => next, + Err(err) => return self.fail_and_return(&mut ctx, err).await, + }; } - let mut status = self.base_status(&run_id, &thread_id, started_at); - status.status = ExecutionStatus::Completed; - status.current_step = steps; - status.checkpoint_id = last_checkpoint.clone(); - status.ended_at = Some(SystemTime::now()); - self.save_status(status.clone()).await; - self.emit(GraphEvent::RunCompleted { - run_id: run_id.clone(), - steps, - }); - - Ok(GraphExecution { - state, - run_id: run_id.clone(), - graph_id: self.graph_id.clone(), - root_run_id, - parent_run_id, - child_runs: all_child_runs, - visited, - steps, - interrupts: Vec::new(), - status, - checkpoint_id: last_checkpoint, - }) - } - - /// Emits a [`GraphEvent::RunFailed`] and records a terminal `Failed` status - /// for a run that aborted with `err`. - /// - /// `checkpoint_id` is the resumable failure-boundary checkpoint when the run - /// left one (a node-handler failure on a checkpointed thread), or `None` for - /// a structural/non-resumable abort. When present it is recorded on the - /// status so an observer can locate the checkpoint to `resume`/`retry` from. - async fn fail_run( - &self, - run_id: &RunId, - thread_id: &Option, - started_at: SystemTime, - steps: usize, - err: &TinyAgentsError, - checkpoint_id: Option, - ) { - self.emit(GraphEvent::RunFailed { - run_id: run_id.clone(), - error: err.to_string(), - }); - let mut status = self.base_status(run_id, thread_id, started_at); - status.status = ExecutionStatus::Failed; - status.current_step = steps; - status.ended_at = Some(SystemTime::now()); - status.error = Some(err.to_string()); - status.checkpoint_id = checkpoint_id; - self.save_status(status).await; - } - - /// Records a terminal `Failed` status for `err` (via [`Self::fail_run`]) and - /// returns it as `Err`. - /// - /// Used at the step boundary so an error raised *after* the node runners — - /// 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) + Ok(self.finish_run(&mut ctx, state).await) } - /// Persists a resumable failure-boundary checkpoint for a node-handler - /// failure that survived the node-retry policy. + /// Runs one superstep, racing it against this run's cooperative + /// cancellation token (I4 part 2) so a long-running node's handlers + /// cannot indefinitely block a cancellation request once requested. /// - /// Mirrors the interrupt boundary: `next_nodes` schedules the failed node - /// (and any not-yet-run members of the step) so `resume`/`retry` re-runs - /// exactly what did not complete, while `completed_tasks` records the - /// branches that already succeeded (their updates are folded into `state` - /// before this is called). The rendered error and failed node id are stamped - /// into the checkpoint metadata for diagnosis. A no-op returning `None` when - /// no checkpointer/thread is configured — the run then aborts without a - /// resumable checkpoint, exactly as before this policy existed. - #[allow(clippy::too_many_arguments)] - async fn persist_failure_checkpoint( + /// Returns `Ok(Some(step_run))` when the step completed first, + /// `Ok(None)` when the token was already cancelled or was cancelled + /// while the step's handlers were still in flight (the step's own future + /// is then dropped, abandoning it — see [`super::run_ctx::RunDropGuard`]'s + /// doc for what that does and does not guarantee for any checkpoint + /// write the abandoned step's handlers had already triggered), and + /// `Err` for an ordinary step failure. + async fn run_step_with_cancel( &self, - thread_id: &Option, - run_id: &RunId, + runner: &StepRunner<'_, State, Update>, + ctx: &mut RunCtx<'_, State, Update>, + active: &[Activation], state: &State, - pending: &[Activation], - completed_tasks: &[Activation], - barrier_arrivals: &HashMap>, - parent: Option, step: usize, - failed_node: &NodeId, - error: &TinyAgentsError, - recursion: &serde_json::Value, - child_runs: &serde_json::Value, - ) -> Result> { - let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { - return Ok(None); - }; - let checkpoint = Checkpoint { - thread_id: thread.to_string(), - checkpoint_id: next_checkpoint_id(), - run_id: Some(run_id.to_string()), - parent_checkpoint_id: parent, - namespace: self.namespace.clone(), - state: state.clone(), - next_nodes: activation_nodes(pending), - completed_tasks: activation_nodes(completed_tasks), - pending_writes: Self::completion_writes(completed_tasks, step), - interrupts: Vec::new(), - pending_activations: Some(pending.iter().map(PendingActivation::from).collect()), - barrier_arrivals: barriers_to_persisted(barrier_arrivals), - metadata: serde_json::json!({ - "source": "loop", - "step": step, - "recursion": recursion, - "child_runs": child_runs, - "failed_node": failed_node.as_str(), - "error": error.to_string(), - }), - }; - let writes = checkpoint.pending_writes.clone(); - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), + ) -> Result>> { + let Some(token) = ctx.cancellation.clone() else { + return runner.run_step(ctx, active, state, step).await.map(Some); }; - let id = checkpointer.put(checkpoint).await?; - // Also record the ledger through the write protocol, so backends that - // implement it can answer "did this task run?" without loading the - // whole state payload. - checkpointer.put_writes(&config, &writes).await?; - self.emit(GraphEvent::CheckpointSaved { - checkpoint_id: id.clone(), - }); - Ok(Some(id)) - } - - /// Builds the per-task [`NodeContext`] for `node_id` at the given branch. - /// - /// `fork` carries the branch identity in a concurrent step (`None` in - /// sequential mode or single-node steps). The resume value for the node is - /// consumed from `resume_map`. - #[allow(clippy::too_many_arguments)] - fn node_context( - &self, - node_id: &NodeId, - run_id: &RunId, - thread_id: &Option, - step: usize, - resume_map: &mut HashMap, - fork: Option, - send_arg: Option, - root_run_id: &RunId, - frames: &[RecursionFrame], - child_runs: &ChildRunSink, - binding: &Option, - ) -> NodeContext { - NodeContext { - graph_id: self.graph_id.clone(), - node_id: node_id.clone(), - run_id: run_id.clone(), - thread_id: thread_id.clone(), - step, - resume: resume_map.remove(node_id), - fork, - send_arg, - root_run_id: Some(root_run_id.clone()), - recursion_frames: frames.to_vec(), - child_runs: Some(child_runs.clone()), - agent_binding: binding.clone(), + if token.is_cancelled() { + return Ok(None); } - } - - /// Wraps a node future in the configured per-node timeout (if any), mapping - /// an elapsed deadline onto [`TinyAgentsError::Timeout`]. - async fn run_node_future( - &self, - node_id: &NodeId, - fut: NodeFuture, - ) -> Result> { - match self.node_timeout { - Some(timeout) => match tokio::time::timeout(timeout, fut).await { - Ok(result) => result, - Err(_) => Err(TinyAgentsError::Timeout(format!( - "node `{node_id}` exceeded its {timeout:?} timeout" - ))), - }, - None => fut.await, + tokio::select! { + biased; + _ = token.cancelled() => Ok(None), + result = runner.run_step(ctx, active, state, step) => result.map(Some), } } - /// Runs one node handler under the graph's node-retry policy. - /// - /// Builds a fresh handler future (and re-clones the context) for each - /// attempt, so a retried node re-runs from its start — matching the durable - /// execution model, where a node is never suspended mid-flight. On a - /// [retryable][tinyagents_harness::retry::is_retryable] error, when a - /// [`RetryPolicy`](tinyagents_harness::retry::RetryPolicy) is configured and - /// permits another attempt, it emits - /// [`GraphEvent::NodeRetryScheduled`], sleeps the (opt-in) backoff, and - /// retries. Non-retryable errors, absence of a policy, or an exhausted - /// attempt budget return the error unchanged. The per-node timeout still - /// bounds every individual attempt via [`Self::run_node_future`]. - async fn run_node_with_retry( + /// Checks the recursion-limit, wall-clock-deadline, and per-node + /// visit-count guards for the next superstep, then advances `ctx.steps`, + /// assigns any missing task ids in `active` (a failure checkpoint + /// carries these with its pending activations, letting a later resume + /// skip only the completed fan-out task), and emits `StepStarted`. + /// Returns the step number on success. + async fn begin_step( &self, - node_id: &NodeId, - handler: &Arc>, - state: &State, - ctx: NodeContext, - step: usize, - ) -> Result> { - let mut attempt = 0usize; - loop { - let fut = handler(state.clone(), ctx.clone()); - match self.run_node_future(node_id, fut).await { - Ok(result) => return Ok(result), - Err(error) => { - let retry = self - .node_retry - .as_ref() - .filter(|policy| policy.should_retry(attempt) && is_retryable(&error)); - let Some(policy) = retry else { - return Err(error); - }; - attempt += 1; - self.emit(GraphEvent::NodeRetryScheduled { - node: node_id.clone(), - step, - attempt, - }); - policy.sleep_backoff(attempt).await; - } - } + ctx: &mut RunCtx<'_, State, Update>, + active: &mut [Activation], + ) -> Result { + // The effective step cap is the smaller of the builder's recursion + // limit and the policy's `max_total_steps`, so a policy never + // loosens an existing limit. Both surface a `RecursionLimit`. + let step_limit = self + .recursion_limit + .min(self.recursion_policy.max_total_steps); + if ctx.steps >= step_limit { + return Err(TinyAgentsError::RecursionLimit(step_limit)); } - } - - /// Folds a single successful branch result into the step accumulators. - /// - /// Pushes the node to `visited`, records updates/goto, emits the matching - /// events, and returns the interrupt (with its branch index) when the branch - /// paused. Shared by the sequential and parallel run paths so both fold - /// results identically; only the *running* of handlers differs. - #[allow(clippy::too_many_arguments)] - fn fold_result( - &self, - index: usize, - node_id: &NodeId, - step: usize, - result: NodeResult, - updates: &mut Vec, - goto_map: &mut HashMap>, - visited: &mut Vec, - ) -> Option<(usize, Interrupt)> { - visited.push(node_id.clone()); - match result { - NodeResult::Update(update) => { - updates.push(update); - self.emit(GraphEvent::StateUpdated { - node: node_id.clone(), - step, - }); - } - NodeResult::Command(command) => { - if let Some(update) = command.update { - updates.push(update); - self.emit(GraphEvent::StateUpdated { - node: node_id.clone(), - step, - }); - } - if !command.goto.is_empty() { - goto_map.insert(index, command.goto); - } - } - NodeResult::Interrupt(emitted) => { - self.emit(GraphEvent::InterruptEmitted { - interrupt: emitted.clone(), - }); - return Some((index, emitted)); - } - } - self.emit(GraphEvent::NodeCompleted { - node: node_id.clone(), - step, - }); - None - } - - /// Runs the active node set one node at a time (default behavior). - /// - /// Short-circuits on the first error (run aborts) or interrupt (later nodes - /// in the step are not started), exactly preserving milestone-1 semantics. - #[allow(clippy::too_many_arguments)] - async fn run_active_sequential( - &self, - active: &[Activation], - state: &State, - run_id: &RunId, - thread_id: &Option, - step: usize, - resume_map: &mut HashMap, - visited: &mut Vec, - root_run_id: &RunId, - frames: &[RecursionFrame], - child_runs: &ChildRunSink, - binding: &Option, - ) -> Result> { - let mut updates: Vec = Vec::new(); - let mut goto_map: HashMap> = HashMap::new(); - let mut interrupt: Option<(usize, Interrupt)> = None; - let mut failure: Option = None; - - for (index, activation) in active.iter().enumerate() { - let node_id = &activation.node; - let node = self - .nodes - .get(node_id) - .ok_or_else(|| TinyAgentsError::MissingNode(node_id.to_string()))?; - - self.emit(GraphEvent::TaskScheduled { - node: node_id.clone(), - step, - }); - self.emit(GraphEvent::NodeStarted { - node: node_id.clone(), - step, - }); - - let ctx = self.node_context( - node_id, - run_id, - thread_id, - step, - resume_map, - None, - activation.send_arg.clone(), - root_run_id, - frames, - child_runs, - binding, - ); - let result = match self - .run_node_with_retry(node_id, &node.handler, state, ctx, step) - .await - { - Ok(result) => result, - Err(error) => { - self.emit(GraphEvent::NodeFailed { - node: node_id.clone(), - step, - error: error.to_string(), - }); - // Preserve the progress of the branches that already ran: - // the executor records them as completed and schedules their - // successors plus this node and the not-yet-run tail for a - // resumable retry. - failure = Some(StepFailure { - failed_index: index, - error, - }); - break; - } - }; - - if let Some(found) = self.fold_result( - index, - node_id, - step, - result, - &mut updates, - &mut goto_map, - visited, - ) { - interrupt = Some(found); - break; + // Whole-run wall-clock deadline: stop *between* super-steps once the + // elapsed run time reaches it, leaving the last committed boundary + // checkpoint intact (unlike an external `tokio::time::timeout`, which + // aborts mid-super-step and cannot). The already-completed super-steps + // and their checkpoints are preserved; the run fails with `Timeout`. + if let Some(deadline) = self.run_deadline { + let elapsed = ctx.started_instant.elapsed(); + if elapsed >= deadline { + return Err(TinyAgentsError::Timeout(format!( + "graph run exceeded its {deadline:?} deadline after {} super-step(s) \ + ({elapsed:?} elapsed)", + ctx.steps + ))); } } - - Ok(StepRun { - updates, - goto_map, - interrupt, - failure, - }) - } - - /// Runs the active node set concurrently (opt-in via `with_parallel`). - /// - /// Each branch executes on its own cloned `State` snapshot and a distinct - /// [`ForkId`], optionally with the [`Send`] argument that scheduled it. With - /// no `max_concurrency` bound every branch starts before any is awaited and - /// all are driven via [`futures::future::join_all`]; with a bound the active - /// set is run in chunks of at most that many futures, so at most that many - /// node handlers are in flight at once. Results are folded in active-set - /// index order — the reducer is the join/fan-in — so the merged state is - /// reproducible regardless of completion order. The lowest-index branch that - /// errors or interrupts is the step's terminal outcome; lower-index - /// successful branches still contribute their updates. - #[allow(clippy::too_many_arguments)] - async fn run_active_parallel( - &self, - active: &[Activation], - state: &State, - run_id: &RunId, - thread_id: &Option, - step: usize, - resume_map: &mut HashMap, - visited: &mut Vec, - root_run_id: &RunId, - frames: &[RecursionFrame], - child_runs: &ChildRunSink, - binding: &Option, - ) -> Result> { - // Build one forked context + future per branch. Node lookup and resume - // consumption happen up front so the futures borrow nothing mutable; each - // branch drives its handler through the node-retry policy (which also - // applies the per-node timeout), so a transient failure in one branch is - // retried without disturbing its siblings. - let mut futures = Vec::with_capacity(active.len()); - for (index, activation) in active.iter().enumerate() { - let node_id = &activation.node; - let node = self - .nodes - .get(node_id) - .ok_or_else(|| TinyAgentsError::MissingNode(node_id.to_string()))?; - - self.emit(GraphEvent::TaskScheduled { - node: node_id.clone(), - step, - }); - self.emit(GraphEvent::NodeStarted { - node: node_id.clone(), - step, - }); - - self.emit(GraphEvent::ContextForked { - node: node_id.clone(), - fork: index, - step, - }); - let fork = Some(ForkId::new(index, node_id.clone())); - let ctx = self.node_context( - node_id, - run_id, - thread_id, - step, - resume_map, - fork, - activation.send_arg.clone(), - root_run_id, - frames, - child_runs, - binding, - ); - let handler = node.handler.clone(); - let owned_node = node_id.clone(); - // Box each branch future behind a concrete `Send` bound. This keeps - // the `buffer_unordered` rolling window below (used for a - // `max_concurrency` bound) from requiring a higher-ranked `Send` - // proof over the borrowed recursion frames, which the compiler - // cannot discharge for the bare `async` blocks. - let fut: std::pin::Pin< - Box>> + Send + '_>, - > = Box::pin(async move { - self.run_node_with_retry(&owned_node, &handler, state, ctx, step) - .await - }); - futures.push(fut); + // Node-loop recursion: enforce `max_visits_per_node` per activation. + for activation in active.iter() { + ctx.recursion + .record_node_visit(&mut ctx.node_visits, &activation.node)?; } - - // Drive branches to completion, bounding in-flight count when configured. - // With a bound, keep a rolling window of `limit` branches in flight - // instead of fixed `join_all` chunks. A chunked join runs each chunk to - // completion before starting the next, so a single slow branch - // head-of-line blocks the whole chunk; the rolling window starts a new - // branch as soon as *any* in-flight one finishes. `select_all` reports - // which pending future completed; a parallel index Vec maps it back to - // the branch's active-set position, so results are re-ordered into - // deterministic order for the fold below. - let results = match self.max_concurrency { - Some(limit) if limit < futures.len() => { - let total = futures.len(); - let mut slots: Vec>>> = - (0..total).map(|_| None).collect(); - let mut source = futures.into_iter().enumerate(); - let mut running = Vec::with_capacity(limit); - let mut running_index = Vec::with_capacity(limit); - for (index, fut) in source.by_ref().take(limit) { - running.push(fut); - running_index.push(index); - } - while !running.is_empty() { - let (result, completed, rest) = futures::future::select_all(running).await; - let index = running_index.remove(completed); - slots[index] = Some(result); - running = rest; - if let Some((index, fut)) = source.next() { - running.push(fut); - running_index.push(index); - } - } - slots - .into_iter() - .map(|slot| slot.expect("every branch produced a result")) - .collect::>() - } - _ => futures::future::join_all(futures).await, - }; - - // Fold in deterministic active-set index order. - let mut updates: Vec = Vec::new(); - let mut goto_map: HashMap> = HashMap::new(); - let mut interrupt: Option<(usize, Interrupt)> = None; - let mut failure: Option = None; - - for (index, (activation, result)) in active.iter().zip(results).enumerate() { - let node_id = &activation.node; - let result = match result { - Ok(result) => result, - Err(error) => { - self.emit(GraphEvent::NodeFailed { - node: node_id.clone(), - step, - error: error.to_string(), - }); - // The lowest-index failing branch is terminal: fold the - // lower-index successes (already applied above) and schedule - // their successors plus this branch and the rest for a - // resumable retry. - failure = Some(StepFailure { - failed_index: index, - error, - }); - break; - } - }; - - if let Some(found) = self.fold_result( - index, - node_id, - step, - result, - &mut updates, - &mut goto_map, - visited, - ) { - interrupt = Some(found); - break; + ctx.steps += 1; + for (index, activation) in active.iter_mut().enumerate() { + if activation.task_id.as_str().is_empty() { + activation.task_id = + TaskId::from(format!("{}:{}:{}", ctx.steps, index, activation.node)); } } - - Ok(StepRun { - updates, - goto_map, - interrupt, - failure, - }) + ctx.emit(GraphEvent::StepStarted { + step: ctx.steps, + active: activation_nodes(active), + }); + Ok(ctx.steps) } - /// Routes a set of completed activations into their successor activations. - /// - /// Honors per-activation command `goto` (keyed by active-set index), static - /// and conditional edges, barrier gating (a waiting node is held until every - /// required predecessor has arrived, accumulating into `barrier_arrivals` - /// across supersteps), and per-node dedup — while preserving each `Send` - /// packet's per-invocation argument. Emits a - /// [`GraphEvent::RouteSelected`] per selected edge. - /// - /// Shared by the normal step boundary (routes the whole active set) and the - /// interrupt/failure boundaries (route just the branches that completed - /// before the pause, so their successors are still scheduled on resume). - #[allow(clippy::too_many_arguments)] - async fn persist_checkpoint( + /// Builds the terminal [`GraphExecution`] for a run that emptied its + /// active set without interrupting or failing: records a `Completed` + /// status and emits `RunCompleted`. + async fn finish_run( &self, - thread_id: &Option, - run_id: &RunId, - state: &State, - pending: &[Activation], - completed_tasks: &[Activation], - interrupts: Vec, - interrupted: &[NodeId], - barrier_arrivals: &HashMap>, - parent: Option, - step: usize, - source: &str, - recursion: &serde_json::Value, - child_runs: &serde_json::Value, - ) -> Result> { - let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { - return Ok(None); - }; - let checkpoint = self.build_loop_checkpoint( - thread, - run_id, - state, - pending, - completed_tasks, - interrupts, - interrupted, - barrier_arrivals, - parent, - step, - source, - recursion, - child_runs, - ); - let writes = checkpoint.pending_writes.clone(); - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; - let id = checkpointer.put(checkpoint).await?; - checkpointer.put_writes(&config, &writes).await?; - self.emit(GraphEvent::CheckpointSaved { - checkpoint_id: id.clone(), + ctx: &mut RunCtx<'_, State, Update>, + state: State, + ) -> GraphExecution { + ctx.disarm_drop_guard(); + let mut status = ctx.base_status(); + status.status = ExecutionStatus::Completed; + status.current_step = ctx.steps; + status.checkpoint_id = ctx.last_checkpoint.clone(); + status.ended_at = Some(SystemTime::now()); + ctx.save_status(status.clone()).await; + ctx.emit(GraphEvent::RunCompleted { + run_id: ctx.run_id.clone(), + steps: ctx.steps, }); - Ok(Some(id)) - } - /// Persists a boundary checkpoint without blocking the superstep loop - /// ([`DurabilityMode::Async`]). - /// - /// The checkpoint id is minted up front and returned immediately so the - /// loop keeps chaining lineage onto it, while the actual `put` (and the - /// [`GraphEvent::CheckpointSaved`] emitted on its success) runs on a - /// spawned background task tracked in `writes`. - /// - /// # Failure semantics - /// - /// A background write error is never dropped: it is recorded in `writes` - /// and surfaced by the executor at the next durability boundary, or at the - /// latest when the run drains all in-flight writes at its terminal / - /// interrupt boundary — so the run result reflects persistence failures. - /// Because the `CheckpointSaved` event is emitted from the background - /// task, its ordering relative to subsequent step events is not - /// deterministic under `Async` durability. - /// - /// Outside a tokio runtime there is nothing to spawn onto, so the write - /// happens inline — degrading to [`DurabilityMode::Sync`] behavior. - #[allow(clippy::too_many_arguments)] - async fn persist_checkpoint_nonblocking( - &self, - writes: &mut AsyncCheckpointWrites, - thread_id: &Option, - run_id: &RunId, - state: &State, - pending: &[Activation], - completed_tasks: &[Activation], - barrier_arrivals: &HashMap>, - parent: Option, - step: usize, - recursion: &serde_json::Value, - child_runs: &serde_json::Value, - ) -> Result> { - let (Some(checkpointer), Some(thread)) = (&self.checkpointer, thread_id) else { - return Ok(None); - }; - let checkpoint = self.build_loop_checkpoint( - thread, - run_id, + GraphExecution { state, - pending, - completed_tasks, - Vec::new(), - &[], - barrier_arrivals, - parent, - step, - "loop", - recursion, - child_runs, - ); - let id = CheckpointId::new(checkpoint.checkpoint_id.clone()); - - match tokio::runtime::Handle::try_current() { - Ok(handle) => { - let checkpointer = Arc::clone(checkpointer); - let sink = self.event_sink.clone(); - writes.spawn_ordered(&handle, async move { - let id = checkpointer.put(checkpoint).await?; - if let Some(sink) = sink { - sink.emit(GraphEvent::CheckpointSaved { - checkpoint_id: id.clone(), - }); - } - Ok(id) - }); - Ok(Some(id)) - } - Err(_) => { - let id = checkpointer.put(checkpoint).await?; - self.emit(GraphEvent::CheckpointSaved { - checkpoint_id: id.clone(), - }); - Ok(Some(id)) - } - } - } - - /// Records completion markers for the tasks that finished in the step a - /// boundary checkpoint closes. - /// - /// A graph's `Update` carries no `Serialize` bound, so the executor cannot - /// persist *what* a task wrote — but it does not need to: the applied value - /// is already durable in the checkpoint's `state`. What was missing was the - /// other half, the per-task record of *that* it ran, which is what lets a - /// resume distinguish "already done" from "not yet started". See - /// [`PendingWrite`](crate::checkpoint::PendingWrite)'s docs for why - /// that distinction is the whole point of - /// the ledger. - /// - /// The task id is persisted on the activation itself, so a resume can - /// match a marker to one fan-out task rather than every task with its node. - fn completion_writes( - completed_tasks: &[Activation], - _step: usize, - ) -> Vec { - completed_tasks - .iter() - .map(|activation| { - crate::checkpoint::PendingWrite::completion_marker( - activation.node.clone(), - activation.task_id.clone(), - ) - }) - .collect() - } - - /// Builds the loop-boundary [`Checkpoint`] record shared by the sync and - /// async persist paths, minting a fresh checkpoint id. - #[allow(clippy::too_many_arguments)] - fn build_loop_checkpoint( - &self, - thread: &ThreadId, - run_id: &RunId, - state: &State, - pending: &[Activation], - completed_tasks: &[Activation], - interrupts: Vec, - interrupted: &[NodeId], - barrier_arrivals: &HashMap>, - parent: Option, - step: usize, - source: &str, - 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(), - run_id: Some(run_id.to_string()), - parent_checkpoint_id: parent, - namespace: self.namespace.clone(), - state: state.clone(), - next_nodes: activation_nodes(pending), - completed_tasks: activation_nodes(completed_tasks), - pending_writes: Self::completion_writes(completed_tasks, step), - pending_activations: Some(pending.iter().map(PendingActivation::from).collect()), - barrier_arrivals: barriers_to_persisted(barrier_arrivals), - interrupts, - metadata, + run_id: ctx.run_id.clone(), + graph_id: self.graph_id.clone(), + root_run_id: ctx.root_run_id.clone(), + parent_run_id: ctx.parent_run_id.clone(), + child_runs: std::mem::take(&mut ctx.all_child_runs), + visited: std::mem::take(&mut ctx.visited), + steps: ctx.steps, + interrupts: Vec::new(), + status, + checkpoint_id: ctx.last_checkpoint.clone(), + drained: false, } } - - fn base_status( - &self, - run_id: &RunId, - thread_id: &Option, - started_at: SystemTime, - ) -> GraphRunStatus { - let mut status = GraphRunStatus::new( - run_id.clone(), - self.graph_id.clone(), - ExecutionStatus::Running, - ); - status.thread_id = thread_id.clone(); - status.checkpoint_namespace = self.namespace.clone(); - status.started_at = started_at; - status.updated_at = SystemTime::now(); - status - } } diff --git a/crates/tinyagents-graph/src/compiled/interrupt_selectors_test.rs b/crates/tinyagents-graph/src/compiled/interrupt_selectors_test.rs new file mode 100644 index 00000000..51aad28d --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/interrupt_selectors_test.rs @@ -0,0 +1,505 @@ +//! Executor-level interrupt injection (`GraphBuilder::interrupt_before` / +//! `interrupt_after`, with `mark_interrupt` as the `before` alias) and +//! fail-closed `Interrupt::response_schema` validation on resume. +//! +//! Every test counts handler invocations through a shared `AtomicUsize`: +//! the load-bearing contract of both selectors is that the paused node's +//! handler runs **exactly once** across the pause and the resume — never +//! zero times (a lost write) and never twice (a repeated side effect). + +use super::*; +use crate::builder::{GraphBuilder, NodeContext}; +use crate::checkpoint::{Checkpointer, InMemoryCheckpointer}; +use crate::command::{Command, Interrupt, NodeResult}; +use serde_json::json; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use tinyagents_harness::ids::ExecutionStatus; + +fn memory() -> Arc> { + Arc::new(InMemoryCheckpointer::::new()) +} + +/// `a (+1) -> b (+10, counted, adds any resume `bump`) -> c (+100)`, with +/// the interrupt selectors applied by `configure`. +fn chain( + b_runs: Arc, + configure: impl FnOnce(GraphBuilder) -> GraphBuilder, +) -> CompiledGraph { + let builder = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", move |s, ctx: NodeContext| { + let b_runs = b_runs.clone(); + async move { + b_runs.fetch_add(1, AtomicOrdering::SeqCst); + let bump = ctx + .resume + .as_ref() + .and_then(|v| v.get("bump")) + .and_then(|v| v.as_i64()) + .unwrap_or(0) as i32; + Ok(NodeResult::Update(s + 10 + bump)) + } + }) + .add_node("c", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 100)) + }) + .add_sequence(["a", "b", "c"]) + .set_entry("a") + .set_finish("c"); + configure(builder) + .compile() + .unwrap() + .with_checkpointer(memory()) +} + +fn phase(interrupt: &Interrupt) -> &str { + interrupt.payload["phase"].as_str().unwrap_or("") +} + +#[tokio::test] +async fn interrupt_before_pauses_without_running_the_handler_and_resume_runs_it_once() { + let b_runs = Arc::new(AtomicUsize::new(0)); + let graph = chain(b_runs.clone(), |b| b.interrupt_before(["b"])); + + let paused = graph.run_with_thread("before", 0).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(paused.status.status, ExecutionStatus::Interrupted); + assert_eq!(paused.interrupts.len(), 1); + assert_eq!(paused.interrupts[0].node.as_str(), "b"); + assert_eq!(phase(&paused.interrupts[0]), "before"); + assert!( + paused.interrupts[0].task_id.is_some(), + "stamped with its task id" + ); + // `a` committed; `b` never ran. + assert_eq!(paused.state, 1); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 0); + + // The pause is a checkpoint carrying the interrupt. + let history = graph.get_state_history("before", None).await.unwrap(); + assert!(history[0].metadata.has_interrupts); + assert_eq!(history[0].pending_interrupts.len(), 1); + assert_eq!( + history[0] + .next_nodes + .iter() + .map(|n| n.as_str()) + .collect::>(), + vec!["b"] + ); + assert_eq!(history[0].values, 1); + + // Resume runs `b` normally (with the resume value) and finishes. + let resumed = graph + .resume("before", Command::resume(json!({ "bump": 5 }))) + .await + .unwrap(); + assert!(!resumed.is_interrupted()); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(resumed.state, 1 + 10 + 5 + 100); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); +} + +#[tokio::test] +async fn mark_interrupt_is_an_alias_for_interrupt_before() { + let b_runs = Arc::new(AtomicUsize::new(0)); + let graph = chain(b_runs.clone(), |b| b.mark_interrupt("b")); + + // Export marker still set... + let topology = graph.topology(); + let b = topology.nodes.iter().find(|n| n.id == "b").unwrap(); + assert!(b.interrupt); + + // ...and the runtime pause is real. + let paused = graph.run_with_thread("alias", 0).await.unwrap(); + assert_eq!(phase(&paused.interrupts[0]), "before"); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 0); + let resumed = graph.retry("alias").await.unwrap(); + assert_eq!(resumed.state, 111); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); +} + +#[tokio::test] +async fn interrupt_after_runs_the_handler_once_and_holds_its_update_until_resume() { + let b_runs = Arc::new(AtomicUsize::new(0)); + let graph = chain(b_runs.clone(), |b| b.interrupt_after(["b"])); + + let paused = graph.run_with_thread("after", 0).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(paused.interrupts.len(), 1); + assert_eq!(paused.interrupts[0].node.as_str(), "b"); + assert_eq!(phase(&paused.interrupts[0]), "after"); + // The handler ran exactly once, but its `+10` is not yet committed. + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); + assert_eq!(paused.state, 1); + + // The checkpoint also holds the pre-update state, the interrupt, and + // `b` as the pending task with its deferred result in the ledger. + let snapshot = graph.get_state("after", None).await.unwrap().unwrap(); + assert_eq!(snapshot.values, 1); + assert!(snapshot.metadata.has_interrupts); + assert_eq!( + snapshot + .next_nodes + .iter() + .map(|n| n.as_str()) + .collect::>(), + vec!["b"] + ); + let tuple = graph + .checkpointer + .as_ref() + .unwrap() + .get_tuple(snapshot.config.clone()) + .await + .unwrap() + .unwrap(); + let deferred: Vec<_> = tuple + .pending_writes + .iter() + .filter(|w| w.is_interrupt_after()) + .collect(); + assert_eq!(deferred.len(), 1); + assert_eq!(deferred[0].payload["update"], json!(11)); + let history = graph.get_state_history("after", None).await.unwrap(); + assert!(history[0].metadata.has_interrupts); + + // Resume replays the stored result: no second handler run, update applied. + let resumed = graph.retry("after").await.unwrap(); + assert!(!resumed.is_interrupted()); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(resumed.state, 111); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); +} + +#[tokio::test] +async fn interrupt_after_replays_the_deferred_command_goto_on_resume() { + let b_runs = Arc::new(AtomicUsize::new(0)); + let runs = b_runs.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", move |s, _c: NodeContext| { + let runs = runs.clone(); + async move { + runs.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Command( + Command::update(s + 10).with_goto(["d"]), + )) + } + }) + .add_node("c", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 100)) + }) + .add_node("d", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1000)) + }) + .set_entry("a") + .add_edge("a", "b") + .mark_command_routing("b") + .set_finish("c") + .set_finish("d") + .interrupt_after(["b"]) + .compile() + .unwrap() + .with_checkpointer(memory()); + + let paused = graph.run_with_thread("goto", 0).await.unwrap(); + assert_eq!(phase(&paused.interrupts[0]), "after"); + assert_eq!(paused.state, 1); + + let resumed = graph.retry("goto").await.unwrap(); + // `b`'s +10 applied, and its explicit `goto d` honoured (not `c`). + assert_eq!(resumed.state, 1 + 10 + 1000); + assert_eq!( + resumed + .visited + .iter() + .map(|n| n.as_str()) + .collect::>(), + vec!["b", "d"] + ); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); +} + +#[tokio::test] +async fn interrupt_before_and_after_on_one_node_pause_twice_and_run_it_once() { + let b_runs = Arc::new(AtomicUsize::new(0)); + let graph = chain(b_runs.clone(), |b| { + b.interrupt_before(["b"]).interrupt_after(["b"]) + }); + + let first = graph.run_with_thread("both", 0).await.unwrap(); + assert_eq!(phase(&first.interrupts[0]), "before"); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 0); + + let second = graph.retry("both").await.unwrap(); + assert!(second.is_interrupted()); + assert_eq!(phase(&second.interrupts[0]), "after"); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); + assert_eq!(second.state, 1); + + let done = graph.retry("both").await.unwrap(); + assert_eq!(done.status.status, ExecutionStatus::Completed); + assert_eq!(done.state, 111); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); +} + +#[tokio::test] +async fn interrupt_after_in_a_parallel_step_defers_every_branch() { + // Two fan-out branches, both `interrupt_after`: both handlers run once, + // both updates are held, both are replayed on resume through the + // additive reducer. + let runs = Arc::new(AtomicUsize::new(0)); + let r1 = runs.clone(); + let r2 = runs.clone(); + let graph = GraphBuilder::::new() + .set_reducer(crate::reducer::ClosureStateReducer::new( + |s: i32, u: i32| Ok(s + u), + )) + .with_parallel(true) + .add_node("fan", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update(0)) + }) + .add_node("x", move |_s, _c: NodeContext| { + let r1 = r1.clone(); + async move { + r1.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(10)) + } + }) + .add_node("y", move |_s, _c: NodeContext| { + let r2 = r2.clone(); + async move { + r2.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(20)) + } + }) + .set_entry("fan") + .add_edge("fan", "x") + .add_edge("fan", "y") + .set_finish("x") + .set_finish("y") + .interrupt_after(["x", "y"]) + .compile() + .unwrap() + .with_checkpointer(memory()); + + let paused = graph.run_with_thread("par", 0).await.unwrap(); + assert_eq!(paused.interrupts.len(), 2); + assert!(paused.interrupts.iter().all(|i| phase(i) == "after")); + assert_eq!(paused.state, 0); + assert_eq!(runs.load(AtomicOrdering::SeqCst), 2); + + let resumed = graph.retry("par").await.unwrap(); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(resumed.state, 30); + assert_eq!(runs.load(AtomicOrdering::SeqCst), 2); +} + +#[test] +fn interrupt_selectors_must_name_real_nodes() { + let err = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .set_entry("a") + .set_finish("a") + .interrupt_before(["ghost"]) + .compile() + .unwrap_err(); + assert!( + matches!(&err, TinyAgentsError::MissingNode(n) if n == "ghost"), + "got {err:?}" + ); +} + +// ── Interrupt::response_schema ─────────────────────────────────────────── + +/// The schema every test below resumes against. +fn approval_schema() -> serde_json::Value { + json!({ + "type": "object", + "required": ["approved"], + "properties": { "approved": { "type": "boolean" } } + }) +} + +/// `approve` pauses with a schema-bearing interrupt; once resumed it commits +/// `+1` when approved, `-1` otherwise. +fn approval_graph() -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node("approve", |s, ctx: NodeContext| async move { + match ctx.resume { + Some(value) => { + let approved = value["approved"].as_bool().unwrap_or(false); + Ok(NodeResult::Update(if approved { s + 1 } else { s - 1 })) + } + None => Ok(NodeResult::Interrupt( + Interrupt::new("approve", json!({ "ask": "approve?" })) + .with_response_schema(approval_schema()), + )), + } + }) + .set_entry("approve") + .set_finish("approve") + .compile() + .unwrap() + .with_checkpointer(memory()) +} + +#[tokio::test] +async fn response_schema_accepts_a_conforming_resume_value() { + let graph = approval_graph(); + let paused = graph.run_with_thread("ok", 10).await.unwrap(); + assert_eq!( + paused.interrupts[0].response_schema, + Some(approval_schema()) + ); + let resumed = graph + .resume("ok", Command::resume(json!({ "approved": true }))) + .await + .unwrap(); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(resumed.state, 11); +} + +#[tokio::test] +async fn response_schema_rejects_bad_values_before_touching_the_checkpoint() { + let graph = approval_graph(); + let paused = graph.run_with_thread("bad", 10).await.unwrap(); + let before = graph.get_state_history("bad", None).await.unwrap(); + let latest_id = before[0].config.checkpoint_id.clone(); + + for bad in [json!({ "approved": "yes" }), json!({}), json!("approve")] { + let err = graph + .resume("bad", Command::resume(bad.clone())) + .await + .unwrap_err(); + assert!( + matches!(&err, TinyAgentsError::Validation(msg) if msg.contains("response_schema")), + "{bad}: got {err:?}" + ); + } + + // Fail-closed: no new checkpoint, same latest id, same pending interrupt. + let after = graph.get_state_history("bad", None).await.unwrap(); + assert_eq!(after.len(), before.len()); + assert_eq!(after[0].config.checkpoint_id, latest_id); + assert_eq!(after[0].values, 10); + assert_eq!(after[0].pending_interrupts, paused.interrupts); + assert!(after[0].metadata.has_interrupts); + + // A conforming value still resumes the untouched checkpoint. + let resumed = graph + .resume("bad", Command::resume(json!({ "approved": false }))) + .await + .unwrap(); + assert_eq!(resumed.state, 9); + assert_eq!( + graph.get_state_history("bad", None).await.unwrap().len(), + before.len() + 1 + ); +} + +#[tokio::test] +async fn response_schema_validates_per_task_resume_values() { + let graph = approval_graph(); + let paused = graph.run_with_thread("by-task", 0).await.unwrap(); + let task = paused.interrupts[0].task_id.clone().unwrap(); + + let err = graph + .resume( + "by-task", + Command::resume_tasks([(task.clone(), json!({ "approved": 1 }))]), + ) + .await + .unwrap_err(); + assert!(matches!(err, TinyAgentsError::Validation(_)), "got {err:?}"); + + let resumed = graph + .resume( + "by-task", + Command::resume_tasks([(task, json!({ "approved": true }))]), + ) + .await + .unwrap(); + assert_eq!(resumed.state, 1); +} + +#[tokio::test] +async fn retry_without_a_resume_value_skips_schema_validation() { + // A bare `retry` delivers no value, so there is nothing to validate; + // the node simply pauses again. + let graph = approval_graph(); + graph.run_with_thread("retry", 0).await.unwrap(); + let again = graph.retry("retry").await.unwrap(); + assert!(again.is_interrupted()); +} + +#[tokio::test] +async fn interrupt_before_ack_survives_a_later_node_emitted_pause() { + // `b` is `interrupt_before` *and* emits its own interrupt on its first + // real run: the executor's `before` pause must not fire again after the + // node's own pause is resumed. + let b_runs = Arc::new(AtomicUsize::new(0)); + let runs = b_runs.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("b", move |s, ctx: NodeContext| { + let runs = runs.clone(); + async move { + runs.fetch_add(1, AtomicOrdering::SeqCst); + match ctx.resume { + Some(_) => Ok(NodeResult::Update(s + 10)), + None => Ok(NodeResult::Interrupt(Interrupt::new("b", json!({})))), + } + } + }) + .set_entry("b") + .set_finish("b") + .interrupt_before(["b"]) + .compile() + .unwrap() + .with_checkpointer(memory()); + + let first = graph.run_with_thread("ack", 0).await.unwrap(); + assert_eq!(phase(&first.interrupts[0]), "before"); + let second = graph.retry("ack").await.unwrap(); + assert!(second.is_interrupted()); + assert_eq!(phase(&second.interrupts[0]), "", "the node's own interrupt"); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); + let done = graph + .resume("ack", Command::resume(json!({ "ok": true }))) + .await + .unwrap(); + assert_eq!(done.status.status, ExecutionStatus::Completed); + assert_eq!(done.state, 10); + assert_eq!( + b_runs.load(AtomicOrdering::SeqCst), + 2, + "re-run once for the node's own pause" + ); +} + +#[tokio::test] +async fn update_state_carries_a_deferred_interrupt_after_result_forward() { + // Inspect -> `update_state` (not attributed to the paused node) -> + // resume: the manual write must not lose `b`'s held-back result. + let b_runs = Arc::new(AtomicUsize::new(0)); + let graph = chain(b_runs.clone(), |b| b.interrupt_after(["b"])); + graph.run_with_thread("edit", 0).await.unwrap(); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); + + // Overwrite reducer: the committed state becomes 1000, `b` still pending. + graph.update_state("edit", 1000, None).await.unwrap(); + let resumed = graph.retry("edit").await.unwrap(); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + // `b`'s deferred update was computed against state 1 (-> 11), replayed + // verbatim (overwrite), then `c` adds 100. + assert_eq!(resumed.state, 111); + assert_eq!(b_runs.load(AtomicOrdering::SeqCst), 1); +} diff --git a/crates/tinyagents-graph/src/compiled/mod.rs b/crates/tinyagents-graph/src/compiled/mod.rs index fa4cfb9f..66d475a4 100644 --- a/crates/tinyagents-graph/src/compiled/mod.rs +++ b/crates/tinyagents-graph/src/compiled/mod.rs @@ -29,15 +29,41 @@ //! - All active branches in a parallel step start before any is awaited, and all //! are driven to completion (`join_all`) before the step boundary runs. //! - Branch results are then folded in active-set index order. The reducer is -//! the fan-in / join: lower-index branches' updates are applied first. +//! the fan-in / join: every branch's update is applied — **every** branch +//! that completed this step, not only the ones with a lower index than a +//! sibling that errored or interrupted (see the C1 fix in +//! `docs/runtime-comparison/code-review-graph.md`: a completed higher-index +//! branch is no longer discarded and silently re-run on resume). //! - The *lowest-index* branch that errors or interrupts is the step's terminal -//! outcome. Updates produced by lower-index successful branches are still -//! applied/persisted; an error persists a resumable failure boundary (see -//! below) and aborts, an interrupt persists a checkpoint whose pending nodes -//! are that branch and every later active node. +//! outcome; any other branch that also errored/interrupted is still recorded +//! (not dropped, not mistaken for completed) but does not become *the* +//! surfaced failure/interrupt. Every branch that completed is folded into +//! committed state, but its *routing* is deferred rather than resolved +//! immediately (the C2 fix): resolving a completed branch's successor before +//! its stalled siblings are known would let that successor observe a state +//! missing whatever those siblings eventually write, which is exactly the +//! ordering bug an uninterrupted run never has. The deferred branches' +//! node ids are persisted (`Checkpoint::completed_tasks`) and carried +//! forward across however many times this step interrupts/fails and gets +//! resumed/retried; only once every branch of the step has completed does +//! the executor route them all together, in one call, against one +//! committed state — see [`boundary::CompiledGraph::advance`]'s +//! `carried_completed` handling. One caveat: a deferred branch's routing +//! is re-resolved via static/conditional edges only (an explicit +//! `Command::goto` it returned is not itself persisted across the +//! boundary — see `StepRun::completed`). //! - Because branches run on cloned snapshots and never share mutable state, //! concurrency is data-race free; the reducer alone resolves conflicting //! writes (deterministically, by index). +//! - Sequential steps have their own cousin of C1: [`step::StepRunner::run_sequential`] +//! stops invoking further branches at the first error/interrupt, so those +//! not-yet-started siblings never appear in that step's raw results at +//! all. [`step::StepRunner::fold_step`] now folds them into `stalled` +//! anyway (by original active-set index), so a failure/interrupt boundary +//! records them as pending tasks (`Checkpoint::tasks`) alongside the +//! branch that stopped the step, instead of silently dropping them from +//! the checkpoint — a resumed/retried sequential run reaches the same +//! final state as an uninterrupted one. //! //! ## Network resilience and resumable failures //! @@ -64,22 +90,29 @@ //! [`CompiledGraph::update_state`] before resuming. Without a checkpointer the //! run aborts immediately, exactly as before. +mod boundary; mod executor; +mod resume; mod routing; +mod run_ctx; mod state_api; +mod step; mod types; -pub use types::{CompiledGraph, GraphExecution, GraphInput, ResumeTarget, StateSnapshot}; +pub use types::{ + CompiledGraph, DrainHandle, DrainSignal, GraphExecution, GraphInput, ResumeTarget, RunOptions, + StateSnapshot, +}; pub(crate) use types::AsyncCheckpointWrites; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, SystemTime}; use crate::builder::{ - BarrierRelief, Branch, BuilderNode, END, ForkId, NodeContext, NodeFuture, NodeHandler, - NodeMeta, START, + BarrierRelief, Branch, BuilderNode, END, ForkId, IdleClock, NodeContext, NodeFuture, + NodeHandler, NodeMeta, NodePolicy, START, UpdateCodec, }; use crate::checkpoint::{ BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointTuple, Checkpointer, DurabilityMode, @@ -89,10 +122,10 @@ use crate::command::{Command, Interrupt, NodeResult, RouteTarget}; use crate::recursion::{ChildRun, ChildRunSink, RecursionFrame, RecursionPolicy, RecursionStack}; use crate::reducer::StateReducer; use crate::status::GraphRunStatus; -use crate::stream::{GraphEvent, GraphEventSink}; +use crate::stream::{GraphEvent, GraphEventEnvelope, GraphEventSink}; use crate::{Result, TinyAgentsError}; use tinyagents_harness::ids::{ - CheckpointId, ExecutionStatus, GraphId, InterruptId, NodeId, RunId, ThreadId, + CheckpointId, ExecutionStatus, GraphId, InterruptId, NodeId, RunId, TaskId, ThreadId, }; use tinyagents_harness::retry::is_retryable; @@ -122,7 +155,10 @@ fn snapshot_from_tuple(tuple: CheckpointTuple) -> StateSnapshot = checkpoint.tasks.iter().map(|t| t.node.clone()).collect(); StateSnapshot { values: checkpoint.state, tasks: next_nodes.clone(), @@ -134,28 +170,6 @@ fn snapshot_from_tuple(tuple: CheckpointTuple) -> StateSnapshot { - /// Branch updates in deterministic active-set index order. - updates: Vec, - /// Explicit routing (plain `goto` nodes and/or [`Send`] packets) keyed by the - /// producing branch's active-set index. - /// - /// Keyed by index rather than node id so repeated [`Send`] activations of - /// the *same* node within a step (map-reduce fanout) each keep their own - /// [`Command::goto`] — a node-keyed map would let a later activation's - /// command clobber an earlier one's routing. - goto_map: HashMap>, - /// The lowest-index branch interrupt, if any (its active-set index + value). - interrupt: Option<(usize, Interrupt)>, - /// A node-handler failure that survived the node-retry policy, if any. When - /// set, `updates` still carries the updates of the branches that completed - /// *before* the failing branch, so the executor can fold that partial - /// progress into committed state and persist a resumable failure boundary. - failure: Option, -} - /// A node-handler failure captured by a runner so the executor can persist a /// resumable failure-boundary checkpoint instead of discarding partial progress. struct StepFailure { @@ -179,8 +193,10 @@ struct StepFailure { #[derive(Clone)] struct Activation { node: NodeId, - send_arg: Option, - task_id: String, + /// `Arc`-wrapped (M2): a `Send` fan-out of the same node, and every + /// retry attempt of one activation, share this allocation. + send_arg: Option>, + task_id: TaskId, } impl Activation { @@ -188,7 +204,7 @@ impl Activation { Self { node, send_arg: None, - task_id: String::new(), + task_id: TaskId::from(String::new()), } } } @@ -231,6 +247,20 @@ fn barriers_from_persisted(persisted: &[BarrierArrivals]) -> HashMap) -> serde_json::Value { + node_visits + .iter() + .map(|(node, count)| (node.to_string(), serde_json::json!(count))) + .collect::>() + .into() +} + /// Maps an [`Activation`] slice to its node ids (for events, status, and /// checkpoint records, which are node-keyed). fn activation_nodes(active: &[Activation]) -> Vec { @@ -283,7 +313,7 @@ impl CompiledGraph { graph_id: GraphId, name: Option, nodes: HashMap>, - edges: HashMap, + edges: HashMap>, branches: HashMap>, command_nodes: HashSet, waiting: HashMap>, @@ -323,9 +353,56 @@ impl CompiledGraph { run_deadline: None, durability: crate::checkpoint::DurabilityMode::default(), node_retry: None, + sequence: Arc::new(std::sync::atomic::AtomicU64::new(0)), + node_policies: Arc::new(HashMap::new()), + node_defaults: None, + task_cache: None, + cached_nodes: Arc::new(HashMap::new()), + interrupt_before: Arc::new(HashSet::new()), + interrupt_after: Arc::new(HashSet::new()), + update_codec: None, } } + /// Installs the `interrupt_before`/`interrupt_after` node selectors and + /// the `Update` codec the latter persists deferred results with (called + /// from `GraphBuilder::compile`). + pub(crate) fn with_interrupt_selectors( + mut self, + interrupt_before: HashSet, + interrupt_after: HashSet, + update_codec: Option>, + ) -> Self { + self.interrupt_before = Arc::new(interrupt_before); + self.interrupt_after = Arc::new(interrupt_after); + self.update_codec = update_codec; + self + } + + /// Installs the per-node policies and graph-wide default policy the + /// builder accumulated (called from `GraphBuilder::compile`). + pub(crate) fn with_node_policies( + mut self, + node_policies: HashMap>, + node_defaults: Option>, + ) -> Self { + self.node_policies = Arc::new(node_policies); + self.node_defaults = node_defaults.map(Arc::new); + self + } + + /// Resolves the effective [`NodePolicy`] for `node`: per-node field → + /// `set_node_defaults` field → legacy graph-wide `node_retry` / + /// `node_timeout`. + pub(crate) fn effective_policy(&self, node: &NodeId) -> NodePolicy { + NodePolicy::resolve( + self.node_policies.get(node), + self.node_defaults.as_deref(), + self.node_retry.as_ref(), + self.node_timeout, + ) + } + /// The graph id. pub fn graph_id(&self) -> &GraphId { &self.graph_id @@ -396,6 +473,55 @@ impl CompiledGraph { self } + /// Attaches the [`TaskCache`](crate::cache::TaskCache) backend used by + /// any node configured through [`Self::with_cached_node`]. + /// + /// Without a task cache, [`NodeCachePolicy`](crate::NodeCachePolicy) + /// entries installed by `with_cached_node` are inert: the executor never + /// looks anything up or stores anything, and every node runs exactly as + /// it would with no cache configured at all. + pub fn with_task_cache(mut self, cache: Arc) -> Self { + self.task_cache = Some(cache); + self + } + + /// Opts `node` into result caching under `policy`. + /// + /// A cache hit (an unexpired entry under `policy.key`'s computed hash) + /// skips the node's handler entirely and replays the stored `Update`, + /// emitting [`GraphEvent::TaskCompleted`](crate::stream::GraphEvent::TaskCompleted) + /// with `cached: true` in place of the handler's normal + /// `NodeStarted`/`NodeCompleted` pair. A miss runs the handler as usual + /// and, on success, stores the resulting `Update` and emits + /// `TaskCompleted { cached: false, .. }`. + /// + /// Actually persisting a cached value needs `Update: Serialize + + /// DeserializeOwned`; that bound lives on this method rather than on + /// [`CompiledGraph`] itself, so a graph with no cached nodes at all never + /// has to satisfy it. This method has no effect until a backend is also + /// installed via [`Self::with_task_cache`]. + pub fn with_cached_node( + mut self, + node: impl Into, + policy: crate::builder::NodeCachePolicy, + ) -> Self + where + Update: serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static, + { + let mut cached_nodes = (*self.cached_nodes).clone(); + cached_nodes.insert( + node.into(), + crate::cache::CachedNode { + key: policy.key, + ttl: policy.ttl, + encode: Arc::new(|update: &Update| serde_json::to_value(update)), + decode: Arc::new(|value: serde_json::Value| serde_json::from_value(value)), + }, + ); + self.cached_nodes = Arc::new(cached_nodes); + self + } + /// Bounds the whole run by a wall-clock `deadline`, checked at every /// super-step boundary. /// @@ -479,22 +605,99 @@ impl CompiledGraph { self } - fn emit(&self, event: GraphEvent) { - if let Some(sink) = &self.event_sink { - // Durable sinks persist asynchronously off the executor thread. On a - // terminal run event, flush so a caller that reads the journal right - // after the run returns sees a complete log. - let terminal = matches!( - event, - GraphEvent::RunCompleted { .. } | GraphEvent::RunFailed { .. } - ); - sink.emit(event); - if terminal { - sink.flush(); - } + /// Wraps `event` in a [`crate::stream::GraphEventEnvelope`] stamped with + /// `run_id`, this graph instance's checkpoint namespace, and the next + /// value of its [`CompiledGraph::sequence`] counter, then delivers it to + /// the configured sink (a no-op without one). The envelope's `task_id` is + /// `None` — use [`Self::emit_task`] for an event that belongs to one + /// task/activation. + fn emit(&self, run_id: &RunId, event: GraphEvent) { + self.emit_task(run_id, None, event); + } + + /// [`Self::emit`] for an event that belongs to one task/activation: + /// stamps the envelope's `task_id` with `task_id` so a multi-task stream + /// (e.g. a parallel superstep) can be correlated back to the activation + /// that produced each event. + pub(crate) fn emit_task(&self, run_id: &RunId, task_id: Option<&TaskId>, event: GraphEvent) { + let Some(sink) = &self.event_sink else { + return; + }; + // Durable sinks persist asynchronously off the executor thread. On a + // terminal run event, flush so a caller that reads the journal right + // after the run returns sees a complete log. + let terminal = matches!( + event, + GraphEvent::RunCompleted { .. } + | GraphEvent::RunFailed { .. } + | GraphEvent::RunCancelled { .. } + | GraphEvent::RunDrained { .. } + ); + let envelope = self.envelope_task(run_id, task_id, event); + sink.emit(envelope); + if terminal { + sink.flush(); + } + } + + /// Builds a [`crate::stream::GraphEventEnvelope`] for `event` without + /// delivering it — the shared stamping logic behind [`Self::emit`] and + /// the async-checkpoint-write path in `boundary.rs`, which must build an + /// envelope on the calling thread (to keep `seq` ordered) before handing + /// the write off to a spawned task. `task_id` is `None`; use + /// [`Self::envelope_task`] to stamp one. + pub(crate) fn envelope(&self, run_id: &RunId, event: GraphEvent) -> GraphEventEnvelope { + self.envelope_task(run_id, None, event) + } + + /// [`Self::envelope`] that also stamps the envelope's `task_id`. + pub(crate) fn envelope_task( + &self, + run_id: &RunId, + task_id: Option<&TaskId>, + event: GraphEvent, + ) -> GraphEventEnvelope { + let seq = self + .sequence + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + GraphEventEnvelope { + run_id: run_id.clone(), + task_id: task_id.cloned(), + ns: self.namespace.clone(), + seq, + event, } } + + /// [`Self::emit`] for the handful of manual/out-of-band checkpoint APIs + /// (`update_state`, `bulk_update_state`, …) that run outside any live + /// [`RunCtx`], and so have no real [`RunId`] to stamp — the envelope + /// carries an empty one rather than a fabricated live run. + pub(crate) fn emit_unscoped(&self, event: GraphEvent) { + self.emit(&RunId::from(String::new()), event); + } + + /// Resets this graph instance's sequence counter to a fresh, independent + /// one. Used when embedding a graph as a subgraph node + /// ([`crate::subgraph`]): the embedded instance gets its own namespace + /// already (see [`crate::stream::GraphEventEnvelope`]), and sharing the + /// parent's counter would only entangle two otherwise-independent + /// sequences for no benefit. + pub(crate) fn with_fresh_sequence(mut self) -> Self { + self.sequence = Arc::new(std::sync::atomic::AtomicU64::new(0)); + self + } } +#[cfg(test)] +mod drain_test; +#[cfg(test)] +mod durable_task_test; +#[cfg(test)] +mod durable_test; +#[cfg(test)] +mod interrupt_selectors_test; +#[cfg(test)] +mod policy_test; #[cfg(test)] mod test; diff --git a/crates/tinyagents-graph/src/compiled/policy_test.rs b/crates/tinyagents-graph/src/compiled/policy_test.rs new file mode 100644 index 00000000..42b504c1 --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/policy_test.rs @@ -0,0 +1,556 @@ +//! Unit tests for per-node execution policy (Unit A): per-node +//! retry/timeout overrides, idle timeouts with heartbeats, task caching, +//! `on_error` recovery, and real `defer` scheduling. + +use super::*; +use crate::builder::{GraphBuilder, NodeCachePolicy, NodeContext, NodePolicy}; +use crate::cache::InMemoryTaskCache; +use crate::command::{Command, NodeResult, RouteTarget, Send}; +use crate::reducer::ClosureStateReducer; +use crate::stream::{CollectingSink, GraphEvent}; +use serde_json::json; +use tinyagents_harness::retry::RetryPolicy; + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use std::time::Duration; + +/// A single-node graph whose handler fails (with a retryable model error) +/// the first `fail_times` invocations, then succeeds with `+1`. +fn flaky_builder(fail_times: usize, attempts: Arc) -> GraphBuilder { + GraphBuilder::::overwrite() + .add_node("flaky", move |s, _c: NodeContext| { + let attempts = attempts.clone(); + async move { + let n = attempts.fetch_add(1, AtomicOrdering::SeqCst); + if n < fail_times { + Err(TinyAgentsError::Model(format!("transient blip {n}"))) + } else { + Ok(NodeResult::Update(s + 1)) + } + } + }) + .set_entry("flaky") + .set_finish("flaky") +} + +// ── A.1: per-node retry / timeout overrides ────────────────────────────────── + +/// A node's own `NodePolicy::retry` wins over the graph-wide +/// `with_node_retry` policy: the node retries per its own attempt cap even +/// though the graph-wide policy would have given up earlier. +#[tokio::test] +async fn per_node_retry_policy_overrides_graph_wide_retry() { + let attempts = Arc::new(AtomicUsize::new(0)); + // Fails 3 times; graph-wide budget is 2 attempts (would fail), per-node + // budget is 5 attempts (recovers on the 4th). + let graph = flaky_builder(3, attempts.clone()) + .with_node_policy( + "flaky", + NodePolicy { + retry: Some( + RetryPolicy::default() + .with_max_attempts(5) + .with_backoff_sleep(false), + ), + ..NodePolicy::default() + }, + ) + .compile() + .unwrap() + .with_node_retry( + RetryPolicy::default() + .with_max_attempts(2) + .with_backoff_sleep(false), + ); + + let run = graph.run(10).await.unwrap(); + assert_eq!(run.state, 11); + assert_eq!( + attempts.load(AtomicOrdering::SeqCst), + 4, + "1 try + 3 retries" + ); +} + +/// With no graph-wide retry at all, a per-node retry policy still applies. +#[tokio::test] +async fn per_node_retry_policy_applies_without_graph_wide_retry() { + let attempts = Arc::new(AtomicUsize::new(0)); + let graph = flaky_builder(1, attempts.clone()) + .with_node_policy( + "flaky", + NodePolicy { + retry: Some( + RetryPolicy::default() + .with_max_attempts(2) + .with_backoff_sleep(false), + ), + ..NodePolicy::default() + }, + ) + .compile() + .unwrap(); + + let run = graph.run(10).await.unwrap(); + assert_eq!(run.state, 11); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 2); +} + +/// A node's own `NodePolicy::timeout` (shorter than the graph-wide +/// `with_node_timeout`) is what bounds it: the handler times out at the +/// per-node value even though the graph default would have let it finish. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn per_node_timeout_overrides_graph_wide_timeout() { + let graph = GraphBuilder::::overwrite() + .with_node_timeout(Duration::from_secs(5)) + .add_node("slow", |s: i32, _c: NodeContext| async move { + tokio::time::sleep(Duration::from_millis(300)).await; + Ok(NodeResult::Update(s)) + }) + .with_node_policy( + "slow", + NodePolicy::default().with_timeout(Duration::from_millis(20)), + ) + .set_entry("slow") + .set_finish("slow") + .compile() + .unwrap(); + + let started = std::time::Instant::now(); + let err = graph.run(0).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); + assert!( + started.elapsed() < Duration::from_millis(250), + "the per-node 20ms timeout fired, not the 5s graph-wide one" + ); +} + +/// `set_node_defaults` is the middle precedence layer: a node with no +/// per-node timeout uses the defaults' timeout over the legacy graph-wide +/// `with_node_timeout`. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn node_defaults_timeout_beats_legacy_graph_wide_timeout() { + let graph = GraphBuilder::::overwrite() + .with_node_timeout(Duration::from_millis(20)) + .set_node_defaults(NodePolicy::default().with_timeout(Duration::from_secs(5))) + .add_node("slow", |s: i32, _c: NodeContext| async move { + tokio::time::sleep(Duration::from_millis(60)).await; + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("slow") + .set_finish("slow") + .compile() + .unwrap(); + + let run = graph.run(0).await.unwrap(); + assert_eq!( + run.state, 1, + "the 5s default timeout let the 60ms node finish" + ); +} + +// ── A.1: idle timeout + heartbeat ──────────────────────────────────────────── + +/// A handler that heartbeats more often than its `idle_timeout` survives +/// well past what a flat timeout of that same duration would have allowed. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn heartbeat_keeps_idle_timeout_from_firing() { + let graph = GraphBuilder::::overwrite() + .add_node("worker", |s: i32, ctx: NodeContext| async move { + // Runs for 200ms total, heartbeating every 20ms — far inside the + // 60ms idle window, but 3x longer than a flat 60ms timeout. + for _ in 0..10 { + tokio::time::sleep(Duration::from_millis(20)).await; + ctx.heartbeat(); + } + Ok(NodeResult::Update(s + 1)) + }) + .with_node_policy( + "worker", + NodePolicy::default().with_idle_timeout(Duration::from_millis(60)), + ) + .set_entry("worker") + .set_finish("worker") + .compile() + .unwrap(); + + let run = graph.run(0).await.unwrap(); + assert_eq!(run.state, 1); +} + +/// With only `idle_timeout` set and no heartbeat ever sent, the node times +/// out at (approximately) exactly the idle duration after it starts — the +/// idle timeout degrades to a flat timeout. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn idle_timeout_without_heartbeats_is_a_flat_timeout() { + let idle = Duration::from_millis(80); + let graph = GraphBuilder::::overwrite() + .add_node("silent", |s: i32, _c: NodeContext| async move { + tokio::time::sleep(Duration::from_secs(5)).await; + Ok(NodeResult::Update(s)) + }) + .with_node_policy("silent", NodePolicy::default().with_idle_timeout(idle)) + .set_entry("silent") + .set_finish("silent") + .compile() + .unwrap(); + + let started = std::time::Instant::now(); + let err = graph.run(0).await.unwrap_err(); + let elapsed = started.elapsed(); + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); + assert!( + elapsed >= idle && elapsed < idle + Duration::from_millis(150), + "expected the idle timeout to fire in [{idle:?}, {idle:?} + slop), got {elapsed:?}" + ); +} + +/// A flat `timeout` and an `idle_timeout` on the same node are independent +/// ceilings: a handler that heartbeats forever is still cut off by the +/// flat timeout. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn flat_timeout_still_bounds_a_heartbeating_handler() { + let graph = GraphBuilder::::overwrite() + .add_node("chatty", |s: i32, ctx: NodeContext| async move { + loop { + tokio::time::sleep(Duration::from_millis(10)).await; + ctx.heartbeat(); + if false { + break; + } + } + #[allow(unreachable_code)] + Ok(NodeResult::Update(s)) + }) + .with_node_policy( + "chatty", + NodePolicy::default() + .with_idle_timeout(Duration::from_millis(100)) + .with_timeout(Duration::from_millis(60)), + ) + .set_entry("chatty") + .set_finish("chatty") + .compile() + .unwrap(); + + let started = std::time::Instant::now(); + let err = graph.run(0).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Timeout(_)), "got {err:?}"); + assert!(started.elapsed() < Duration::from_millis(300)); +} + +// ── A.2: task cache ────────────────────────────────────────────────────────── + +/// A single-node graph counting handler invocations, cached on the input +/// state's value. +fn counting_cached_graph( + calls: Arc, + ttl: Option, +) -> (CompiledGraph, Arc) { + let cache = Arc::new(InMemoryTaskCache::new()); + let mut policy = NodeCachePolicy::new(|s: &i32, _arg| format!("state={s}")); + if let Some(ttl) = ttl { + policy = policy.with_ttl(ttl); + } + let graph = GraphBuilder::::overwrite() + .add_node("compute", move |s: i32, _c: NodeContext| { + let calls = calls.clone(); + async move { + calls.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(s * 10)) + } + }) + .set_entry("compute") + .set_finish("compute") + .compile() + .unwrap() + .with_task_cache(cache.clone()) + .with_cached_node("compute", policy); + (graph, cache) +} + +/// A second run with the same cache key skips the handler entirely, replays +/// the cached update, and reports the hit as `TaskCompleted { cached: true }`. +#[tokio::test] +async fn cache_hit_skips_handler_and_emits_cached_task_completed() { + let calls = Arc::new(AtomicUsize::new(0)); + let (graph, _cache) = counting_cached_graph(calls.clone(), None); + let sink = Arc::new(CollectingSink::new()); + let graph = graph.with_event_sink(sink.clone()); + + let first = graph.run(4).await.unwrap(); + assert_eq!(first.state, 40); + assert_eq!(calls.load(AtomicOrdering::SeqCst), 1); + assert!( + !sink + .events() + .iter() + .any(|e| matches!(e, GraphEvent::TaskCompleted { cached: true, .. })), + "the first run was a miss" + ); + + let second = graph.run(4).await.unwrap(); + assert_eq!(second.state, 40, "the cached update was replayed"); + assert_eq!( + calls.load(AtomicOrdering::SeqCst), + 1, + "the handler was not invoked on the cache hit" + ); + assert_eq!( + second.visited, + vec![NodeId::from("compute")], + "a cached node still counts as visited" + ); + let hit = sink.events().into_iter().find(|e| { + matches!( + e, + GraphEvent::TaskCompleted { + cached: true, + step: 1, + .. + } + ) + }); + assert!(hit.is_some(), "expected a cached TaskCompleted event"); + + // A different key is a miss again. + let third = graph.run(5).await.unwrap(); + assert_eq!(third.state, 50); + assert_eq!(calls.load(AtomicOrdering::SeqCst), 2); +} + +/// Once an entry's TTL elapses the handler runs again (and repopulates). +#[tokio::test] +async fn cache_ttl_expiry_reruns_the_handler() { + let calls = Arc::new(AtomicUsize::new(0)); + let (graph, _cache) = counting_cached_graph(calls.clone(), Some(Duration::from_millis(40))); + + graph.run(1).await.unwrap(); + graph.run(1).await.unwrap(); + assert_eq!( + calls.load(AtomicOrdering::SeqCst), + 1, + "second run was a hit" + ); + + tokio::time::sleep(Duration::from_millis(80)).await; + let run = graph.run(1).await.unwrap(); + assert_eq!(run.state, 10); + assert_eq!( + calls.load(AtomicOrdering::SeqCst), + 2, + "the expired entry forced a re-run" + ); +} + +/// The cache key function receives each activation's `send_arg`, so a +/// `Send` fan-out of one node keys (and hits) per argument. +#[tokio::test] +async fn cache_key_receives_send_arg_per_fanout_activation() { + let calls = Arc::new(AtomicUsize::new(0)); + let seen_args = Arc::new(std::sync::Mutex::new(Vec::::new())); + let cache = Arc::new(InMemoryTaskCache::new()); + let key_args = seen_args.clone(); + let graph = GraphBuilder::, Vec>::new() + .set_reducer(ClosureStateReducer::new( + |mut s: Vec, u: Vec| { + s.extend(u); + Ok(s) + }, + )) + .add_node("fan", |_s, _c: NodeContext| async move { + Ok(NodeResult::Command(Command { + update: None, + goto: vec![ + RouteTarget::Send(Send::new("work", json!("a"))), + RouteTarget::Send(Send::new("work", json!("b"))), + ], + resume: None, + resume_by_task: Default::default(), + })) + }) + .add_node("work", { + let calls = calls.clone(); + move |_s, ctx: NodeContext| { + let calls = calls.clone(); + async move { + calls.fetch_add(1, AtomicOrdering::SeqCst); + let arg = ctx.send_arg.and_then(|v| v.as_str().map(String::from)); + Ok(NodeResult::Update(vec![format!( + "work:{}", + arg.unwrap_or_default() + )])) + } + } + }) + .mark_command_routing("fan") + .set_entry("fan") + .set_finish("work") + .compile() + .unwrap() + .with_task_cache(cache.clone()) + .with_cached_node( + "work", + NodeCachePolicy::new(move |_s: &Vec, arg: Option<&serde_json::Value>| { + let arg = arg.map(|v| v.to_string()).unwrap_or_default(); + key_args.lock().unwrap().push(arg.clone()); + format!("arg={arg}") + }), + ); + + let first = graph.run(vec![]).await.unwrap(); + let mut got = first.state.clone(); + got.sort(); + assert_eq!(got, vec!["work:a", "work:b"]); + assert_eq!(calls.load(AtomicOrdering::SeqCst), 2); + { + let mut args = seen_args.lock().unwrap(); + args.sort(); + assert_eq!(*args, vec!["\"a\"", "\"b\""], "key fn saw each send_arg"); + } + + // Second run: both fan-out activations are cache hits, the fan node + // itself (uncached) still runs. + let second = graph.run(vec![]).await.unwrap(); + let mut got = second.state.clone(); + got.sort(); + assert_eq!( + got, + vec!["work:a", "work:b"], + "cached updates were replayed" + ); + assert_eq!( + calls.load(AtomicOrdering::SeqCst), + 2, + "neither fan-out activation invoked the handler" + ); +} + +// ── A.3: on_error recovery ──────────────────────────────────────────────── + +/// `on_error` is consulted only after the retry budget is exhausted (or the +/// error is non-retryable); returning `Some(command)` recovers the node with +/// that command's update instead of failing the run. +#[tokio::test] +async fn on_error_recovers_after_retries_are_exhausted() { + let attempts = Arc::new(AtomicUsize::new(0)); + let graph = flaky_builder(usize::MAX, attempts.clone()) + .with_node_policy( + "flaky", + NodePolicy { + retry: Some( + RetryPolicy::default() + .with_max_attempts(2) + .with_backoff_sleep(false), + ), + on_error: Some(Arc::new(|state: &i32, _err: &TinyAgentsError| { + Some(Command { + update: Some(state + 100), + goto: vec![], + resume: None, + resume_by_task: Default::default(), + }) + })), + ..NodePolicy::default() + }, + ) + .compile() + .unwrap(); + + let run = graph.run(10).await.unwrap(); + assert_eq!(run.state, 110, "on_error's command update was applied"); + assert_eq!( + attempts.load(AtomicOrdering::SeqCst), + 2, + "1 try + 1 retry, then on_error recovered instead of a 3rd attempt" + ); +} + +/// `on_error` returning `None` falls through to the ordinary escalation — +/// the run still fails with the underlying error. +#[tokio::test] +async fn on_error_none_falls_through_to_the_original_error() { + let attempts = Arc::new(AtomicUsize::new(0)); + let graph = flaky_builder(usize::MAX, attempts.clone()) + .with_node_policy( + "flaky", + NodePolicy { + on_error: Some(Arc::new(|_state: &i32, _err: &TinyAgentsError| None)), + ..NodePolicy::default() + }, + ) + .compile() + .unwrap(); + + let err = graph.run(10).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + assert_eq!( + attempts.load(AtomicOrdering::SeqCst), + 1, + "no retry policy set" + ); +} + +// ── A.4: real defer ────────────────────────────────────────────────────── + +/// A deferred node fanned out to alongside a non-deferred sibling is held +/// back: it runs in a later superstep than its sibling, once the sibling's +/// own successor leaves nothing non-deferred in the frontier — not +/// concurrently with it, which is what a purely cosmetic `mark_deferred` +/// marker (metadata-only, no scheduling effect) would have produced. +#[tokio::test] +async fn deferred_node_runs_only_once_the_frontier_has_no_other_work() { + let graph = GraphBuilder::, Vec>::new() + .set_reducer(ClosureStateReducer::new( + |mut s: Vec, u: Vec| { + s.extend(u); + Ok(s) + }, + )) + .add_node("start", |_s, _c: NodeContext| async move { + Ok(NodeResult::Command(Command { + update: None, + goto: vec![ + RouteTarget::Node(NodeId::from("worker")), + RouteTarget::Node(NodeId::from("synth")), + ], + resume: None, + resume_by_task: Default::default(), + })) + }) + .add_node("worker", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update(vec!["worker".to_string()])) + }) + .add_node("synth", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update(vec!["synth".to_string()])) + }) + .mark_command_routing("start") + .mark_deferred("synth") + .set_entry("start") + .set_finish("worker") + .set_finish("synth") + .compile() + .unwrap(); + + let sink = Arc::new(CollectingSink::new()); + let graph = graph.with_event_sink(sink.clone()); + let run = graph.run(vec![]).await.unwrap(); + let mut got = run.state.clone(); + got.sort(); + assert_eq!(got, vec!["synth", "worker"], "both branches still ran"); + + let started_step = |name: &str| { + sink.events().into_iter().find_map(|e| match e { + GraphEvent::NodeStarted { node, step } if node.as_str() == name => Some(step), + _ => None, + }) + }; + let worker_step = started_step("worker").expect("worker started"); + let synth_step = started_step("synth").expect("synth started"); + assert!( + synth_step > worker_step, + "the deferred node must run in a later superstep than its \ + non-deferred sibling, got worker={worker_step} synth={synth_step}" + ); +} diff --git a/crates/tinyagents-graph/src/compiled/resume.rs b/crates/tinyagents-graph/src/compiled/resume.rs new file mode 100644 index 00000000..d856543f --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/resume.rs @@ -0,0 +1,401 @@ +//! Resume: loading a checkpoint, filtering out already-completed tasks, and +//! building the resume-value map handed to the re-run node(s). +//! +//! Split out of `executor.rs`; see that module's doc comment for the public +//! `resume`/`resume_from`/`retry` entry points that call into +//! [`CompiledGraph::resume_from_inner`]. + +use super::*; + +use crate::compiled::executor::RunSeed; + +impl CompiledGraph +where + State: Clone + Send + Sync + 'static, + Update: Send + 'static, +{ + pub(super) async fn resume_from_inner( + &self, + thread_id: ThreadId, + target: ResumeTarget, + command: Command, + binding: Option, + options: RunOptions, + ) -> Result> { + let checkpointer = self + .checkpointer + .as_ref() + .ok_or_else(|| TinyAgentsError::Resume("no checkpointer configured".to_string()))?; + + let checkpoint_id = match &target { + ResumeTarget::Latest => None, + ResumeTarget::Checkpoint(id) => Some(id.as_str()), + }; + let checkpoint = checkpointer + .get_scoped(thread_id.as_str(), checkpoint_id, &self.namespace) + .await? + .ok_or_else(|| match &target { + ResumeTarget::Latest => { + TinyAgentsError::Resume(format!("no checkpoint found for thread `{thread_id}`")) + } + ResumeTarget::Checkpoint(id) => TinyAgentsError::Resume(format!( + "no checkpoint `{id}` found for thread `{thread_id}`" + )), + })?; + // Resume *loads* this checkpoint — it is a read, not a write — so emit a + // restore event, not `CheckpointSaved` (which would falsely inflate + // persisted-checkpoint counts and mislead durability observers). + // No new run id has been minted yet at this point in the resume + // path (see `executor.rs::resume`) — the checkpoint being restored + // is the one whose own `run_id` this read is about, so that (rather + // than the not-yet-existing resumed run's id) is what the envelope + // is stamped with. + let restored_run_id = checkpoint + .run_id + .clone() + .map(RunId::from) + .unwrap_or_else(|| RunId::from(String::new())); + self.emit( + &restored_run_id, + GraphEvent::CheckpointRestored { + checkpoint_id: CheckpointId::new(checkpoint.checkpoint_id.clone()), + }, + ); + + // `checkpoint` was already normalized on read (every backend's decode + // path calls `Checkpoint::normalize`), so `tasks` is always the + // single source of truth here, regardless of the stored record's + // original format version. + let active: Vec = checkpoint.tasks.iter().map(Activation::from).collect(); + if active.is_empty() { + return Err(TinyAgentsError::Resume( + "checkpoint has no pending nodes to resume".to_string(), + )); + } + + // Partial-failure guard. The boundary that produced this checkpoint + // recorded a completion marker per task that had already finished; a + // node named by *both* the pending set and that ledger has therefore + // already run, and re-running it would repeat its side effects. On a + // checkpoint the executor itself wrote the two sets are disjoint, so + // this is a no-op — it earns its keep on a checkpoint that was + // hand-built, time-travelled to, or edited through `update_state`, + // where `next_nodes` can legitimately disagree with what ran. + let completed_config = CheckpointConfig { + thread_id: thread_id.to_string(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: self.namespace.clone(), + }; + let recorded = checkpointer.get_writes(&completed_config).await?; + let ledger: &[crate::checkpoint::PendingWrite] = if recorded.is_empty() { + &checkpoint.pending_writes + } else { + &recorded + }; + // Only a completion marker (or any other non-replay write) says a + // task ran. A replay memo — a `durable_task` write or a deferred + // `interrupt_after` result — belongs to a task that is still + // *pending* and must re-run (consuming the memo), so it is excluded + // here and grouped per task below instead. + let done: HashSet = ledger + .iter() + .filter(|w| !w.is_task_replay()) + .map(|w| w.task_id.as_str().to_string()) + .collect(); + let mut task_writes: HashMap> = HashMap::new(); + for write in ledger.iter().filter(|w| w.is_task_replay()) { + task_writes + .entry(write.task_id.as_str().to_string()) + .or_default() + .push(write.clone()); + } + // Executor-injected `interrupt_before`/`interrupt_after` pauses this + // checkpoint recorded: resuming acknowledges them, so the re-run of + // that task skips the same phase instead of pausing again. + // Plus the acks an earlier pause of a still-pending task already + // carried into this checkpoint's metadata (see + // `boundary::with_carried_acks`). + let mut acknowledged_interrupts: HashSet = checkpoint + .interrupts + .iter() + .filter_map(|interrupt| { + let phase = interrupt.payload.get("phase")?.as_str()?; + if phase != "before" && phase != "after" { + return None; + } + let task_id = interrupt.task_id.as_ref()?; + Some( + crate::compiled::run_ctx::RunCtx::::interrupt_ack_key( + phase, task_id, + ), + ) + }) + .collect(); + if let Some(carried) = checkpoint + .metadata + .get("acknowledged_interrupts") + .and_then(serde_json::Value::as_array) + { + acknowledged_interrupts.extend( + carried + .iter() + .filter_map(serde_json::Value::as_str) + .map(str::to_string), + ); + } + let active: Vec = if done.is_empty() { + active + } else { + let filtered: Vec = active + .iter() + // A node name is not a task identity: a Send fan-out can have + // several live activations of one node. Legacy checkpoints + // have no persisted task id, so leave them runnable. + .filter(|a| a.task_id.as_str().is_empty() || !done.contains(a.task_id.as_str())) + .cloned() + .collect(); + if filtered.is_empty() { + // Every pending node claims to have run. Trust the pending set + // rather than turning a resumable checkpoint into a hard error: + // a wrong re-run is recoverable, a stuck thread is not. + tracing::warn!( + "[graph:resume] every pending node of checkpoint `{}` has a completion \ + marker; resuming them anyway rather than stranding the thread", + checkpoint.checkpoint_id + ); + active + } else { + if filtered.len() != active.len() { + tracing::debug!( + "[graph:resume] checkpoint `{}`: skipping {} already-completed task(s)", + checkpoint.checkpoint_id, + active.len() - filtered.len() + ); + } + filtered + } + }; + + // The resume value belongs to the node(s) that actually interrupted. + // Interrupt/failure boundaries persist `pending` as exactly the + // stalled (interrupted/failed) branches of that step — a completed + // sibling's routing is deferred rather than folded into `pending` + // (see `boundary::advance`'s `carried_completed` handling) — but a + // checkpoint could still carry a wider pending set (a hand-built one, + // or one written before this policy), so this still keys off + // `interrupted_nodes` rather than assuming `active` is exactly the + // interrupted set. A boundary that recorded no interrupt (a failure + // boundary, resumed via `retry` with no value) keeps the old + // fan-across-pending behaviour. + // I1/R5: keyed by task id (falling back to node id) so a `Send` + // fan-out of the same node — several live activations sharing one + // `NodeId` — each receive their own resume value instead of every + // same-node activation racing for a single node-keyed slot. Prefer + // the persisted interrupts' own `task_id` (stamped by the interrupt + // boundary, R5) when present; fall back to `interrupted_nodes` (node + // names only — a checkpoint written before task identity existed, or + // a re-emitted subgraph interrupt whose task id was not stamped), + // keying by every active activation of that node. + let mut resume_map: HashMap = HashMap::new(); + // A caller-supplied per-task map (I1) always wins: it names its + // targets explicitly, so there is nothing to infer. + for (task_id, value) in &command.resume_by_task { + resume_map.insert(task_id.as_str().to_string(), value.clone()); + } + if let Some(value) = command.resume { + let task_targets: Vec = checkpoint + .interrupts + .iter() + .filter_map(|i| i.task_id.as_ref()) + .map(|t| t.as_str().to_string()) + .filter(|t| active.iter().any(|a| a.task_id.as_str() == t)) + .collect(); + if !task_targets.is_empty() { + for task_id in task_targets { + resume_map.entry(task_id).or_insert_with(|| value.clone()); + } + } else { + let interrupted = interrupted_nodes(&checkpoint, &active); + if interrupted.is_empty() { + for activation in &active { + resume_map + .entry(resume_key(activation)) + .or_insert_with(|| value.clone()); + } + } else { + for node in interrupted { + for activation in active.iter().filter(|a| a.node == node) { + resume_map + .entry(resume_key(activation)) + .or_insert_with(|| value.clone()); + } + } + } + } + } + + // Fail closed on `Interrupt::response_schema`: every value this resume + // would hand to a schema-bearing interrupt's task is validated *now*, + // before `execute` claims the thread or any boundary writes a + // checkpoint, so a rejected value leaves the thread's checkpoint + // exactly as it was. + self.validate_resume_values(&checkpoint, &active, &resume_map)?; + + // Restore accumulated barrier arrivals so a join's precondition survives + // the interrupt/failure boundary this checkpoint recorded. + let initial_barriers = barriers_from_persisted(&checkpoint.barrier_arrivals); + // Chain the first post-resume boundary onto the checkpoint we loaded so + // the lineage spine stays connected across the resume. + let initial_parent = Some(checkpoint.checkpoint_id.clone()); + + // A mid-step checkpoint (an interrupt/failure `loop`-source boundary + // — stamped with `interrupted_nodes` or `failed_node`) leaves its + // completed siblings unrouted (see `boundary::advance`'s + // `carried_completed` doc, the C2 fix): carry their node ids forward + // so this resumed run's *first* boundary routes the whole original + // step together, rather than routing only the freshly re-run + // pending set in isolation (which would let a successor observe a + // state missing whatever the other, already-completed siblings + // wrote). + // + // The `source == "loop"` check matters: `update_state` (I2) can + // *also* stamp `interrupted_nodes` onto an `update`-sourced + // checkpoint, purely to preserve resume-value provenance — but + // `update_state` always fully resolves every carried completion's + // routing itself before writing (see `state_api::update_state`), so + // its `completed_tasks` never represents owed work. Treating it as + // mid-step here would route those already-resolved completions a + // second time, scheduling their successors twice. + let source_is_loop = checkpoint + .metadata + .get("source") + .and_then(serde_json::Value::as_str) + == Some("loop"); + let mid_step = source_is_loop + && (checkpoint.metadata.get("interrupted_nodes").is_some() + || checkpoint.metadata.get("failed_node").is_some()); + let carried_completed = if mid_step && !checkpoint.completed.is_empty() { + Some( + checkpoint + .completed + .iter() + .map(|c| (c.node.clone(), c.routes.clone())) + .collect(), + ) + } else { + None + }; + // I3: continue this thread's step counter and per-node visit counts + // from the loaded checkpoint instead of restarting at zero, so + // `metadata.step` (and `get_state_history`) stays monotonic and + // `RecursionPolicy::max_visits_per_node` bounds the whole thread's + // lifetime rather than resetting every resume. + let initial_steps = checkpoint.to_metadata().step; + let initial_node_visits = node_visits_from_persisted(&checkpoint.metadata); + // I5/R3: carry the per-node channel-versions bookkeeping forward + // across the resume, same reasoning as `initial_node_visits`. + let initial_versions_seen = checkpoint.versions_seen.clone().into_iter().collect(); + + self.execute(RunSeed { + state: checkpoint.state, + active, + thread_id: Some(thread_id), + resume_map, + barriers: initial_barriers, + parent: initial_parent, + binding, + resume_seed: crate::compiled::run_ctx::ResumeSeed { + initial_steps, + initial_node_visits, + initial_versions_seen, + carried_completed, + task_writes, + acknowledged_interrupts, + }, + options, + _update: std::marker::PhantomData, + }) + .await + } + + /// Validates the resume value each schema-bearing pending interrupt + /// would receive (see [`Interrupt::response_schema`]) against that + /// schema, via [`tinyagents_harness::tool::validate_against_schema`]. + /// + /// The value is looked up exactly as [`RunCtx::node_context`] will hand + /// it out — by the interrupt's task id, falling back to its node id (the + /// legacy/whole-node key) — so what is validated is what the node would + /// see. An interrupt whose task receives no value at all (a bare + /// `retry`) has nothing to validate. Must be called before any + /// checkpoint write. + /// + /// [`RunCtx::node_context`]: crate::compiled::run_ctx::RunCtx::node_context + fn validate_resume_values( + &self, + checkpoint: &Checkpoint, + active: &[Activation], + resume_map: &HashMap, + ) -> Result<()> { + for interrupt in &checkpoint.interrupts { + let Some(schema) = &interrupt.response_schema else { + continue; + }; + let value = interrupt + .task_id + .as_ref() + .and_then(|task| resume_map.get(task.as_str())) + .or_else(|| { + // Node-keyed fallback, only when the interrupt's node is + // actually pending (matching `node_context`'s lookup). + active + .iter() + .any(|a| a.node == interrupt.node) + .then(|| resume_map.get(interrupt.node.as_str())) + .flatten() + }); + let Some(value) = value else { + continue; + }; + tinyagents_harness::tool::validate_against_schema(schema, value).map_err(|err| { + TinyAgentsError::Validation(format!( + "resume value for interrupt `{}` (node `{}`) rejected by its response_schema: \ + {err}", + interrupt.id, interrupt.node + )) + })?; + } + Ok(()) + } +} + +/// The key an activation's resume value is looked up under: its task id +/// when known (I1/R5 — distinguishes concurrent same-node activations), +/// falling back to its node id (legacy checkpoints, or a value fanned across +/// every pending node with no interrupt provenance). +fn resume_key(activation: &Activation) -> String { + if activation.task_id.as_str().is_empty() { + activation.node.to_string() + } else { + activation.task_id.as_str().to_string() + } +} + +/// Parses a checkpoint's persisted `metadata.node_visits` object (see +/// `boundary`'s checkpoint builders) back into the live per-node visit-count +/// map. Missing/malformed metadata (checkpoints written before this field +/// existed) yields an empty map — the pre-I3 behavior for that checkpoint. +fn node_visits_from_persisted(metadata: &serde_json::Value) -> HashMap { + metadata + .get("node_visits") + .and_then(serde_json::Value::as_object) + .map(|obj| { + obj.iter() + .filter_map(|(node, count)| { + count + .as_u64() + .map(|c| (NodeId::from(node.as_str()), c as usize)) + }) + .collect() + }) + .unwrap_or_default() +} diff --git a/crates/tinyagents-graph/src/compiled/routing.rs b/crates/tinyagents-graph/src/compiled/routing.rs index 2755ec26..a3111cad 100644 --- a/crates/tinyagents-graph/src/compiled/routing.rs +++ b/crates/tinyagents-graph/src/compiled/routing.rs @@ -22,34 +22,51 @@ where /// (see [`Self::route`]); `barrier_arrivals` is mutated in place as /// waiting-node predecessors arrive, so a barrier can still be pending /// across supersteps. + /// + /// `completed` pairs each branch with its *original* active-set index + /// (not necessarily `0..completed.len()` in order — see + /// [`crate::compiled::step::StepRun::completed`] and + /// [`crate::compiled::boundary::CompiledGraph::advance`]'s + /// `carried_completed` handling, both of which can hand this a + /// non-contiguous or reordered set spanning more than one step's + /// original indices). That original index is what `goto_map` is keyed + /// by, so it is threaded through explicitly rather than re-derived from + /// `completed`'s own position. pub(super) fn route_completed( &self, - completed: &[Activation], + run_id: &RunId, + completed: &[(usize, Activation)], goto_map: &HashMap>, state: &State, barrier_arrivals: &mut HashMap>, ) -> Result> { let mut next: Vec = Vec::new(); let mut next_seen: HashSet = HashSet::new(); - // Resolved targets per activation index, captured once here and - // reused by the barrier-relief pass below instead of calling - // `self.route` a second time — a router closure is only guaranteed - // pure/idempotent per the `route`/`add_conditional_edges` contract, - // not safe to invoke twice for the same activation. + // Resolved targets per completed-slice position (not original + // index), captured once here and reused by the barrier-relief pass + // below instead of calling `self.route` a second time — a router + // closure is only guaranteed pure/idempotent per the + // `route`/`add_conditional_edges` contract, not safe to invoke + // twice for the same activation. let mut resolved: Vec> = Vec::with_capacity(completed.len()); - for (index, activation) in completed.iter().enumerate() { + for (orig_index, activation) in completed.iter() { let node_id = &activation.node; - let targets = self.route(node_id, goto_map.get(&index).map(Vec::as_slice), state)?; + let targets = + self.route(node_id, goto_map.get(orig_index).map(Vec::as_slice), state)?; resolved.push(targets.clone()); for target in targets { let tnode = target.node().clone(); if tnode.as_str() == END { continue; } - self.emit(GraphEvent::RouteSelected { - node: node_id.clone(), - target: tnode.clone(), - }); + self.emit_task( + run_id, + Some(&activation.task_id), + GraphEvent::RouteSelected { + node: node_id.clone(), + target: tnode.clone(), + }, + ); // Barrier gating: hold a waiting node until every required // predecessor has arrived (possibly across supersteps). if let Some(required) = self.waiting.get(&tnode) { @@ -62,18 +79,18 @@ where } // `Send` activations may repeat the same node (each carries its // own arg); plain activations are deduplicated by node. - let send_arg = target.send_arg().cloned(); + let send_arg = target.send_arg().cloned().map(Arc::new); if send_arg.is_some() { next.push(Activation { node: tnode, send_arg, - task_id: String::new(), + task_id: TaskId::from(String::new()), }); } else if next_seen.insert(tnode.clone()) { next.push(Activation { node: tnode, send_arg: None, - task_id: String::new(), + task_id: TaskId::from(String::new()), }); } } @@ -107,7 +124,7 @@ where let source_indices: Vec = completed .iter() .enumerate() - .filter(|(_, activation)| activation.node == relief.source) + .filter(|(_, (_, activation))| activation.node == relief.source) .map(|(index, _)| index) .collect(); if source_indices.is_empty() { @@ -144,7 +161,7 @@ where next.push(Activation { node: relief.barrier_node.clone(), send_arg: None, - task_id: String::new(), + task_id: TaskId::from(String::new()), }); } } @@ -171,16 +188,29 @@ where if from == to { return true; } - let mut current = from; + // A static fan-out (`self.edges` mapping to more than one target) is + // still fully deterministic — every target in the list unconditionally + // activates, unlike a conditional branch — so this walks every static + // successor of `from`, not just a single chain, tracking visited nodes + // to stay finite over a cycle. + let mut stack: Vec<&NodeId> = vec![from]; let mut seen: HashSet<&NodeId> = HashSet::new(); - while let Some(next) = self.edges.get(current) { - if next == to { - return true; + while let Some(current) = stack.pop() { + if !seen.insert(current) { + continue; } - if next == stop || !seen.insert(next) { - return false; + let Some(targets) = self.edges.get(current) else { + continue; + }; + for next in targets { + if next == to { + return true; + } + if next == stop { + continue; + } + stack.push(next); } - current = next; } false } @@ -203,11 +233,14 @@ where self.validate_route_targets(node_id, targets)?; return Ok(targets.to_vec()); } - if let Some(target) = self.edges.get(node_id) { - return Ok(vec![RouteTarget::Node(target.clone())]); + if let Some(targets) = self.edges.get(node_id) { + return Ok(targets + .iter() + .map(|target| RouteTarget::Node(target.clone())) + .collect()); } if let Some(branch) = self.branches.get(node_id) { - let route = (branch.router)(state); + let route = (branch.router)(state).to_string(); let target = branch.routes.get(&route).cloned().ok_or_else(|| { TinyAgentsError::MissingRoute { node: node_id.to_string(), diff --git a/crates/tinyagents-graph/src/compiled/run_ctx.rs b/crates/tinyagents-graph/src/compiled/run_ctx.rs new file mode 100644 index 00000000..9e9ce6fa --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/run_ctx.rs @@ -0,0 +1,567 @@ +//! Per-run execution context threaded through the superstep loop. +//! +//! [`RunCtx`] bundles run identity (ids, namespace), clocks/deadlines, +//! recursion bookkeeping, the accumulators a superstep loop carries forward +//! (`node_visits`, `barrier_arrivals`, `visited`, `all_child_runs`, +//! `steps`/checkpoint lineage), and the handles for background checkpoint +//! writes and status/event I/O. +//! +//! It exists so the step-running, boundary, and resume helpers split out of +//! `executor.rs` stop threading a dozen positional parameters between them: +//! every one of those helpers takes `&RunCtx`/`&mut RunCtx` plus the handful +//! of values that are genuinely local to that call (the active set, the +//! state snapshot, a step's folded outcome). `RunCtx` is created once per +//! `execute_run` call and never outlives it — it borrows the owning +//! [`CompiledGraph`] for that duration. + +use super::*; + +use crate::observability::GraphStatusStore; + +/// Run-scoped state for one `execute_run` call. +/// +/// Fields fall into three groups: identity that never changes for the run +/// (`run_id`, `thread_id`, `root_run_id`, `parent_run_id`, `started_at`, +/// `live_frames`, `recursion_meta`, `binding`), accumulators the superstep +/// loop updates every iteration (`recursion`, `node_visits`, +/// `barrier_arrivals`, `resume_map`, `visited`, `all_child_runs`, `steps`, +/// `last_checkpoint`, `parent_checkpoint`), and I/O handles +/// (`child_sink`, `async_writes`). `graph` is the owning [`CompiledGraph`], +/// kept here so the convenience methods below (`emit`, `save_status`, +/// `base_status`, `node_context`) don't need a separate receiver. +pub(super) struct RunCtx<'a, State, Update> { + pub(super) graph: &'a CompiledGraph, + pub(super) run_id: RunId, + pub(super) thread_id: Option, + pub(super) root_run_id: RunId, + pub(super) parent_run_id: Option, + pub(super) started_at: SystemTime, + /// Monotonic start instant used for wall-clock-jump-proof deadline + /// arithmetic (M3): `run_deadline` is checked against + /// [`std::time::Instant::elapsed`] rather than [`SystemTime::elapsed`], + /// so a system clock step (NTP sync, VM pause/resume, manual clock + /// change) cannot make a run time out early or never at all. + /// `started_at` (above) remains the wall-clock stamp surfaced on + /// [`GraphRunStatus`], which is what observers expect. + pub(super) started_instant: std::time::Instant, + pub(super) live_frames: Vec, + pub(super) recursion_meta: serde_json::Value, + pub(super) recursion: RecursionStack, + pub(super) binding: Option, + pub(super) child_sink: ChildRunSink, + pub(super) node_visits: HashMap, + pub(super) barrier_arrivals: HashMap>, + pub(super) async_writes: AsyncCheckpointWrites, + /// Keyed by task id, falling back to node id (I1/R5); see + /// [`super::executor::RunSeed::resume_map`]. + pub(super) resume_map: HashMap, + /// Per-node snapshot of the channel versions as of the last time each + /// node ran (I5/R3), keyed by node id string. Loaded from the resumed + /// checkpoint's [`crate::checkpoint::Checkpoint::versions_seen`] + /// (`ResumeSeed::initial_versions_seen`); updated in [`Self::node_context`] + /// and persisted back onto every boundary checkpoint this run writes + /// (`compiled::boundary`). + pub(super) versions_seen: HashMap>, + pub(super) visited: Vec, + pub(super) all_child_runs: Vec, + pub(super) steps: usize, + pub(super) last_checkpoint: Option, + pub(super) parent_checkpoint: Option, + /// Nodes (with their persisted explicit `Command::goto`, R1) carried + /// forward from a resumed mid-step checkpoint (an interrupt/failure + /// boundary whose completed siblings were never routed) — see + /// [`super::boundary::CompiledGraph::advance`]'s doc. `None` for a fresh + /// run or a resume from a fully-routed (normal) boundary. Consumed + /// (`take`n) by the first `advance` call of this run; + /// [`super::boundary`]'s failure/interrupt boundaries read it (without + /// consuming it) to keep carrying it forward across a step that + /// interrupts or fails more than once in a row. + pub(super) carried_completed: Option)>>, + /// Deferred activations (a node whose [`NodePolicy::defer`] is set) held + /// back from the frontier while a non-deferred activation was also + /// ready, per [`super::boundary::CompiledGraph::apply_defer`]. Released + /// (all at once) the first time a boundary's routed frontier would + /// otherwise be empty — i.e. once nothing *else* is left to run. + /// + /// Not persisted in any checkpoint: a run resumed mid-way through a + /// deferred hold starts this back at empty, so a held deferred + /// activation does not survive a crash/resume. Real durability for + /// deferred scheduling is future work. + pub(super) deferred_pending: Vec, + /// Optional cooperative-cancellation token for this run (I4 part 2), from + /// [`super::RunOptions::cancellation`]. Checked at every superstep + /// boundary and raced against the step's in-flight node handlers by + /// [`super::executor::CompiledGraph::run_step_with_cancel`]. + pub(super) cancellation: Option, + /// Optional graceful-drain signal for this run, from + /// [`super::RunOptions::drain`]. Polled only between supersteps + /// (`execute_run`'s loop top): the step in flight always finishes. + pub(super) drain: Option, + /// Per-task replay memos loaded from the resumed checkpoint's write + /// ledger, keyed by task id: the task's + /// [`NodeContext::durable_task`] memo writes (pre-seeded onto its + /// `NodeContext` by [`Self::node_context`]) and, for an + /// `interrupt_after` pause, its deferred-result write (consumed by + /// [`Self::task_plan`]). Empty for a fresh run. + pub(super) task_writes: HashMap>, + /// Executor-injected interrupts (`interrupt_before`/`interrupt_after`) + /// the resumed checkpoint recorded, as `":"` keys: a + /// task listed here has already paused at that phase and must not be + /// paused there again when it re-runs. Empty for a fresh run. + pub(super) acknowledged_interrupts: HashSet, + /// Guards against the run future being dropped before it reaches a + /// normal terminal state (I4 part 3) — see [`RunDropGuard`]. + pub(super) drop_guard: RunDropGuard, +} + +/// What the step runner must do for one activation beyond invoking its +/// handler, resolved up front by [`RunCtx::task_plan`] so a branch future +/// borrows nothing from the run context. +pub(super) struct TaskPlan { + /// Pause before running the handler (`interrupt_before`, not yet + /// acknowledged for this task). + pub(super) inject_before: bool, + /// Pause after the handler completes, holding its result back + /// (`interrupt_after`, not yet acknowledged for this task). + pub(super) inject_after: bool, + /// The deferred result persisted by an earlier `interrupt_after` pause + /// of this task, to replay instead of running the handler. + pub(super) replay_after: Option, +} + +/// Drop guard (I4 part 3) that guarantees a run's terminal status is +/// durably set to `Cancelled` if the run's future is dropped before it +/// reaches a normal terminal state (completed, failed, interrupted, or an +/// explicit cooperative cancellation the executor already handled) — +/// for example when a caller wraps the run in `tokio::time::timeout` and the +/// deadline fires, or aborts the `JoinHandle` of the task the run was +/// spawned on. Without this guard such a drop leaves the run's last written +/// status stuck at `Running` forever, with nothing to signal that it will +/// never make further progress. +/// +/// Constructed armed by [`RunCtx::start`]; [`Self::disarm`] is called at the +/// top of every one of `execute_run`'s terminal exit paths — success +/// ([`super::executor::CompiledGraph::finish_run`]), an aborting error +/// ([`super::boundary::CompiledGraph::fail_and_return`]), a resumable +/// failure boundary +/// ([`super::boundary::CompiledGraph::handle_failure_boundary`]), an +/// interrupt boundary +/// ([`super::boundary::CompiledGraph::handle_interrupt_boundary`]), and an +/// explicit cooperative-cancellation boundary +/// ([`super::boundary::CompiledGraph::handle_cancel_boundary`]) — so a run +/// that reaches a real terminal state on its own never gets a spurious +/// `Cancelled` overwrite from `Drop` racing (or following) that path. +/// +/// # Best-effort guarantee +/// +/// `Drop::drop` cannot `.await`, so this guard cannot synchronously flush +/// in-flight [`AsyncCheckpointWrites`]. It instead spawns a detached +/// background task (via [`tokio::runtime::Handle::try_current`], a no-op +/// outside a tokio runtime) that persists a `Cancelled` [`GraphRunStatus`]; +/// this can still race a runtime shutdown that happens immediately after the +/// drop, in which case even this best-effort write may not land. Any +/// checkpoint write still in flight under `DurabilityMode::Async` is *not* +/// separately re-awaited by this guard — but it is not abandoned either: a +/// tokio `JoinHandle` being dropped only detaches it, it does not abort the +/// task, so the underlying `checkpointer.put`/`put_writes` call keeps +/// running to completion on its own regardless of whether `RunCtx` is still +/// alive to track it. What this guard cannot restore is the tracker's +/// ability to *observe* that write's outcome (the concern +/// [`AsyncCheckpointWrites`]'s own contract documents) — a write that fails +/// after the run future was dropped is only visible in the checkpointer +/// backend's own logs, not through `GraphRunStatus.error`. The one concrete, +/// testable contract this guard gives is: the run's stored status is never +/// left at `Running` forever. +pub(super) struct RunDropGuard { + armed: bool, + status_store: Option>, + run_id: RunId, + thread_id: Option, + graph_id: GraphId, + namespace: Vec, + started_at: SystemTime, +} + +impl RunDropGuard { + #[allow(clippy::too_many_arguments)] + fn new( + status_store: Option>, + run_id: RunId, + thread_id: Option, + graph_id: GraphId, + namespace: Vec, + started_at: SystemTime, + ) -> Self { + Self { + armed: true, + status_store, + run_id, + thread_id, + graph_id, + namespace, + started_at, + } + } + + /// Disarms the guard so a normal terminal exit does not also trigger the + /// `Drop`-time `Cancelled` write. + pub(super) fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for RunDropGuard { + fn drop(&mut self) { + if !self.armed { + return; + } + let Some(store) = self.status_store.take() else { + return; + }; + let Ok(handle) = tokio::runtime::Handle::try_current() else { + return; + }; + let run_id = self.run_id.clone(); + let thread_id = self.thread_id.clone(); + let graph_id = self.graph_id.clone(); + let namespace = std::mem::take(&mut self.namespace); + let started_at = self.started_at; + handle.spawn(async move { + let mut status = + GraphRunStatus::new(run_id.clone(), graph_id, ExecutionStatus::Cancelled); + status.thread_id = thread_id; + status.checkpoint_namespace = namespace; + status.started_at = started_at; + status.updated_at = SystemTime::now(); + status.ended_at = Some(SystemTime::now()); + if let Err(err) = store.put_status(status).await { + tracing::warn!( + "[graph:drop-guard] failed to persist cancelled status for run `{run_id}` \ + after its future was dropped before completion: {err}" + ); + } + }); + } +} + +/// Everything a resumed run seeds `RunCtx` with beyond a fresh run's +/// defaults, bundled into one optional parameter so [`RunCtx::start`] does +/// not grow a positional argument per resume-only field. +/// +/// A fresh run (`resume_from_inner` was never called) passes `None`, which +/// is equivalent to `ResumeSeed::default()`. +#[derive(Default)] +pub(super) struct ResumeSeed { + /// The loaded checkpoint's own step number (`to_metadata().step`), so + /// this run's `ctx.steps` continues counting up from it instead of + /// restarting at `0` — see the I3 finding in + /// `docs/runtime-comparison/code-review-graph.md`: without this, + /// `metadata.step` (and so `get_state_history`) goes non-monotonic + /// across a resume, and per-node visit caps + /// (`RecursionPolicy::max_visits_per_node`) reset every resume rather + /// than bounding the whole thread's lifetime. + pub(super) initial_steps: usize, + /// The loaded checkpoint's persisted `node_visits` metadata (see + /// [`super::boundary`]'s checkpoint builders), so per-node visit counts + /// accumulate across a resume instead of resetting. + pub(super) initial_node_visits: HashMap, + /// The loaded checkpoint's persisted `versions_seen` (I5/R3), so + /// per-node "have I already seen this channel change" bookkeeping + /// survives a resume instead of resetting. + pub(super) initial_versions_seen: HashMap>, + /// Nodes (with their persisted goto, R1) carried forward from a + /// mid-step (interrupt/failure) checkpoint whose completed siblings + /// were never routed — see [`RunCtx::carried_completed`]. + pub(super) carried_completed: Option)>>, + /// Per-task replay memos from the loaded checkpoint's write ledger — + /// see [`RunCtx::task_writes`]. + pub(super) task_writes: HashMap>, + /// Executor-injected interrupts the loaded checkpoint recorded — see + /// [`RunCtx::acknowledged_interrupts`]. + pub(super) acknowledged_interrupts: HashSet, +} + +impl<'a, State, Update> RunCtx<'a, State, Update> +where + State: Clone + Send + Sync + 'static, + Update: Send + 'static, +{ + /// Forwards to the owning graph's event sink (a no-op without one), + /// stamping the envelope with this run's id. + pub(super) fn emit(&self, event: GraphEvent) { + self.graph.emit(&self.run_id, event); + } + + /// Whether this run's cooperative-cancellation token (if any) has been + /// cancelled (I4 part 2). + pub(super) fn is_cancelled(&self) -> bool { + self.cancellation + .as_ref() + .is_some_and(tinyagents_harness::CancellationToken::is_cancelled) + } + + /// Whether this run's graceful-drain signal (if any) has been raised. + pub(super) fn is_drain_requested(&self) -> bool { + self.drain + .as_ref() + .is_some_and(super::DrainSignal::is_requested) + } + + /// The `acknowledged_interrupts` key for `phase` of `task_id`. + pub(super) fn interrupt_ack_key(phase: &str, task_id: &TaskId) -> String { + format!("{phase}:{}", task_id.as_str()) + } + + /// Resolves the executor-level interrupt handling for `activation`: + /// whether to pause before/after its handler (`interrupt_before` / + /// `interrupt_after` selectors, minus the phases this task already + /// acknowledged in the resumed checkpoint) and whether a deferred + /// `interrupt_after` result is waiting to be replayed instead of + /// running the handler at all. + pub(super) fn task_plan(&self, activation: &Activation) -> TaskPlan { + let node = &activation.node; + let acked = |phase: &str| { + self.acknowledged_interrupts + .contains(&Self::interrupt_ack_key(phase, &activation.task_id)) + }; + let after_acked = acked("after"); + let replay_after = if after_acked { + self.task_writes + .get(activation.task_id.as_str()) + .and_then(|writes| writes.iter().find(|w| w.is_interrupt_after())) + .map(|w| w.payload.clone()) + } else { + None + }; + TaskPlan { + inject_before: self.graph.interrupt_before.contains(node) && !acked("before"), + inject_after: self.graph.interrupt_after.contains(node) && !after_acked, + replay_after, + } + } + + /// Disarms this run's [`RunDropGuard`] — called at the top of every + /// terminal exit path of `execute_run` so a normal completion never + /// races a spurious `Cancelled` write from `Drop`. + pub(super) fn disarm_drop_guard(&mut self) { + self.drop_guard.disarm(); + } + + /// Forwards to the owning graph's status store (a no-op without one). + pub(super) async fn save_status(&self, status: GraphRunStatus) { + self.graph.save_status(status).await; + } + + /// Builds a fresh [`GraphRunStatus`] for this run at `Running` status, + /// stamped with this context's identity and start time. + pub(super) fn base_status(&self) -> GraphRunStatus { + self.graph + .base_status(&self.run_id, &self.thread_id, self.started_at) + } + + /// Builds this run's `RunCtx`: constructs the recursion stack from the + /// inherited parent frames and pushes the frame for this graph call (a + /// push that would exceed `max_depth` fails the run — emitting + /// `RunStarted` and a terminal `Failed` status — before any node + /// executes), then emits `RunStarted`/`RecursionDepthChanged` for a + /// successful push. + #[allow(clippy::too_many_arguments)] + pub(super) async fn start( + graph: &'a CompiledGraph, + run_id: RunId, + thread_id: Option, + resume_map: HashMap, + initial_barriers: HashMap>, + initial_parent: Option, + binding: Option, + resume_seed: ResumeSeed, + options: super::RunOptions, + ) -> Result { + let ResumeSeed { + initial_steps, + initial_node_visits, + initial_versions_seen, + carried_completed, + task_writes, + acknowledged_interrupts, + } = resume_seed; + let super::RunOptions { + cancellation, + drain, + } = options; + let started_at = SystemTime::now(); + let started_instant = std::time::Instant::now(); + // Graph-call depth (the stack) is tracked separately from node-loop + // visits (`node_visits`, below). + let mut recursion = + RecursionStack::with_frames(graph.recursion_frames.clone(), graph.recursion_policy); + let root_run_id = graph + .recursion_frames + .first() + .map(|f| f.run_id.clone()) + .unwrap_or_else(|| run_id.clone()); + let parent_run_id = graph.recursion_frames.last().map(|f| f.run_id.clone()); + let this_frame = RecursionFrame { + graph_id: graph.graph_id.clone(), + node_id: graph.recursion_node.clone(), + run_id: run_id.clone(), + task_id: None, + namespace: graph.namespace.clone(), + depth: recursion.depth(), + parent: parent_run_id.clone(), + }; + if let Err(err) = recursion.push(this_frame) { + graph.emit( + &run_id, + GraphEvent::RunStarted { + run_id: run_id.clone(), + }, + ); + graph + .fail_run(&run_id, &thread_id, started_at, 0, &err, None) + .await; + return Err(err); + } + // Serialized once per run for embedding in every checkpoint's metadata. + let recursion_meta = + serde_json::to_value(recursion.frames()).unwrap_or(serde_json::Value::Null); + let live_frames = recursion.frames().to_vec(); + let drop_guard = RunDropGuard::new( + graph.status_store.clone(), + run_id.clone(), + thread_id.clone(), + graph.graph_id.clone(), + graph.namespace.clone(), + started_at, + ); + + let ctx = Self { + graph, + run_id, + thread_id, + root_run_id, + parent_run_id, + started_at, + started_instant, + live_frames, + recursion_meta, + recursion, + binding, + child_sink: ChildRunSink::new(), + node_visits: initial_node_visits, + barrier_arrivals: initial_barriers, + async_writes: AsyncCheckpointWrites::default(), + resume_map, + versions_seen: initial_versions_seen, + visited: Vec::new(), + all_child_runs: Vec::new(), + steps: initial_steps, + last_checkpoint: None, + parent_checkpoint: initial_parent, + carried_completed, + deferred_pending: Vec::new(), + cancellation, + drain, + task_writes, + acknowledged_interrupts, + drop_guard, + }; + ctx.emit(GraphEvent::RunStarted { + run_id: ctx.run_id.clone(), + }); + // Surface this run's recursion depth so observers can attribute + // nested runs without reconstructing the tree from logs. + ctx.emit(GraphEvent::RecursionDepthChanged { + depth: ctx.recursion.depth(), + }); + Ok(ctx) + } + + /// Drains this step's child-run sink into `all_child_runs` and returns + /// its serialized form for embedding into this boundary's checkpoint + /// metadata. + pub(super) fn take_step_child_runs(&mut self) -> serde_json::Value { + let step_child_runs = self.child_sink.drain(); + self.all_child_runs.extend(step_child_runs.iter().cloned()); + serde_json::to_value(&step_child_runs).unwrap_or(serde_json::Value::Null) + } + + /// Builds the per-task [`NodeContext`] for `activation`, consuming its + /// entry from `resume_map` (a task can only be handed its resume value + /// once). + /// + /// `fork` carries the branch identity in a concurrent step (`None` in + /// sequential mode or single-node steps). `siblings` is the number of + /// activations of `activation.node` in this same step's active set + /// (I1): more than one means a `Send` fan-out of the same node, which is + /// what a subgraph node consults to namespace its child checkpoint by + /// task id instead of sharing one namespace across every fan-out branch. + /// + /// Resume lookup prefers `resume_map`'s task-id key (I1/R5: distinguishes + /// concurrent same-node activations) and falls back to the node-id key + /// (legacy/whole-node resume, or a resume value fanned across every + /// pending node with no interrupt provenance). + pub(super) fn node_context( + &mut self, + activation: &Activation, + step: usize, + fork: Option, + siblings: usize, + state: &State, + ) -> NodeContext { + let node_id = &activation.node; + let resume = self + .resume_map + .remove(activation.task_id.as_str()) + .or_else(|| self.resume_map.remove(node_id.as_str())); + // I5/R3: the channel versions this node's invocation observes are + // whatever the committed `state` reports right now (downcast to + // `ChannelState` when the graph uses the channel model; a plain + // whole-state graph reports nothing here — `changed_since_last_run` + // is only meaningful for a channel graph). This node's own + // last-observed snapshot (`versions_seen`) is recorded *before* + // being overwritten with the current one, so + // `NodeContext::changed_since_last_run` can compare "what I saw last + // time" against "what is current". + let current_versions = (state as &dyn std::any::Any) + .downcast_ref::() + .map(|cs| cs.channel_versions().clone()) + .unwrap_or_default(); + let key = node_id.to_string(); + let seen_before = self.versions_seen.get(&key).cloned().unwrap_or_default(); + self.versions_seen.insert(key, current_versions.clone()); + // Pre-seed the task's `durable_task` memos from the resumed + // checkpoint so a re-run hits instead of repeating the side effect. + let durable_writes: Vec = self + .task_writes + .get(activation.task_id.as_str()) + .map(|writes| { + writes + .iter() + .filter(|w| w.is_durable_task()) + .cloned() + .collect() + }) + .unwrap_or_default(); + NodeContext { + graph_id: self.graph.graph_id.clone(), + node_id: node_id.clone(), + run_id: self.run_id.clone(), + thread_id: self.thread_id.clone(), + step, + resume, + fork, + send_arg: activation.send_arg.clone(), + root_run_id: Some(self.root_run_id.clone()), + recursion_frames: self.live_frames.clone(), + child_runs: Some(self.child_sink.clone()), + agent_binding: self.binding.clone(), + task_id: activation.task_id.clone(), + siblings, + channel_versions: current_versions, + versions_seen: seen_before, + idle_clock: crate::builder::IdleClock::default(), + durable_writes: Arc::new(std::sync::Mutex::new(durable_writes)), + } + } +} diff --git a/crates/tinyagents-graph/src/compiled/state_api.rs b/crates/tinyagents-graph/src/compiled/state_api.rs index a721cbac..5c48f9b6 100644 --- a/crates/tinyagents-graph/src/compiled/state_api.rs +++ b/crates/tinyagents-graph/src/compiled/state_api.rs @@ -120,118 +120,224 @@ where })?; let parent_step = base.to_metadata().step; let parent_id = base.checkpoint_id.clone(); + + // The base checkpoint may itself be mid-step: an interrupt/failure + // boundary whose completed siblings were deferred rather than + // routed (see `boundary::advance`'s `carried_completed` doc, the C2 + // fix). Those node ids live in `base.completed_tasks` with no + // persisted `goto_map` entry of their own, so — exactly like + // `advance`'s carried-completed handling — they route via + // static/conditional edges only, not any `Command::goto` they might + // have returned. Routing them here (rather than silently dropping + // them) is what keeps a manual write from permanently losing a + // step's other branches the moment it touches a mid-step thread. + // Same `source == "loop"` gate as `resume::resume_from_inner`'s + // `mid_step` check, and for the same reason: an `update`-sourced + // checkpoint can carry `interrupted_nodes` forward for provenance + // (I2) without its `completed_tasks` representing owed routing — + // `update_state` always resolves every carried completion itself + // before writing. + let base_source_is_loop = base + .metadata + .get("source") + .and_then(serde_json::Value::as_str) + == Some("loop"); + let carried_completed: Vec = if base_source_is_loop + && (base.metadata.get("interrupted_nodes").is_some() + || base.metadata.get("failed_node").is_some()) + { + base.completed.iter().map(|c| c.node.clone()).collect() + } else { + Vec::new() + }; 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. + // attributed write (or a carried-forward completion routed here) + // records its own arrival into them. let mut arrivals = barriers_from_persisted(&base.barrier_arrivals); - // Pending schedule: the attributed node's successors *merged into* the - // base checkpoint's still-pending work, or the inherited set verbatim. - // - // `next_nodes` and `pending_activations` are derived from one merged - // activation list so they can never disagree — resume prefers the - // activations, so a node named by only one of them would be silently - // dropped (or re-scheduled without its `Send` arg). + // Pending schedule: the attributed node's successors and any + // carried-forward completions' successors, merged into the base + // checkpoint's still-pending work. // - // The merge is unconditional rather than a fallback for the - // nothing-was-scheduled case. `route(node, None, ..)` resolves a static - // or conditional edge, so today it yields at most one target and a - // withheld barrier is the only way to end up with none — but keying the - // merge on that would silently drop the untouched branches the moment a - // single call ever resolves a withheld target *and* a schedulable one. - let (next_nodes, pending_activations): (Vec, Option>) = - match &as_node { - Some(node) => { - // The attributed node counts as completed, so it leaves the - // schedule; every other branch the base checkpoint had in - // flight (with its `Send` arg, when it carried one) stays. - let mut merged: Vec = match &base.pending_activations { - Some(pending) if !pending.is_empty() => pending - .iter() - .map(Activation::from) - .filter(|activation| activation.node != *node) - .collect(), - // Checkpoints written before `pending_activations` - // existed only carry the node-id projection. - _ => base - .next_nodes - .iter() - .filter(|pending| *pending != node) - .cloned() - .map(Activation::node) - .collect(), - }; - let mut seen: HashSet = merged - .iter() - .filter(|activation| activation.send_arg.is_none()) - .map(|activation| activation.node.clone()) - .collect(); - 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. The barrier's other - // predecessors are still scheduled (they are part of - // `merged` above), so they run and clear the join. - 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) { - continue; - } - arrivals.remove(&tnode); - } - // `Send` activations may legitimately repeat a node - // (each carries its own arg); plain ones are - // deduplicated so a successor already pending is not - // scheduled twice. - let send_arg = target.send_arg().cloned(); - if send_arg.is_some() || seen.insert(tnode.clone()) { - merged.push(Activation { - node: tnode, - send_arg, - task_id: String::new(), - }); - } + // `base` was already normalized on read (every backend's decode path + // calls `Checkpoint::normalize`), so `base.tasks` is always the + // single source of truth here regardless of which format version the + // stored record was written in. + let mut merged: Vec = base + .tasks + .iter() + .map(Activation::from) + .filter(|activation| Some(&activation.node) != as_node.as_ref()) + .collect(); + let mut seen: HashSet = merged + .iter() + .filter(|activation| activation.send_arg.is_none()) + .map(|activation| activation.node.clone()) + .collect(); + // Routes one completed node's (static/conditional-only) successors + // into `merged`, applying the same barrier gate `route_completed` + // applies at the normal boundary. The merge is unconditional rather + // than a fallback for the nothing-was-scheduled case: `route(node, + // None, ..)` resolves a static or conditional edge, so today it + // yields at most one target and a withheld barrier is the only way + // to end up with none — but keying the merge on that would silently + // drop the untouched branches the moment a single call ever + // resolves a withheld target *and* a schedulable one. + let mut route_into_merged = |node: &NodeId| -> Result<()> { + for target in self.route(node, None, &new_state)? { + let tnode = target.node().clone(); + if tnode.as_str() == END { + continue; + } + 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) { + continue; } - let nodes = activation_nodes(&merged); - let activations = if merged.is_empty() { - None - } else { - Some(merged.iter().map(PendingActivation::from).collect()) - }; - (nodes, activations) + arrivals.remove(&tnode); + } + // `Send` activations may legitimately repeat a node (each + // carries its own arg); plain ones are deduplicated so a + // successor already pending is not scheduled twice. + let send_arg = target.send_arg().cloned().map(Arc::new); + if send_arg.is_some() || seen.insert(tnode.clone()) { + merged.push(Activation { + node: tnode, + send_arg, + task_id: TaskId::from(String::new()), + }); } - None => (base.next_nodes.clone(), base.pending_activations.clone()), - }; - let completed_tasks: Vec = as_node.iter().cloned().collect(); + } + Ok(()) + }; + for node in &carried_completed { + route_into_merged(node)?; + } + if let Some(node) = &as_node { + route_into_merged(node)?; + } + let tasks: Vec = merged.iter().map(PendingActivation::from).collect(); + // This write resolves every carried-forward completion's routing + // (above), so none of them are still "owed" afterward; only the + // attributed node (if any) is freshly completed by this write. + let completed: Vec = as_node + .iter() + .cloned() + .map(|node| crate::checkpoint::CompletedTask::new(TaskId::from(String::new()), node)) + .collect(); let barrier_arrivals = barriers_to_persisted(&arrivals); + // Per-task replay memos (`durable_task` writes, a deferred + // `interrupt_after` result) and executor-interrupt acknowledgements + // belong to the tasks that are still pending after this write, so + // carry exactly those forward; a task this write completed (or + // dropped) takes its memos with it. + let still_pending: HashSet<&str> = tasks + .iter() + .map(|t| t.task_id.as_str()) + .filter(|id| !id.is_empty()) + .collect(); + let pending_writes: Vec = base + .pending_writes + .iter() + .filter(|w| w.is_task_replay() && still_pending.contains(w.task_id.as_str())) + .cloned() + .collect(); + let carried_acks: Vec = base + .metadata + .get("acknowledged_interrupts") + .and_then(serde_json::Value::as_array) + .map(|acks| { + acks.iter() + .filter_map(serde_json::Value::as_str) + .filter(|key| { + key.split_once(':') + .is_some_and(|(_, task)| still_pending.contains(task)) + }) + .map(str::to_string) + .collect() + }) + .unwrap_or_default(); + + // I2: carry the base checkpoint's interrupt provenance through this + // manual write unless `as_node` names the node that actually + // interrupted — clearing it in exactly that case is how the + // documented "inspect -> update_state -> resume(value)" flow keeps + // working (`resume` fans the value across the pending set when it + // finds no interrupt provenance at all, so blindly erasing + // `interrupts`/`interrupted_nodes` on every manual write — as before + // this fix — handed the resume value to nodes that never paused). + let base_interrupted_stamped: Vec = base + .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(); + let names_interrupted_node = |node: &NodeId| -> bool { + if base_interrupted_stamped.is_empty() { + base.interrupts.iter().any(|i| &i.node == node) + } else { + base_interrupted_stamped.contains(node) + } + }; + let clears_interrupt = as_node.as_ref().is_some_and(names_interrupted_node); + let (interrupts, interrupted_nodes_meta) = if clears_interrupt { + (Vec::new(), Vec::new()) + } else { + (base.interrupts.clone(), base_interrupted_stamped) + }; let checkpoint_id = next_checkpoint_id(); let config = self.config_for(thread_id, Some(&checkpoint_id)); - let checkpoint = Checkpoint { - thread_id: thread_id.to_string(), - checkpoint_id, - run_id: None, - parent_checkpoint_id: Some(parent_id), - namespace: self.namespace.clone(), - state: new_state, - next_nodes, - completed_tasks, - pending_writes: Vec::new(), - interrupts: Vec::new(), - pending_activations, - barrier_arrivals, - metadata: serde_json::json!({ "source": "update", "step": parent_step + 1 }), - }; + let mut metadata = serde_json::json!({ "source": "update", "step": parent_step + 1 }); + if !interrupted_nodes_meta.is_empty() { + metadata["interrupted_nodes"] = serde_json::json!( + interrupted_nodes_meta + .iter() + .map(|n| n.to_string()) + .collect::>() + ); + } + if !carried_acks.is_empty() { + metadata["acknowledged_interrupts"] = serde_json::json!(carried_acks); + } + // I5/R3: the same `channel_bookkeeping` dispatch point the executor + // boundary uses (`compiled::boundary::channel_checkpoint_fields`), + // so a manual write can never disagree with a normal superstep + // boundary about what it persists here (the "one write path" + // contract). `versions_seen` (per-node) carries over from the base + // checkpoint unchanged — a manual write does not run any node. + let (channel_versions, channel_deltas) = + crate::channel::channel_bookkeeping(&new_state, parent_step as u64 + 1); + let checkpoint = Checkpoint::new(new_state, tasks) + .with_thread_id(thread_id.to_string()) + .with_checkpoint_id(checkpoint_id) + .with_parent_checkpoint_id(Some(parent_id)) + .with_namespace(self.namespace.clone()) + .with_completed(completed) + .with_pending_writes(pending_writes) + .with_interrupts(interrupts) + .with_barrier_arrivals(barrier_arrivals) + .with_channel_versions(channel_versions) + .with_channel_deltas(channel_deltas) + .with_versions_seen(base.versions_seen.clone()) + .with_metadata(metadata); + let writes = checkpoint.pending_writes.clone(); let id = checkpointer.put(checkpoint).await?; - self.emit(GraphEvent::CheckpointSaved { checkpoint_id: id }); + if !writes.is_empty() { + checkpointer.put_writes(&config, &writes).await?; + } + self.emit_unscoped(GraphEvent::CheckpointSaved { + checkpoint_id: id, + step: Some(parent_step + 1), + }); Ok(config) } @@ -283,23 +389,30 @@ where let step = source.to_metadata().step; let checkpoint_id = next_checkpoint_id(); let config = self.config_for(target_thread, Some(&checkpoint_id)); - let forked = Checkpoint { - thread_id: target_thread.to_string(), - checkpoint_id, - run_id: None, - parent_checkpoint_id: None, - namespace: source.namespace.clone(), - state: source.state.clone(), - next_nodes: source.next_nodes.clone(), - completed_tasks: source.completed_tasks.clone(), - pending_writes: source.pending_writes.clone(), - interrupts: source.interrupts.clone(), - pending_activations: source.pending_activations.clone(), - barrier_arrivals: source.barrier_arrivals.clone(), - metadata: serde_json::json!({ "source": "fork", "step": step }), - }; + // `source` was already normalized on read, so `.tasks`/`.completed` + // are the single source of truth regardless of the stored record's + // original format version. + // A fork copies the source checkpoint verbatim rather than writing — + // its channel bookkeeping is copied unchanged too, not bumped + // through `channel_bookkeeping` (there is no new write to account + // for). + let forked = Checkpoint::new(source.state.clone(), source.tasks.clone()) + .with_thread_id(target_thread.to_string()) + .with_checkpoint_id(checkpoint_id) + .with_namespace(source.namespace.clone()) + .with_completed(source.completed.clone()) + .with_pending_writes(source.pending_writes.clone()) + .with_interrupts(source.interrupts.clone()) + .with_barrier_arrivals(source.barrier_arrivals.clone()) + .with_channel_versions(source.channel_versions.clone()) + .with_channel_deltas(source.channel_deltas.clone()) + .with_versions_seen(source.versions_seen.clone()) + .with_metadata(serde_json::json!({ "source": "fork", "step": step })); let id = checkpointer.put(forked).await?; - self.emit(GraphEvent::CheckpointSaved { checkpoint_id: id }); + self.emit_unscoped(GraphEvent::CheckpointSaved { + checkpoint_id: id, + step: Some(step), + }); Ok(config) } } diff --git a/crates/tinyagents-graph/src/compiled/step.rs b/crates/tinyagents-graph/src/compiled/step.rs new file mode 100644 index 00000000..47d40c78 --- /dev/null +++ b/crates/tinyagents-graph/src/compiled/step.rs @@ -0,0 +1,1006 @@ +//! Running one superstep's active node set, and folding the results. +//! +//! This is the *execution* half of a superstep, split from the *boundary* +//! half (reducer apply, routing, checkpoint persist — see `boundary.rs`). +//! [`StepRunner`] drives the active node set's handlers, sequentially or +//! concurrently, and hands back a [`StepOutcome`] carrying every +//! `(Activation, Result>)` pair it actually produced. +//! [`StepRunner::fold_step`] then folds that outcome into a [`StepRun`]. +//! +//! Running and folding are deliberately kept as separate steps (rather than +//! folding inline as each branch completes, as the pre-split code did) so a +//! change to the fold policy touches only `fold_step`. Per the C1/C2 +//! findings in `docs/runtime-comparison/code-review-graph.md`, `fold_step` +//! now folds **every** `Ok` result regardless of its position in the active +//! set: a parallel step always drives every branch to completion +//! ([`StepRunner::run_parallel`]), so a higher-index branch that completed +//! before a lower-index one interrupted or failed must not be discarded and +//! re-run on resume. `fold_step` partitions the step's results into +//! `completed` (every branch that produced an `Update`/`Command`, in +//! original active-set-index order) and `stalled` (the branches that +//! errored or interrupted, which become the boundary's `pending` set) — +//! see [`StepRun`]. + +use super::*; + +use crate::cache::TaskCacheKey; +use crate::checkpoint::PendingWrite; +use crate::compiled::run_ctx::{RunCtx, TaskPlan}; + +/// One branch's settled outcome: the handler's (possibly executor-adjusted) +/// result plus the task's replay memos — its [`NodeContext::durable_task`] +/// writes and any deferred `interrupt_after` result — as they stood when it +/// settled. +type TaskOutput = (Result>, Vec); + +/// A boxed branch future; see [`StepRunner::run_parallel`] for why branches +/// are boxed behind a concrete `Send` bound. +type BranchFuture<'a, Update> = + std::pin::Pin> + Send + 'a>>; + +/// Builds the executor-injected interrupt an `interrupt_before` / +/// `interrupt_after` selector records for `node`: payload +/// `{"phase": "before"}` / `{"phase": "after"}`. The boundary stamps the +/// task id on it like any other interrupt. +fn injected_interrupt(node: &NodeId, phase: &str) -> Interrupt { + Interrupt::new(node.clone(), serde_json::json!({ "phase": phase })) +} + +/// Counts how many activations of this step's active set target each node +/// (I1): more than one is a `Send` fan-out of the same node, which +/// [`RunCtx::node_context`] surfaces on [`NodeContext::siblings`] so a +/// subgraph node handler can namespace its child checkpoint by task id +/// instead of sharing one namespace across every fan-out branch. +fn sibling_counts(active: &[Activation]) -> HashMap { + let mut counts: HashMap = HashMap::new(); + for activation in active { + *counts.entry(activation.node.clone()).or_insert(0) += 1; + } + counts +} + +/// The raw, unfolded result of running a superstep's active node set: one +/// `(Activation, Result)` pair per branch that was actually +/// invoked, in active-set index order. +/// +/// [`StepRunner::run_sequential`] stops invoking further branches at the +/// first error or interrupt (so `results` may be a strict prefix of the +/// active set); [`StepRunner::run_parallel`] always drives every branch to +/// completion first (so `results` always covers the whole active set). Ready +/// for [`StepRunner::fold_step`]. +pub(super) struct StepOutcome { + pub(super) results: Vec<(Activation, TaskOutput)>, +} + +/// The folded result of running a superstep's active node set, ready to +/// apply at the step boundary. +pub(super) struct StepRun { + /// Branch updates in deterministic active-set index order, from *every* + /// branch that produced one (an `Update` or a `Command` carrying one), + /// regardless of whether a lower-index sibling errored or interrupted. + pub(super) updates: Vec, + /// Explicit routing (plain `goto` nodes and/or [`Send`] packets) keyed by + /// the producing branch's active-set index. + /// + /// Keyed by index rather than node id so repeated [`Send`] activations of + /// the *same* node within a step (map-reduce fanout) each keep their own + /// [`Command::goto`] — a node-keyed map would let a later activation's + /// command clobber an earlier one's routing. + pub(super) goto_map: HashMap>, + /// Every branch that completed (produced an `Update`/`Command`, not an + /// error or interrupt), paired with its original active-set index — + /// needed so a later `route_completed` call can look its `goto_map` + /// entry back up by that same index. Superset of what the pre-C1/C2 fold + /// kept (the index-ascending prefix): a higher-index branch that + /// completed despite a lower-index sibling erroring/interrupting is + /// included here rather than dropped. + pub(super) completed: Vec<(usize, Activation)>, + /// Every branch that errored or interrupted this step, in ascending + /// original-index order — the boundary's `pending` set (re-run from + /// scratch on resume/retry). The first entry is always the branch named + /// by `interrupt`/`failure` below, when either is set. + pub(super) stalled: Vec<(usize, Activation)>, + /// Every branch that interrupted this step, active-set-index-paired, in + /// ascending index order (I1). Empty when nothing interrupted. Unlike + /// the pre-I1 fold (which surfaced only the lowest-index interrupt), + /// every interrupted branch is carried through to the boundary — a + /// `Send` fan-out of one node interrupting on every concurrent + /// activation surfaces all of them on + /// [`GraphExecution::interrupts`](super::GraphExecution), each stamped + /// with its own branch's task id. + pub(super) interrupted: Vec<(usize, Interrupt)>, + /// A node-handler failure that survived the node-retry policy, if any — + /// always the lowest-index error this step. When set, `updates` still + /// carries the updates of every branch that completed (not just those + /// with a lower index), so the executor can fold that partial progress + /// into committed state and persist a resumable failure boundary. + pub(super) failure: Option, + /// The replay memos of every *stalled* branch (durable-task writes plus + /// any deferred `interrupt_after` result), for the boundary to persist + /// next to the step's completion markers. A completed branch's memos are + /// dropped: the task will never re-run, so nothing needs replaying. + pub(super) task_writes: Vec, +} + +/// The two accumulators [`StepRunner::fold_result`] fills in as it walks a +/// step's results: branch updates and explicit routing. Bundled so +/// `fold_result` takes one accumulator instead of two separate `&mut` +/// parameters. +struct FoldAccum { + updates: Vec, + goto_map: HashMap>, +} + +/// The per-branch identity [`StepRunner::fold_result`] needs to emit its +/// events and update its accumulators. Bundled (rather than four separate +/// parameters) to keep `fold_result`'s signature small. +struct FoldBranch<'a> { + run_id: &'a RunId, + task_id: &'a TaskId, + index: usize, + node_id: &'a NodeId, + step: usize, +} + +/// Runs one superstep's active node set against a [`CompiledGraph`]. +/// +/// A thin wrapper around a `&CompiledGraph` borrow — it exists to give the +/// step-running/folding methods a home distinct from the boundary and +/// entry-point methods on `CompiledGraph` itself. +pub(super) struct StepRunner<'g, State, Update> { + pub(super) graph: &'g CompiledGraph, +} + +impl<'g, State, Update> StepRunner<'g, State, Update> +where + State: Clone + Send + Sync + 'static, + Update: Send + 'static, +{ + /// Wraps a node future in panic safety and the node's effective + /// timeouts (if any), mapping an elapsed deadline onto + /// [`TinyAgentsError::Timeout`]. + /// + /// Two independent ceilings race the handler: the flat `timeout` (max + /// wall time for the attempt, regardless of heartbeats) and the + /// `idle_timeout` (max gap between two [`NodeContext::heartbeat`] + /// calls, re-armed by each one via [`IdleClock::idle_elapsed`]). Either + /// firing first fails the attempt; a node that never heartbeats sees + /// its idle timeout fire exactly `idle_timeout` after start, i.e. as a + /// flat timeout. + /// + /// A node handler that panics unwinds through `join_all`/`fut.await` + /// unless caught here (I4 part 1): [`futures::FutureExt::catch_unwind`] + /// converts an unwind into an ordinary `Err`, so the panic flows through + /// the same failure boundary (checkpoint write, `RunFailed` event, status + /// `Failed`) as any other node error, instead of poisoning the whole run + /// future and leaving the status store stuck at `Running`. + async fn run_node_future( + &self, + node_id: &NodeId, + fut: NodeFuture, + policy: &NodePolicy, + idle_clock: &IdleClock, + ) -> Result> { + let node_id_owned = node_id.clone(); + let guarded = async move { + match futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(fut)).await { + Ok(result) => result, + Err(payload) => Err(Self::panic_error(&node_id_owned, payload)), + } + }; + let flat = async { + match policy.timeout { + Some(timeout) => tokio::time::sleep(timeout).await, + None => std::future::pending::<()>().await, + } + }; + let idle = async { + match policy.idle_timeout { + Some(idle) => idle_clock.idle_elapsed(idle).await, + None => std::future::pending::<()>().await, + } + }; + tokio::pin!(guarded); + tokio::select! { + result = &mut guarded => result, + _ = flat => Err(TinyAgentsError::Timeout(format!( + "node `{node_id}` exceeded its {:?} timeout", + policy.timeout.unwrap_or_default() + ))), + _ = idle => Err(TinyAgentsError::Timeout(format!( + "node `{node_id}` exceeded its {:?} idle timeout without a heartbeat", + policy.idle_timeout.unwrap_or_default() + ))), + } + } + + /// Extracts a printable message from a caught panic payload, preferring a + /// `&str` then a `String` downcast, and produces the + /// [`TinyAgentsError::Graph`] that stands in for the panic at the normal + /// failure boundary. + fn panic_error(node_id: &NodeId, payload: Box) -> TinyAgentsError { + let message = if let Some(s) = payload.downcast_ref::<&str>() { + (*s).to_string() + } else if let Some(s) = payload.downcast_ref::() { + s.clone() + } else { + "non-string panic payload".to_string() + }; + TinyAgentsError::Graph(format!("node `{node_id}` panicked: {message}")) + } + + /// Runs one node handler under the node's effective retry policy. + /// + /// Builds a fresh handler future (re-cloning the context, and — M2 — + /// sharing this step's `Arc` via `Arc::clone` rather than + /// re-cloning `State`) for each attempt, so a retried node re-runs from + /// its start — matching the durable execution model, where a node is + /// never suspended mid-flight. On a + /// [retryable][tinyagents_harness::retry::is_retryable] error, when + /// a [`RetryPolicy`](tinyagents_harness::retry::RetryPolicy) is + /// configured and permits another attempt, it emits + /// [`GraphEvent::NodeRetryScheduled`], sleeps the (opt-in) backoff, and + /// retries. Non-retryable errors, absence of a policy, or an exhausted + /// attempt budget return the error unchanged. The per-node timeout still + /// bounds every individual attempt via [`Self::run_node_future`]. + async fn run_node_with_retry( + &self, + node_id: &NodeId, + handler: &Arc>, + state: &Arc, + ctx: NodeContext, + step: usize, + policy: &NodePolicy, + ) -> Result> { + let mut attempt = 0usize; + loop { + // M2: every attempt shares this step's `Arc` via a cheap + // `Arc::clone` rather than re-cloning `State` itself — a retried + // activation no longer pays a fresh `State` clone per attempt. + let fut = handler(Arc::clone(state), ctx.clone()); + match self + .run_node_future(node_id, fut, policy, &ctx.idle_clock) + .await + { + Ok(result) => return Ok(result), + Err(error) => { + let retry_policy = policy + .retry + .as_ref() + .filter(|retry| retry.should_retry(attempt) && is_retryable(&error)); + if let Some(retry_policy) = retry_policy { + attempt += 1; + self.graph.emit_task( + &ctx.run_id, + Some(&ctx.task_id), + GraphEvent::NodeRetryScheduled { + node: node_id.clone(), + step, + attempt, + }, + ); + retry_policy.sleep_backoff(attempt).await; + continue; + } + // Retries (if any) are exhausted, or the error is not + // retryable at all: give `on_error` a last chance to + // recover the node's result before the error escalates. + if let Some(on_error) = policy.on_error.as_ref() + && let Some(command) = on_error(state.as_ref(), &error) + { + return Ok(NodeResult::Command(command)); + } + return Err(error); + } + } + } + } + + /// Runs one activation end to end under its [`TaskPlan`]: an + /// `interrupt_before` pause short-circuits to an injected interrupt + /// without touching the handler; a pending `interrupt_after` replay + /// decodes the deferred result instead of running the handler; otherwise + /// the handler runs under the node's retry policy and, for an + /// `interrupt_after` node, its result is deferred (see + /// [`Self::defer_result`]). Always returns the task's replay memos as + /// they stand afterwards, for the boundary to persist if the task + /// stalled. + #[allow(clippy::too_many_arguments)] + async fn run_task( + &self, + node_id: &NodeId, + handler: &Arc>, + state: &Arc, + ctx: NodeContext, + step: usize, + policy: &NodePolicy, + plan: TaskPlan, + ) -> TaskOutput { + if plan.inject_before { + return ( + Ok(NodeResult::Interrupt(injected_interrupt(node_id, "before"))), + ctx.durable_writes_snapshot(), + ); + } + if let Some(payload) = plan.replay_after { + return ( + self.replay_deferred_result(node_id, payload), + ctx.durable_writes_snapshot(), + ); + } + self.graph.emit_task( + &ctx.run_id, + Some(&ctx.task_id), + GraphEvent::NodeStarted { + node: node_id.clone(), + step, + }, + ); + self.graph.emit_task( + &ctx.run_id, + Some(&ctx.task_id), + GraphEvent::TaskStarted { + node: node_id.clone(), + step, + }, + ); + let memo_ctx = ctx.clone(); + let result = self + .run_node_with_retry(node_id, handler, state, ctx, step, policy) + .await; + let result = if plan.inject_after { + self.defer_result(node_id, &memo_ctx, result) + } else { + result + }; + (result, memo_ctx.durable_writes_snapshot()) + } + + /// Holds an `interrupt_after` node's completed result back from this + /// step: encodes its `Update` (if any) with the graph's + /// [`UpdateCodec`] and its `goto` as a + /// [`PendingWrite::interrupt_after`] memo on the task's buffer (so the + /// interrupt boundary persists it), and substitutes an injected + /// `{"phase": "after"}` interrupt as the branch result. An `Err`, or a + /// node-emitted interrupt, passes through untouched — the node did not + /// complete, so there is nothing to defer and no second pause. + fn defer_result( + &self, + node_id: &NodeId, + ctx: &NodeContext, + result: Result>, + ) -> Result> { + let (update, goto) = match result { + Ok(NodeResult::Update(update)) => (Some(update), Vec::new()), + Ok(NodeResult::Command(command)) => (command.update, command.goto), + other => return other, + }; + let codec = self.graph.update_codec.as_ref().ok_or_else(|| { + TinyAgentsError::Graph(format!( + "node `{node_id}` is an interrupt_after node but the graph has no Update codec" + )) + })?; + let encoded = match &update { + Some(update) => (codec.encode)(update).map_err(TinyAgentsError::Serialization)?, + None => serde_json::Value::Null, + }; + let payload = serde_json::json!({ "update": encoded, "goto": goto }); + ctx.lock_durable_writes() + .push(PendingWrite::interrupt_after( + node_id.clone(), + ctx.task_id.clone(), + payload, + )); + Ok(NodeResult::Interrupt(injected_interrupt(node_id, "after"))) + } + + /// Decodes a deferred `interrupt_after` result persisted by + /// [`Self::defer_result`] back into the `Command` the node originally + /// produced (update through the codec, `goto` verbatim), so the resumed + /// step applies it exactly as if the handler had just returned it. + fn replay_deferred_result( + &self, + node_id: &NodeId, + payload: serde_json::Value, + ) -> Result> { + let codec = self.graph.update_codec.as_ref().ok_or_else(|| { + TinyAgentsError::Graph(format!( + "node `{node_id}` has a deferred interrupt_after result but the graph has no \ + Update codec" + )) + })?; + let update = match payload.get("update") { + None | Some(serde_json::Value::Null) => None, + Some(value) => { + Some((codec.decode)(value.clone()).map_err(TinyAgentsError::Serialization)?) + } + }; + let goto: Vec = match payload.get("goto") { + None | Some(serde_json::Value::Null) => Vec::new(), + Some(value) => serde_json::from_value(value.clone())?, + }; + let mut command = Command::new(); + command.update = update; + command.goto = goto; + Ok(NodeResult::Command(command)) + } + + /// Computes the [`TaskCacheKey`] for `node_id`'s activation, when it has + /// a [`crate::NodeCachePolicy`] installed (via + /// [`crate::CompiledGraph::with_cached_node`]). The policy's key + /// function is called at most once per activation — its result is + /// reused for both the lookup and, on a miss, the store — since a key + /// function is documented to observe `send_arg` and may have caller- + /// visible side effects (see `cache_key_receives_send_arg_per_fanout_activation`). + fn cache_key_for( + &self, + node_id: &NodeId, + state: &State, + send_arg: Option<&serde_json::Value>, + ) -> Option { + let cached = self.graph.cached_nodes.get(node_id)?; + let hash = (cached.key)(state, send_arg); + Some(TaskCacheKey::new( + self.graph.graph_id.clone(), + node_id.clone(), + hash, + )) + } + + /// Looks up a live cache entry under `key`, when a + /// [`crate::cache::TaskCache`] backend is attached (via + /// [`crate::CompiledGraph::with_task_cache`]). + /// + /// A cache error, a missing entry, or a value that fails to decode into + /// `Update` are all treated as a miss (`None`) — caching is an + /// optimization, never a correctness requirement (see the module docs on + /// [`crate::cache::TaskCache`]). + async fn cache_get(&self, node_id: &NodeId, key: &TaskCacheKey) -> Option { + let cache = self.graph.task_cache.as_ref()?; + let value = cache.get(key).await.ok().flatten()?; + let cached = self.graph.cached_nodes.get(node_id)?; + (cached.decode)(value).ok() + } + + /// Synchronously encodes a cache-miss result for storage, without ever + /// awaiting — so nothing derived from `Update` (which is not necessarily + /// `Sync`) is ever live across a suspension point. Returns `None` for a + /// node with no cache policy, an error result, an interrupt, or a + /// `Command` with no update to store. + fn prepare_cache_put( + &self, + node_id: &NodeId, + result: &Result>, + ) -> Option<(serde_json::Value, Option)> { + let cached = self.graph.cached_nodes.get(node_id)?; + let result = result.as_ref().ok()?; + let update = match result { + NodeResult::Update(update) => Some(update), + NodeResult::Command(command) => command.update.as_ref(), + NodeResult::Interrupt(_) => None, + }?; + let value = (cached.encode)(update).ok()?; + Some((value, cached.ttl)) + } + + /// Writes a prepared cache-miss entry (see [`Self::prepare_cache_put`]) + /// under the activation's already-computed `key` (see + /// [`Self::cache_key_for`]) and emits [`GraphEvent::TaskCompleted`] + /// (`cached: false`). Takes only owned, unconditionally `Send + Sync` + /// values, so this is safe to await from a context that must itself stay + /// `Send` regardless of `Update`'s auto-trait bounds. + #[allow(clippy::too_many_arguments)] + async fn store_cache_entry( + &self, + run_id: &RunId, + task_id: &TaskId, + key: &TaskCacheKey, + value: serde_json::Value, + ttl: Option, + node_id: &NodeId, + step: usize, + ) { + let Some(cache) = self.graph.task_cache.as_ref() else { + return; + }; + let _ = cache.put(key, value, ttl).await; + self.graph.emit_task( + run_id, + Some(task_id), + GraphEvent::TaskCompleted { + node: node_id.clone(), + step, + cached: false, + }, + ); + } + + /// Runs one superstep's active node set — concurrently when the graph + /// opts into it (`with_parallel`) and more than one node is active, else + /// sequentially — and folds the result. This is the single entry point + /// `execute_run` calls per step. + pub(super) async fn run_step( + &self, + ctx: &mut RunCtx<'_, State, Update>, + active: &[Activation], + state: &State, + step: usize, + ) -> Result> { + // M2: clone `State` exactly once per superstep, into an `Arc` that + // every branch/attempt below shares via a cheap `Arc::clone` — the + // boundary above and below this call still deal in a plain `&State` + // (`RunCtx`/`boundary`/`executor` are unchanged), so this is the one + // place the per-attempt clone the review flagged is eliminated. + let state = Arc::new(state.clone()); + let outcome = if self.graph.parallel && active.len() > 1 { + self.run_parallel(ctx, active, &state, step).await? + } else { + self.run_sequential(ctx, active, &state, step).await? + }; + Ok(self.fold_step(&ctx.run_id, outcome, active, step, &mut ctx.visited)) + } + + /// Runs the active node set one node at a time (default behavior). + /// + /// Stops invoking further branches at the first error (the run aborts) + /// or interrupt (later nodes in the step are not started), exactly + /// preserving milestone-1 semantics: `outcome.results` ends at that + /// branch. + async fn run_sequential( + &self, + ctx: &mut RunCtx<'_, State, Update>, + active: &[Activation], + state: &Arc, + step: usize, + ) -> Result> { + let siblings = sibling_counts(active); + let mut results = Vec::with_capacity(active.len()); + for activation in active { + let node_id = &activation.node; + let node = self + .graph + .nodes + .get(node_id) + .ok_or_else(|| TinyAgentsError::MissingNode(node_id.to_string()))?; + + self.graph.emit_task( + &ctx.run_id, + Some(&activation.task_id), + GraphEvent::TaskScheduled { + node: node_id.clone(), + step, + }, + ); + + let cache_key = self.cache_key_for(node_id, state, activation.send_arg.as_deref()); + let cache_hit = match &cache_key { + Some(key) => self.cache_get(node_id, key).await, + None => None, + }; + let output = if let Some(update) = cache_hit { + self.graph.emit_task( + &ctx.run_id, + Some(&activation.task_id), + GraphEvent::TaskCompleted { + node: node_id.clone(), + step, + cached: true, + }, + ); + (Ok(NodeResult::Update(update)), Vec::new()) + } else { + let plan = ctx.task_plan(activation); + let node_ctx = ctx.node_context( + activation, + step, + None, + siblings.get(node_id).copied().unwrap_or(1), + state, + ); + let policy = self.graph.effective_policy(node_id); + let output = self + .run_task(node_id, &node.handler, state, node_ctx, step, &policy, plan) + .await; + if let (Some(key), Some((value, ttl))) = + (&cache_key, self.prepare_cache_put(node_id, &output.0)) + { + self.store_cache_entry( + &ctx.run_id, + &activation.task_id, + key, + value, + ttl, + node_id, + step, + ) + .await; + } + output + }; + let stop = matches!(output.0, Err(_) | Ok(NodeResult::Interrupt(_))); + results.push((activation.clone(), output)); + if stop { + break; + } + } + Ok(StepOutcome { results }) + } + + /// Runs the active node set concurrently (opt-in via `with_parallel`). + /// + /// Each branch executes against this step's shared `Arc` snapshot + /// (M2: every branch and retry attempt clones the `Arc`, not `State` + /// itself — see [`StepRunner::run_step`]) and a distinct [`ForkId`], + /// optionally with the [`Send`] argument that scheduled it. With no + /// `max_concurrency` bound every branch starts before any is awaited and + /// all are driven via [`futures::future::join_all`]; with a bound the + /// active set is run in chunks of at most that many futures, so at most + /// that many node handlers are in flight at once. Every branch is driven + /// to completion before this returns, regardless of whether an earlier + /// branch errored or interrupted — `outcome.results` always covers the + /// whole active set; [`Self::fold_step`] is what stops at the + /// lowest-index error/interrupt. + async fn run_parallel( + &self, + ctx: &mut RunCtx<'_, State, Update>, + active: &[Activation], + state: &Arc, + step: usize, + ) -> Result> { + // Build one forked context + future per branch. Node lookup and + // resume consumption happen up front so the futures borrow nothing + // mutable; each branch drives its handler through the node-retry + // policy (which also applies the per-node timeout), so a transient + // failure in one branch is retried without disturbing its siblings. + let siblings = sibling_counts(active); + let mut futures = Vec::with_capacity(active.len()); + // Parallel with a cache hit: `true` at `index` means that branch's + // slot in `futures` is an already-resolved replay, not a handler + // invocation — the post-loop pass below must not re-cache it (that + // would spuriously refresh its TTL on every hit). + let mut cache_hits = vec![false; active.len()]; + // The activation's cache key, computed once here and reused by the + // post-loop miss-store pass below (see `cache_key_for`'s doc on why + // the key function is called at most once per activation). + let mut cache_keys: Vec> = Vec::with_capacity(active.len()); + for (index, activation) in active.iter().enumerate() { + let node_id = &activation.node; + let node = self + .graph + .nodes + .get(node_id) + .ok_or_else(|| TinyAgentsError::MissingNode(node_id.to_string()))?; + + self.graph.emit_task( + &ctx.run_id, + Some(&activation.task_id), + GraphEvent::TaskScheduled { + node: node_id.clone(), + step, + }, + ); + + let cache_key = self.cache_key_for(node_id, state, activation.send_arg.as_deref()); + let cache_hit = match &cache_key { + Some(key) => self.cache_get(node_id, key).await, + None => None, + }; + cache_keys.push(cache_key); + if let Some(update) = cache_hit { + self.graph.emit_task( + &ctx.run_id, + Some(&activation.task_id), + GraphEvent::TaskCompleted { + node: node_id.clone(), + step, + cached: true, + }, + ); + cache_hits[index] = true; + let fut: BranchFuture<'_, Update> = + Box::pin(async move { (Ok(NodeResult::Update(update)), Vec::new()) }); + futures.push(fut); + continue; + } + + self.graph.emit_task( + &ctx.run_id, + Some(&activation.task_id), + GraphEvent::ContextForked { + node: node_id.clone(), + fork: index, + step, + }, + ); + + let plan = ctx.task_plan(activation); + let fork = Some(ForkId::new(index, node_id.clone())); + let node_ctx = ctx.node_context( + activation, + step, + fork, + siblings.get(node_id).copied().unwrap_or(1), + state, + ); + let handler = node.handler.clone(); + let owned_node = node_id.clone(); + let policy = self.graph.effective_policy(node_id); + // Box each branch future behind a concrete `Send` bound. This + // keeps the `select_all` rolling window below (used for a + // `max_concurrency` bound) from requiring a higher-ranked `Send` + // proof over the borrowed recursion frames, which the compiler + // cannot discharge for the bare `async` blocks. + let fut: BranchFuture<'_, Update> = Box::pin(async move { + self.run_task(&owned_node, &handler, state, node_ctx, step, &policy, plan) + .await + }); + futures.push(fut); + } + + // Drive branches to completion, bounding in-flight count when + // configured. With a bound, keep a rolling window of `limit` + // branches in flight instead of fixed `join_all` chunks. A chunked + // join runs each chunk to completion before starting the next, so a + // single slow branch head-of-line blocks the whole chunk; the + // rolling window starts a new branch as soon as *any* in-flight one + // finishes. `select_all` reports which pending future completed; a + // parallel index Vec maps it back to the branch's active-set + // position, so results are re-ordered into deterministic order for + // the fold below. + let results = match self.graph.max_concurrency { + Some(limit) if limit < futures.len() => { + let total = futures.len(); + let mut slots: Vec>> = (0..total).map(|_| None).collect(); + let mut source = futures.into_iter().enumerate(); + let mut running = Vec::with_capacity(limit); + let mut running_index = Vec::with_capacity(limit); + for (index, fut) in source.by_ref().take(limit) { + running.push(fut); + running_index.push(index); + } + while !running.is_empty() { + let (result, completed, rest) = futures::future::select_all(running).await; + let index = running_index.remove(completed); + slots[index] = Some(result); + running = rest; + if let Some((index, fut)) = source.next() { + running.push(fut); + running_index.push(index); + } + } + slots + .into_iter() + .map(|slot| slot.expect("every branch produced a result")) + .collect::>() + } + _ => futures::future::join_all(futures).await, + }; + + // Cache-miss branches (not the already-replayed hits above) store + // their result now that every branch has settled, reusing the key + // computed for each activation's lookup above. + for (index, activation) in active.iter().enumerate() { + if cache_hits[index] { + continue; + } + if let (Some(key), Some((value, ttl))) = ( + &cache_keys[index], + self.prepare_cache_put(&activation.node, &results[index].0), + ) { + self.store_cache_entry( + &ctx.run_id, + &activation.task_id, + key, + value, + ttl, + &activation.node, + step, + ) + .await; + } + } + + let results = active.iter().cloned().zip(results).collect::>(); + Ok(StepOutcome { results }) + } + + /// Folds a single successful branch result into the step accumulators. + /// + /// Pushes the node to `visited`, records updates/goto, emits the + /// matching events, and returns the interrupt (with its branch index) + /// when the branch paused. Returning `Some` means the branch did *not* + /// complete (it is a `stalled` branch, not a `completed` one) even + /// though it is not an `Err`. + fn fold_result( + &self, + branch: FoldBranch<'_>, + result: NodeResult, + accum: &mut FoldAccum, + visited: &mut Vec, + ) -> Option<(usize, Interrupt)> { + let FoldBranch { + run_id, + task_id, + index, + node_id, + step, + } = branch; + visited.push(node_id.clone()); + match result { + NodeResult::Update(update) => { + accum.updates.push(update); + self.graph.emit_task( + run_id, + Some(task_id), + GraphEvent::StateUpdated { + node: node_id.clone(), + step, + }, + ); + } + NodeResult::Command(command) => { + if let Some(update) = command.update { + accum.updates.push(update); + self.graph.emit_task( + run_id, + Some(task_id), + GraphEvent::StateUpdated { + node: node_id.clone(), + step, + }, + ); + } + if !command.goto.is_empty() { + accum.goto_map.insert(index, command.goto); + } + } + NodeResult::Interrupt(emitted) => { + self.graph.emit_task( + run_id, + Some(task_id), + GraphEvent::InterruptEmitted { + interrupt: emitted.clone(), + }, + ); + return Some((index, emitted)); + } + } + self.graph.emit_task( + run_id, + Some(task_id), + GraphEvent::NodeCompleted { + node: node_id.clone(), + step, + }, + ); + self.graph.emit_task( + run_id, + Some(task_id), + GraphEvent::TaskCompleted { + node: node_id.clone(), + step, + cached: false, + }, + ); + None + } + + /// Folds a [`StepOutcome`] into a [`StepRun`]. + /// + /// Per the module doc (C1/C2), this walks *every* result in + /// `outcome.results` — never stopping early — and partitions each + /// branch into `completed` (an `Update`/`Command` result) or `stalled` + /// (an error or an interrupt). The first error and the first interrupt + /// encountered (in ascending original-index order) are recorded as this + /// step's `failure`/`interrupt`; every stalled branch, including any + /// later error/interrupt beyond the first, still lands in `stalled` so + /// the boundary can schedule it for resume rather than silently + /// dropping it or mistaking it for completed. + /// + /// [`Self::run_sequential`] stops invoking further branches at the first + /// stop condition, so `outcome.results` may be a strict prefix of + /// `active` there; [`Self::run_parallel`] always drives the whole active + /// set, so `outcome.results` covers it completely. This is the + /// sequential-mode cousin of the C1 fix above `outcome.results` itself: + /// every `active` entry with no entry in `outcome.results` (because + /// `run_sequential` never started it) is a not-yet-started sibling of + /// the branch that stopped the step, and is folded into `stalled` here + /// too — using its *own* original active-set index, one past the + /// highest index `outcome.results` covers — so the boundary schedules it + /// as a pending task exactly like an errored/interrupted branch, instead + /// of silently dropping it from the checkpoint's pending set. Without + /// this, a sequential branch that interrupts or fails strands its + /// unstarted siblings: they never run on resume/retry, and the final + /// state permanently diverges from an uninterrupted run. + fn fold_step( + &self, + run_id: &RunId, + outcome: StepOutcome, + active: &[Activation], + step: usize, + visited: &mut Vec, + ) -> StepRun { + let mut accum = FoldAccum { + updates: Vec::new(), + goto_map: HashMap::new(), + }; + let mut completed: Vec<(usize, Activation)> = Vec::new(); + let mut stalled: Vec<(usize, Activation)> = Vec::new(); + let mut interrupted: Vec<(usize, Interrupt)> = Vec::new(); + let mut failure: Option = None; + let mut task_writes: Vec = Vec::new(); + + let ran = outcome.results.len(); + for (index, (activation, (result, writes))) in outcome.results.into_iter().enumerate() { + let node_id = activation.node.clone(); + match result { + Err(error) => { + self.graph.emit_task( + run_id, + Some(&activation.task_id), + GraphEvent::NodeFailed { + node: node_id, + step, + error: error.to_string(), + }, + ); + self.graph.emit_task( + run_id, + Some(&activation.task_id), + GraphEvent::TaskCompleted { + node: activation.node.clone(), + step, + cached: false, + }, + ); + if failure.is_none() { + failure = Some(StepFailure { + failed_index: index, + error, + }); + } + task_writes.extend(writes); + stalled.push((index, activation)); + } + Ok(result) => { + let branch = FoldBranch { + run_id, + task_id: &activation.task_id, + index, + node_id: &node_id, + step, + }; + match self.fold_result(branch, result, &mut accum, visited) { + Some(found) => { + interrupted.push(found); + task_writes.extend(writes); + stalled.push((index, activation)); + } + None => completed.push((index, activation)), + } + } + } + } + + // Sequential-mode sibling drop fix: any `active` entries beyond + // what `outcome.results` covers were never started this step + // (`run_sequential` stopped at the first error/interrupt). Carry + // them into `stalled` unexecuted, keyed by their own original + // active-set index, so they become pending tasks at the + // failure/interrupt boundary and are run exactly once on + // resume/retry rather than being dropped. + for (index, activation) in active.iter().enumerate().skip(ran) { + stalled.push((index, activation.clone())); + } + + StepRun { + updates: accum.updates, + goto_map: accum.goto_map, + completed, + stalled, + interrupted, + failure, + task_writes, + } + } +} diff --git a/crates/tinyagents-graph/src/compiled/test.rs b/crates/tinyagents-graph/src/compiled/test.rs index 1b5e52a7..6f4ec986 100644 --- a/crates/tinyagents-graph/src/compiled/test.rs +++ b/crates/tinyagents-graph/src/compiled/test.rs @@ -140,6 +140,44 @@ async fn conditional_routing_selects_branch() { assert_eq!(run.state, 100); } +#[tokio::test] +async fn static_edge_fan_out_activates_every_target() { + // `add_edge("start", "a").add_edge("start", "b")` must schedule BOTH "a" + // and "b" as successors of "start" in the same superstep (I10), not + // silently overwrite the first edge with the second. + let graph = GraphBuilder::, String>::new() + .set_reducer(ClosureStateReducer::new(|mut s: Vec, u: String| { + s.push(u); + Ok(s) + })) + .add_node("start", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update("start".to_string())) + }) + .add_node("a", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update("a".to_string())) + }) + .add_node("b", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update("b".to_string())) + }) + .set_entry("start") + .add_edge("start", "a") + .add_edge("start", "b") + .set_finish("a") + .set_finish("b") + .compile() + .unwrap(); + + let run = graph.run(vec![]).await.unwrap(); + assert_eq!(run.state, vec!["start", "a", "b"]); + assert_eq!( + run.visited + .iter() + .map(ToString::to_string) + .collect::>(), + vec!["start", "a", "b"] + ); +} + #[tokio::test] async fn command_goto_overrides_edges() { let graph = GraphBuilder::::overwrite() @@ -414,7 +452,7 @@ async fn resume_emits_restore_not_save_for_the_loaded_checkpoint() { assert!( !resume_events.iter().any(|e| matches!( e, - GraphEvent::CheckpointSaved { checkpoint_id } if *checkpoint_id == loaded + GraphEvent::CheckpointSaved { checkpoint_id, .. } if *checkpoint_id == loaded )), "loading a checkpoint on resume must not re-emit it as saved" ); @@ -740,6 +778,233 @@ async fn update_state_as_command_node_is_rejected() { .unwrap(); } +/// I2 regression: `update_state` (with `as_node: None`, so it does not touch +/// the interrupted node) followed by `resume(value)` must hand the resume +/// value only to the node that actually interrupted — not to every node the +/// checkpoint's pending set happens to carry, including one `update_state` +/// itself just scheduled via a carried-forward completion's routing. +/// Before the fix, `update_state` unconditionally wrote `interrupts: +/// Vec::new()` with no `interrupted_nodes` metadata, so `resume` found no +/// provenance and fanned the value across every pending node instead. +#[tokio::test] +async fn update_state_preserves_interrupt_provenance_for_a_later_resume() { + type ResumesSeen = Arc)>>>; + + let cp = Arc::new(InMemoryCheckpointer::::new()); + let resumes_seen: ResumesSeen = Arc::new(std::sync::Mutex::new(Vec::new())); + let lo_resumes = resumes_seen.clone(); + let y_resumes = resumes_seen.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(["lo", "hi"]), + )) + }) + .add_node("lo", move |_s: Counter, c: NodeContext| { + let seen = lo_resumes.clone(); + async move { + seen.lock() + .unwrap() + .push(("lo".to_string(), c.resume.clone())); + match c.resume { + Some(_) => Ok(NodeResult::Update(2)), + None => Ok(NodeResult::Interrupt(Interrupt::new("lo", json!({})))), + } + } + }) + .add_node("hi", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Update(20)) + }) + .add_node("y", move |_s: Counter, c: NodeContext| { + let seen = y_resumes.clone(); + async move { + seen.lock() + .unwrap() + .push(("y".to_string(), c.resume.clone())); + Ok(NodeResult::Update(5)) + } + }) + .set_entry("super") + .mark_command_routing("super") + .add_edge("hi", "y") + .set_finish("lo") + .set_finish("y") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let paused = graph + .run_with_thread( + "t-i2", + Counter { + value: 0, + log: vec![], + }, + ) + .await + .unwrap(); + assert!(paused.is_interrupted()); + + // A manual write with `as_node: None` — it must not clear the interrupt + // provenance. It also resolves `hi`'s deferred routing (a carried + // completion), scheduling `y` into the pending set alongside `lo`. + graph.update_state("t-i2", 0, None).await.unwrap(); + let mid = cp.get("t-i2", None).await.unwrap().unwrap(); + let mid_next_nodes: Vec = mid.tasks.iter().map(|t| t.node.clone()).collect(); + assert!( + mid_next_nodes.iter().any(|n| n.as_str() == "lo") + && mid_next_nodes.iter().any(|n| n.as_str() == "y"), + "both lo (still interrupted) and y (hi's deferred successor) must \ + be pending, got {mid_next_nodes:?}" + ); + + let resume_value = json!("only-for-lo"); + let done = graph + .resume("t-i2", Command::resume(resume_value.clone())) + .await + .unwrap(); + assert!(!done.is_interrupted()); + + let seen = resumes_seen.lock().unwrap(); + let lo_saw: Vec<_> = seen + .iter() + .filter(|(node, _)| node == "lo") + .map(|(_, r)| r.clone()) + .collect(); + let y_saw: Vec<_> = seen + .iter() + .filter(|(node, _)| node == "y") + .map(|(_, r)| r.clone()) + .collect(); + assert!( + lo_saw.iter().any(|r| *r == Some(resume_value.clone())), + "lo (the node that actually interrupted) must receive the resume value: {lo_saw:?}" + ); + assert!( + y_saw.iter().all(|r| r.is_none()), + "y (merely pending, never interrupted) must not receive the resume value: {y_saw:?}" + ); +} + +/// I3 regression: the checkpoint `step` metadata (and so +/// `get_state_history`) must stay monotonically increasing across a resume, +/// instead of restarting at `1`. Before the fix, `ctx.steps` was always +/// seeded at `0` in `RunCtx::start`, so a resumed run's boundaries +/// re-numbered from `1` again, making a checkpoint's `step` field disagree +/// with its position in the thread's actual lineage. +#[tokio::test] +async fn resume_continues_step_counter_monotonically() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let interrupted_once = Arc::new(AtomicBool::new(false)); + let flag = interrupted_once.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", move |s, c: NodeContext| { + let flag = flag.clone(); + async move { + if c.resume.is_none() && !flag.swap(true, AtomicOrdering::SeqCst) { + return Ok(NodeResult::Interrupt(Interrupt::new("b", json!({})))); + } + 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()); + + // Step 1: a. Step 2: b (interrupts). + let paused = graph.run_with_thread("t-i3-steps", 0).await.unwrap(); + assert!(paused.is_interrupted()); + assert_eq!(paused.status.current_step, 2); + + // Resumed: step 3 finishes b, step 4 runs c. + let done = graph + .resume("t-i3-steps", Command::resume(json!(null))) + .await + .unwrap(); + assert_eq!(done.state, 3, "a(+1) + b(+1) + c(+1)"); + + let history = graph.get_state_history("t-i3-steps", None).await.unwrap(); + let mut steps: Vec = history.iter().map(|snap| snap.metadata.step).collect(); + // History is newest-first; reverse to check monotonicity forward. + steps.reverse(); + for pair in steps.windows(2) { + assert!( + pair[1] > pair[0], + "step must be strictly increasing across the whole lineage, got {steps:?}" + ); + } + assert_eq!( + steps.last().copied(), + Some(4), + "the final boundary's step must continue from where the interrupt \ + left off (2), not restart at 1 after the resume, got {steps:?}" + ); +} + +/// I3 regression: `RecursionPolicy::max_visits_per_node` must bound a node's +/// visits across the whole thread's lifetime, not reset every resume. Before +/// the fix, `node_visits` was always seeded empty in `RunCtx::start`, so an +/// interrupt-then-resume loop could revisit a node past the configured limit +/// without ever tripping it. +#[tokio::test] +async fn resume_accumulates_node_visit_limit_across_resume() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let interrupted_once = Arc::new(AtomicBool::new(false)); + let flag = interrupted_once.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("loop", move |s, c: NodeContext| { + let flag = flag.clone(); + async move { + if c.resume.is_none() && !flag.swap(true, AtomicOrdering::SeqCst) { + return Ok(NodeResult::Interrupt(Interrupt::new("loop", json!({})))); + } + Ok(NodeResult::Update(s + 1)) + } + }) + .set_entry("loop") + .add_edge("loop", "loop") + .compile() + .unwrap() + .with_checkpointer(cp.clone()) + .with_recursion_policy(RecursionPolicy { + max_depth: 25, + max_visits_per_node: Some(2), + max_total_steps: 1000, + }); + + // Visit 1: interrupts (still within the limit of 2). + let paused = graph.run_with_thread("t-i3-visits", 0).await.unwrap(); + assert!(paused.is_interrupted()); + + // Visit 2 (post-resume): completes and self-loops, within the limit. + // Visit 3: must trip the *cumulative* limit of 2 — if node_visits reset + // on resume, this would incorrectly be seen as only the second visit. + let err = graph + .resume("t-i3-visits", Command::resume(json!(null))) + .await + .unwrap_err(); + assert!( + matches!(err, TinyAgentsError::NodeVisitLimit { limit: 2, .. }), + "got {err:?}" + ); +} + #[tokio::test] async fn bulk_update_state_applies_successive_updates() { use crate::CheckpointSource; @@ -1343,48 +1608,62 @@ async fn parallel_interrupt_schedules_completed_branch_successors() { assert_eq!(done.state.value, 111); } +/// R2/C1 regression: a higher-index parallel sibling that completed with a +/// visible side effect (`hi_calls`) must not be re-run when a lower-index +/// sibling interrupts and the thread is later resumed. Before the fix, +/// `fold_step` stopped folding at the first stalled branch by *position*, +/// so a completed higher-index branch was discarded and unconditionally +/// re-scheduled — a second call here would double the side effect and (for +/// a non-idempotent handler) double-apply its update. #[tokio::test] -async fn send_args_survive_interrupt_and_resume() { - // A `Send` fanout schedules three workers (args 1, 2, 3); the arg-1 worker - // interrupts on its first activation. On resume every pending worker must - // still carry its own send arg — before the fix they resumed with `None`. +async fn higher_index_completed_sibling_not_rerun_after_interrupt_then_resume() { let cp = Arc::new(InMemoryCheckpointer::::new()); + let hi_calls = Arc::new(AtomicUsize::new(0)); + let interrupted_once = Arc::new(AtomicBool::new(false)); + let hi_calls_for_node = hi_calls.clone(); + let interrupted_once_for_node = interrupted_once.clone(); let graph = GraphBuilder::::new() .with_parallel(true) .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { s.value += u; - s.log.push(format!("w:{u}")); + s.log.push(format!("+{u}")); Ok(s) })) - .add_node("dispatch", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Command(Command::send([ - Send::new("worker", json!(1)), - Send::new("worker", json!(2)), - Send::new("worker", json!(3)), - ]))) + .add_node("super", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Command( + Command::default().with_goto(["lo", "hi"]), + )) }) - .add_node("worker", |_s: Counter, c: NodeContext| async move { - let arg = c - .send_arg - .clone() - .expect("worker scheduled via Send must carry its arg") - .as_i64() - .unwrap() as i32; - if arg == 1 && c.resume.is_none() { - return Ok(NodeResult::Interrupt(Interrupt::new("worker", json!({})))); + .add_node("lo", move |_s: Counter, c: NodeContext| { + let once = interrupted_once_for_node.clone(); + async move { + match c.resume { + Some(_) => Ok(NodeResult::Update(2)), + None => { + once.store(true, AtomicOrdering::SeqCst); + Ok(NodeResult::Interrupt(Interrupt::new("lo", json!({})))) + } + } } - Ok(NodeResult::Update(arg)) }) - .set_entry("dispatch") - .mark_command_routing("dispatch") - .set_finish("worker") + .add_node("hi", move |_s: Counter, _c: NodeContext| { + let calls = hi_calls_for_node.clone(); + async move { + calls.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(20)) + } + }) + .set_entry("super") + .mark_command_routing("super") + .set_finish("lo") + .set_finish("hi") .compile() .unwrap() .with_checkpointer(cp.clone()); let paused = graph .run_with_thread( - "fan", + "t-higher-index", Counter { value: 0, log: vec![], @@ -1393,25 +1672,37 @@ async fn send_args_survive_interrupt_and_resume() { .await .unwrap(); assert!(paused.is_interrupted()); + assert!(interrupted_once.load(AtomicOrdering::SeqCst)); + // hi (index 1, higher than lo's index 0) still completed and its update + // applied, despite lo (lower index) interrupting the same step. + assert_eq!(hi_calls.load(AtomicOrdering::SeqCst), 1); + assert_eq!(paused.state.value, 20, "hi's update must be applied"); - // Resume: the arg-1 worker unblocks and the other two re-run with their - // preserved args. With the arg lost, `expect(...)` above would panic. let done = graph - .resume("fan", Command::resume(json!(null))) + .resume("t-higher-index", Command::resume(json!(null))) .await .unwrap(); - assert_eq!(done.state.value, 6, "all three worker args (1+2+3) applied"); - let mut log = done.state.log.clone(); - log.sort(); - assert_eq!(log, vec!["w:1", "w:2", "w:3"]); + assert_eq!( + hi_calls.load(AtomicOrdering::SeqCst), + 1, + "hi must not be re-run by the resume" + ); + assert_eq!(done.state.value, 22, "20 (hi) + 2 (lo's resume value)"); } +/// R1 regression: a carried-forward completed sibling's explicit +/// `Command::goto` must survive the interrupt + resume round trip. Before +/// the fix, `RouteTarget`/`Command::goto` were not serializable and +/// `Checkpoint` had nowhere to persist them, so a carried branch's routing +/// was re-resolved via static/conditional edges only on resume — silently +/// diverging from what an uninterrupted run would have routed to. #[tokio::test] -async fn barrier_arrivals_survive_interrupt_and_resume() { - // Diamond join: p1 arrives at the barrier before an interrupt; p2 arrives - // only after resume. The join must still fire — the p1 arrival has to - // survive the checkpoint boundary or the join's precondition is never met. +async fn carried_completed_sibling_goto_survives_resume() { let cp = Arc::new(InMemoryCheckpointer::::new()); + let x_calls = Arc::new(AtomicUsize::new(0)); + let y_calls = Arc::new(AtomicUsize::new(0)); + let x_calls_for_node = x_calls.clone(); + let y_calls_for_node = y_calls.clone(); let graph = GraphBuilder::::new() .with_parallel(true) .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { @@ -1421,39 +1712,49 @@ async fn barrier_arrivals_survive_interrupt_and_resume() { })) .add_node("super", |_s: Counter, _c: NodeContext| async move { Ok(NodeResult::Command( - Command::default().with_goto(["p1", "hold"]), + Command::default().with_goto(["lo", "hi"]), )) }) - .add_node("p1", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(1)) + .add_node("lo", |_s: Counter, c: NodeContext| async move { + match c.resume { + Some(_) => Ok(NodeResult::Update(2)), + None => Ok(NodeResult::Interrupt(Interrupt::new("lo", json!({})))), + } }) - .add_node("p2", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(2)) + // `hi` (the higher-index, already-completed sibling) has no static + // edge at all: it only reaches `x` via its explicit `Command::goto`. + // If that goto is lost across the interrupt boundary (the R1 bug), + // `hi` routes to nothing on resume and `x`/`y` never run. + .add_node("hi", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Command(Command::update(20).with_goto(["x"]))) }) - // `hold` interrupts first; on resume it routes to p2 (the second - // barrier predecessor). - .add_node("hold", |_s: Counter, c: NodeContext| async move { - match c.resume { - Some(_) => Ok(NodeResult::Command(Command::new().with_goto(["p2"]))), - None => Ok(NodeResult::Interrupt(Interrupt::new("hold", json!({})))), + .add_node("x", move |_s: Counter, _c: NodeContext| { + let calls = x_calls_for_node.clone(); + async move { + calls.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(100)) } }) - .add_node("join", |_s: Counter, _c: NodeContext| async move { - Ok(NodeResult::Update(100)) + .add_node("y", move |_s: Counter, _c: NodeContext| { + let calls = y_calls_for_node.clone(); + async move { + calls.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(1000)) + } }) .set_entry("super") .mark_command_routing("super") - .mark_command_routing("hold") - .add_waiting_edge("p1", "join") - .add_waiting_edge("p2", "join") - .set_finish("join") + .mark_command_routing("hi") + .set_finish("lo") + .set_finish("x") + .set_finish("y") .compile() .unwrap() .with_checkpointer(cp.clone()); let paused = graph .run_with_thread( - "diamond", + "t-carried-goto", Counter { value: 0, log: vec![], @@ -1463,180 +1764,711 @@ async fn barrier_arrivals_survive_interrupt_and_resume() { .unwrap(); assert!(paused.is_interrupted()); assert_eq!( - paused.state.value, 1, - "p1 committed (arrived at the barrier)" + paused.state.value, 20, + "hi's update committed before the pause" ); let done = graph - .resume("diamond", Command::resume(json!(null))) + .resume("t-carried-goto", Command::resume(json!(null))) .await .unwrap(); - assert!( - done.visited.iter().any(|n| n.as_str() == "join"), - "join must fire once both barrier predecessors have arrived across the resume" + + assert_eq!( + x_calls.load(AtomicOrdering::SeqCst), + 1, + "hi's persisted goto(\"x\") must run exactly once after resume" ); - // 1 (p1) + 2 (p2) + 100 (join). - assert_eq!(done.state.value, 103); + assert_eq!( + y_calls.load(AtomicOrdering::SeqCst), + 0, + "the static hi -> y edge must not fire once an explicit goto was persisted" + ); + // 20 (hi) + 2 (lo's resume value) + 100 (x) + assert_eq!(done.state.value, 122); } +/// R2/C2 regression: an interrupted-then-resumed run must reach the same +/// final state as the same graph run straight through, with each node +/// completing exactly once in both cases. Before the fix, a completed +/// sibling's successor was routed immediately at the interrupt boundary — +/// before the interrupted sibling's own eventual update was known — so a +/// downstream node could observe a state missing that update, an ordering +/// an uninterrupted run never produces. #[tokio::test] -async fn barrier_relief_fires_when_source_skips_relief_node() { - // Mixed fan-in: `m` waits on both `a` and `c`. `condition` never routes to - // `a` — it always takes the `skip` route to END, simulating an untaken - // conditional branch — so without a barrier relief `m` would deadlock - // forever waiting on a predecessor that never runs. - // `add_barrier_relief("condition", "a", "m")` registers `a`'s phantom - // arrival at `m` whenever `condition` completes without activating `a`, - // so `m` still fires once `c`'s real arrival lands. - let graph = GraphBuilder::, String>::new() - .with_parallel(true) - .set_reducer(ClosureStateReducer::new(|mut s: Vec, u: String| { - s.push(u); - Ok(s) - })) - .add_node("start", |_s, _c: NodeContext| async move { - Ok(NodeResult::Command( - Command::default().with_goto(["condition", "c"]), - )) - }) - .add_node("condition", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update("condition".to_string())) - }) - .add_node("a", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update("a".to_string())) - }) - .add_node("c", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update("c".to_string())) - }) - .add_node("m", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update("m".to_string())) +async fn interrupted_and_uninterrupted_runs_reach_the_same_state() { + fn build( + enable_interrupt: bool, + hi_completions: Arc, + lo_completions: Arc, + y_completions: Arc, + y_observed_value: Arc, + ) -> CompiledGraph { + 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(["lo", "hi"]), + )) + }) + .add_node("lo", move |_s: Counter, c: NodeContext| { + let completions = lo_completions.clone(); + async move { + if enable_interrupt && c.resume.is_none() { + return Ok(NodeResult::Interrupt(Interrupt::new("lo", json!({})))); + } + completions.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(2)) + } + }) + .add_node("hi", move |_s: Counter, _c: NodeContext| { + let completions = hi_completions.clone(); + async move { + completions.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(20)) + } + }) + .add_node("y", move |s: Counter, _c: NodeContext| { + let completions = y_completions.clone(); + let observed = y_observed_value.clone(); + async move { + completions.fetch_add(1, AtomicOrdering::SeqCst); + // The C2 property: `y` (the shared successor of both + // `hi` and `lo`) must observe a state that already + // includes *both* their updates, in either run — not + // just whichever of them completed first. + observed.store(s.value, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(5)) + } + }) + .set_entry("super") + .mark_command_routing("super") + .add_edge("hi", "y") + .set_finish("lo") + .set_finish("y") + .compile() + .unwrap() + } + + // Baseline: no interrupt, straight through. + let baseline_hi = Arc::new(AtomicUsize::new(0)); + let baseline_lo = Arc::new(AtomicUsize::new(0)); + let baseline_y = Arc::new(AtomicUsize::new(0)); + let baseline_y_observed = Arc::new(std::sync::atomic::AtomicI32::new(-1)); + let baseline_graph = build( + false, + baseline_hi.clone(), + baseline_lo.clone(), + baseline_y.clone(), + baseline_y_observed.clone(), + ); + let baseline = baseline_graph + .run(Counter { + value: 0, + log: vec![], }) - .set_entry("start") - .mark_command_routing("start") - // `condition` always takes the `skip` route to END — `a` is never - // reached via a real edge. - .add_conditional_edges( - "condition", - |_s: &Vec| "skip".to_string(), - [("skip", END)], - ) - .add_waiting_edge("a", "m") - .add_waiting_edge("c", "m") - .add_barrier_relief("condition", "a", "m") - .set_finish("m") - .compile() + .await .unwrap(); - let run = graph.run(Vec::new()).await.unwrap(); - + // Interrupted at `lo`, then resumed with the value it would otherwise + // have produced on its own. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let resumed_hi = Arc::new(AtomicUsize::new(0)); + let resumed_lo = Arc::new(AtomicUsize::new(0)); + let resumed_y = Arc::new(AtomicUsize::new(0)); + let resumed_y_observed = Arc::new(std::sync::atomic::AtomicI32::new(-1)); + let resumed_graph = build( + true, + resumed_hi.clone(), + resumed_lo.clone(), + resumed_y.clone(), + resumed_y_observed.clone(), + ) + .with_checkpointer(cp.clone()); + let paused = resumed_graph + .run_with_thread( + "t-equivalence", + Counter { + value: 0, + log: vec![], + }, + ) + .await + .unwrap(); + assert!(paused.is_interrupted()); + let resumed = resumed_graph + .resume("t-equivalence", Command::resume(json!(null))) + .await + .unwrap(); assert!( - run.visited.iter().any(|n| n.as_str() == "m"), - "m must activate via the barrier relief even though `a` never ran" + !resumed.is_interrupted(), + "the resume must carry `lo` past its interrupt check, not pause it again" ); - assert!( - !run.visited.iter().any(|n| n.as_str() == "a"), - "a must never have run (condition always skipped it)" + + // The reducer fan-in *order* of `lo` vs `hi` is allowed to differ (`hi` + // is folded into the original step's state; `lo`'s update lands one + // superstep later, once it actually completes on resume) — durable + // execution never suspends mid-superstep, so an interrupted branch's + // update necessarily commits later than an uninterrupted run's would. + // What must be identical is the *value* every node's update commits + // (the multiset of applied updates) and the final merged state's sum: + // `y`, the successor of both, must see both updates either way (the C2 + // property) — not just whichever completed first. + assert_eq!( + resumed.state.value, baseline.state.value, + "an interrupted-then-resumed run must reach the same final summed \ + state as the same graph run straight through" ); + let mut resumed_log = resumed.state.log.clone(); + let mut baseline_log = baseline.state.log.clone(); + resumed_log.sort(); + baseline_log.sort(); assert_eq!( - run.state, - vec!["condition".to_string(), "c".to_string(), "m".to_string()], - "m fires off condition+c's real contributions, with no phantom `a` update" + resumed_log, baseline_log, + "every node's update must be applied exactly once in both runs, \ + regardless of fan-in order" + ); + // The core C2 property: in both runs, `y` observed a state that already + // included both `hi`'s (20) and `lo`'s (2) updates. + assert_eq!( + baseline_y_observed.load(AtomicOrdering::SeqCst), + 22, + "baseline: y must observe both hi's and lo's updates" + ); + assert_eq!( + resumed_y_observed.load(AtomicOrdering::SeqCst), + 22, + "resumed: y must observe both hi's and lo's updates, not just hi's" ); + // Every node completed exactly once in both runs — no double-execution + // and no missing execution introduced by the interrupt/resume path. + assert_eq!(baseline_hi.load(AtomicOrdering::SeqCst), 1); + assert_eq!(resumed_hi.load(AtomicOrdering::SeqCst), 1); + assert_eq!(baseline_lo.load(AtomicOrdering::SeqCst), 1); + assert_eq!(resumed_lo.load(AtomicOrdering::SeqCst), 1); + assert_eq!(baseline_y.load(AtomicOrdering::SeqCst), 1); + assert_eq!(resumed_y.load(AtomicOrdering::SeqCst), 1); } +/// R2/C1 regression, failure/`retry` variant: a higher-index parallel +/// sibling that completed must not be re-run when a lower-index sibling +/// fails (survives no retry policy, so the run aborts with a resumable +/// failure-boundary checkpoint) and the thread is later retried. #[tokio::test] -async fn reducer_error_at_boundary_transitions_run_to_failed() { - // A reducer error raised at the step boundary (after the node ran) must - // still fail the run — emit RunFailed / a Failed status — rather than - // unwinding and leaving observers to see the run stuck in Running. - let sink = Arc::new(CollectingSink::new()); - let graph = GraphBuilder::::new() - .set_reducer(ClosureStateReducer::new(|_s: i32, u: i32| { - if u == 999 { - Err(TinyAgentsError::Graph("reducer boom".to_string())) - } else { - Ok(u) - } +async fn higher_index_completed_sibling_not_rerun_after_failure_then_retry() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let hi_calls = Arc::new(AtomicUsize::new(0)); + let failed_once = Arc::new(AtomicBool::new(false)); + let hi_calls_for_node = hi_calls.clone(); + let failed_once_for_node = failed_once.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("boom", |_s, _c: NodeContext| async move { - Ok(NodeResult::Update(999)) + .add_node("super", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Command( + Command::default().with_goto(["lo", "hi"]), + )) }) - .set_entry("boom") - .set_finish("boom") + .add_node("lo", move |_s: Counter, _c: NodeContext| { + let once = failed_once_for_node.clone(); + async move { + if once.swap(true, AtomicOrdering::SeqCst) { + Ok(NodeResult::Update(2)) + } else { + Err(TinyAgentsError::Graph("transient boom".to_string())) + } + } + }) + .add_node("hi", move |_s: Counter, _c: NodeContext| { + let calls = hi_calls_for_node.clone(); + async move { + calls.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(20)) + } + }) + .set_entry("super") + .mark_command_routing("super") + .set_finish("lo") + .set_finish("hi") .compile() .unwrap() - .with_event_sink(sink.clone()); + .with_checkpointer(cp.clone()); - let err = graph.run(0).await.unwrap_err(); - assert!(matches!(err, TinyAgentsError::Graph(_)), "got {err:?}"); - assert!( - sink.events() + let failed = graph + .run_with_thread( + "t-higher-index-fail", + Counter { + value: 0, + log: vec![], + }, + ) + .await; + assert!(failed.is_err(), "lo's failure must abort the run"); + assert_eq!(hi_calls.load(AtomicOrdering::SeqCst), 1); + + let checkpoint = cp + .get("t-higher-index-fail", None) + .await + .unwrap() + .expect("a resumable failure-boundary checkpoint must be persisted"); + assert_eq!( + checkpoint + .completed .iter() - .any(|e| matches!(e, GraphEvent::RunFailed { .. })), - "a boundary reducer error must transition the run to Failed (RunFailed emitted)" + .map(|c| c.node.clone()) + .collect::>(), + vec![NodeId::from("hi")], + "hi's completion must be recorded so retry does not re-run it" + ); + + let done = graph.retry("t-higher-index-fail").await.unwrap(); + assert_eq!( + hi_calls.load(AtomicOrdering::SeqCst), + 1, + "hi must not be re-run by retry" ); + assert_eq!(done.state.value, 22, "20 (hi) + 2 (lo, on retry)"); } +/// Sequential-mode cousin of the parallel C1 sibling-drop bug: in the +/// default sequential step runner, `run_sequential` stops invoking further +/// branches at the first interrupt, so a not-yet-started sibling of the +/// interrupting branch never even appears in that step's raw results. +/// Before the `fold_step` fix, such a sibling silently vanished from the +/// checkpoint's pending set (`Checkpoint::tasks`) instead of being carried +/// over, so it never ran on resume. This was pinned by a probe +/// (`drain_test::probe_sequential_stall_keeps_unstarted_siblings_pending`) +/// that has since been converted into this real assertion. #[tokio::test] -async fn status_snapshot_reports_run() { - let graph = adding_graph(); - let run = graph - .run(Counter { - value: 0, - log: vec![], +async fn sequential_stall_keeps_unstarted_sibling_pending() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let c_calls = Arc::new(AtomicUsize::new(0)); + let c_calls_for_node = c_calls.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s)) + }) + .add_node("b", |_s, c: NodeContext| async move { + match c.resume { + Some(_) => Ok(NodeResult::Update(10)), + None => Ok(NodeResult::Interrupt(Interrupt::new("b", json!({})))), + } + }) + .add_node("c", move |s, _c: NodeContext| { + let calls = c_calls_for_node.clone(); + async move { + calls.fetch_add(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(s + 1)) + } }) + .set_entry("a") + .add_edge("a", "b") + .add_edge("a", "c") + .set_finish("b") + .set_finish("c") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let paused = graph.run_with_thread("seq-sibling", 0).await.unwrap(); + assert!(paused.is_interrupted()); + // `c` never ran: `b` (index 0) interrupted before `run_sequential` ever + // started `c` (index 1). + assert_eq!(c_calls.load(AtomicOrdering::SeqCst), 0); + + let snapshot = graph + .get_state("seq-sibling", None) + .await + .unwrap() + .expect("an interrupted thread has a resumable checkpoint"); + let mut pending: Vec = snapshot.next_nodes.iter().map(|n| n.to_string()).collect(); + pending.sort(); + assert_eq!( + pending, + vec!["b".to_string(), "c".to_string()], + "the unstarted sibling `c` must be carried into the checkpoint's pending set \ + alongside the interrupted `b`, not dropped" + ); + + let done = graph + .resume("seq-sibling", Command::resume(json!(null))) .await .unwrap(); - let status = &run.status; - assert_eq!(status.status, ExecutionStatus::Completed); - assert_eq!(status.current_step, 2); - assert!(status.ended_at.is_some()); - assert!(status.error.is_none()); - assert_eq!(status.graph_id, *graph.graph_id()); + assert!(!done.is_interrupted()); + assert_eq!( + c_calls.load(AtomicOrdering::SeqCst), + 1, + "c must run exactly once, on resume" + ); + assert_eq!( + done.state, 1, + "overwrite reducer: c (index 1) applied last, overwriting b's 10" + ); } -/// A `Send` fan-out delivers a distinct per-branch argument to N parallel -/// activations of the *same* node, and the reducer merges their results. +/// Equivalence regression for the fix above: an interrupted-then-resumed +/// sequential run must reach the exact same final state as an uninterrupted +/// run of the same graph, and every node must run exactly once either way. #[tokio::test] -async fn send_fanout_delivers_distinct_args_to_parallel_branches() { +async fn sequential_interrupted_then_resumed_matches_uninterrupted_run() { + fn build( + cp: Arc>, + interrupt_once: Arc, + ) -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node("a", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .add_node("b", move |s, c: NodeContext| { + let interrupt_once = interrupt_once.clone(); + async move { + if c.resume.is_none() && !interrupt_once.swap(true, AtomicOrdering::SeqCst) { + return Ok(NodeResult::Interrupt(Interrupt::new("b", json!({})))); + } + Ok(NodeResult::Update(s + 10)) + } + }) + .add_node("c", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 100)) + }) + .set_entry("a") + .add_edge("a", "b") + .add_edge("a", "c") + .set_finish("b") + .set_finish("c") + .compile() + .unwrap() + .with_checkpointer(cp) + } + + // Baseline: no interrupt ever fires (the flag starts pre-tripped), so + // this is an ordinary uninterrupted sequential run. + let baseline_cp = Arc::new(InMemoryCheckpointer::::new()); + let baseline = build(baseline_cp, Arc::new(AtomicBool::new(true))); + let baseline_run = baseline.run_with_thread("baseline", 0).await.unwrap(); + assert!(!baseline_run.is_interrupted()); + + // Interrupted variant: `b` interrupts on its first (non-resume) call, + // stranding unstarted sibling `c`; resuming must reach the same state. + let interrupted_cp = Arc::new(InMemoryCheckpointer::::new()); + let interrupted = build(interrupted_cp, Arc::new(AtomicBool::new(false))); + let paused = interrupted.run_with_thread("interrupted", 0).await.unwrap(); + assert!(paused.is_interrupted()); + let done = interrupted + .resume("interrupted", Command::resume(json!(null))) + .await + .unwrap(); + assert!(!done.is_interrupted()); + + assert_eq!( + done.state, baseline_run.state, + "an interrupted-then-resumed sequential run must reach the same final \ + state as an uninterrupted run" + ); +} + +#[tokio::test] +async fn send_args_survive_interrupt_and_resume() { + // A `Send` fanout schedules three workers (args 1, 2, 3); the arg-1 worker + // interrupts on its first activation. On resume every pending worker must + // still carry its own send arg — before the fix they resumed with `None`. + 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!("worker:{u}")); + s.log.push(format!("w:{u}")); Ok(s) })) - .with_parallel(true) - // dispatch fans out three custom inputs to the same worker node. .add_node("dispatch", |_s: Counter, _c: NodeContext| async move { Ok(NodeResult::Command(Command::send([ - Send::new("worker", json!(10)), - Send::new("worker", json!(20)), - Send::new("worker", json!(30)), + Send::new("worker", json!(1)), + Send::new("worker", json!(2)), + Send::new("worker", json!(3)), ]))) }) - // each worker invocation consumes its own send arg as the update. .add_node("worker", |_s: Counter, c: NodeContext| async move { let arg = c .send_arg - .expect("worker scheduled via Send carries an arg"); - let v = arg.as_i64().unwrap() as i32; - Ok(NodeResult::Update(v)) + .clone() + .expect("worker scheduled via Send must carry its arg") + .as_i64() + .unwrap() as i32; + if arg == 1 && c.resume.is_none() { + return Ok(NodeResult::Interrupt(Interrupt::new("worker", json!({})))); + } + Ok(NodeResult::Update(arg)) }) - .mark_command_routing("dispatch") .set_entry("dispatch") + .mark_command_routing("dispatch") .set_finish("worker") .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let paused = graph + .run_with_thread( + "fan", + Counter { + value: 0, + log: vec![], + }, + ) + .await .unwrap(); + assert!(paused.is_interrupted()); - let run = graph - .run(Counter { - value: 0, - log: vec![], - }) + // Resume: the arg-1 worker unblocks and the other two re-run with their + // preserved args. With the arg lost, `expect(...)` above would panic. + let done = graph + .resume("fan", Command::resume(json!(null))) .await .unwrap(); + assert_eq!(done.state.value, 6, "all three worker args (1+2+3) applied"); + let mut log = done.state.log.clone(); + log.sort(); + assert_eq!(log, vec!["w:1", "w:2", "w:3"]); +} - // All three distinct args merged: 10 + 20 + 30. +#[tokio::test] +async fn barrier_arrivals_survive_interrupt_and_resume() { + // Diamond join: p1 arrives at the barrier before an interrupt; p2 arrives + // only after resume. The join must still fire — the p1 arrival has to + // survive the checkpoint boundary or the join's precondition is never met. + 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(["p1", "hold"]), + )) + }) + .add_node("p1", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Update(1)) + }) + .add_node("p2", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Update(2)) + }) + // `hold` interrupts first; on resume it routes to p2 (the second + // barrier predecessor). + .add_node("hold", |_s: Counter, c: NodeContext| async move { + match c.resume { + Some(_) => Ok(NodeResult::Command(Command::new().with_goto(["p2"]))), + None => Ok(NodeResult::Interrupt(Interrupt::new("hold", json!({})))), + } + }) + .add_node("join", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Update(100)) + }) + .set_entry("super") + .mark_command_routing("super") + .mark_command_routing("hold") + .add_waiting_edge("p1", "join") + .add_waiting_edge("p2", "join") + .set_finish("join") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let paused = graph + .run_with_thread( + "diamond", + Counter { + value: 0, + log: vec![], + }, + ) + .await + .unwrap(); + assert!(paused.is_interrupted()); + assert_eq!( + paused.state.value, 1, + "p1 committed (arrived at the barrier)" + ); + + let done = graph + .resume("diamond", Command::resume(json!(null))) + .await + .unwrap(); + assert!( + done.visited.iter().any(|n| n.as_str() == "join"), + "join must fire once both barrier predecessors have arrived across the resume" + ); + // 1 (p1) + 2 (p2) + 100 (join). + assert_eq!(done.state.value, 103); +} + +#[tokio::test] +async fn barrier_relief_fires_when_source_skips_relief_node() { + // Mixed fan-in: `m` waits on both `a` and `c`. `condition` never routes to + // `a` — it always takes the `skip` route to END, simulating an untaken + // conditional branch — so without a barrier relief `m` would deadlock + // forever waiting on a predecessor that never runs. + // `add_barrier_relief("condition", "a", "m")` registers `a`'s phantom + // arrival at `m` whenever `condition` completes without activating `a`, + // so `m` still fires once `c`'s real arrival lands. + let graph = GraphBuilder::, String>::new() + .with_parallel(true) + .set_reducer(ClosureStateReducer::new(|mut s: Vec, u: String| { + s.push(u); + Ok(s) + })) + .add_node("start", |_s, _c: NodeContext| async move { + Ok(NodeResult::Command( + Command::default().with_goto(["condition", "c"]), + )) + }) + .add_node("condition", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update("condition".to_string())) + }) + .add_node("a", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update("a".to_string())) + }) + .add_node("c", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update("c".to_string())) + }) + .add_node("m", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update("m".to_string())) + }) + .set_entry("start") + .mark_command_routing("start") + // `condition` always takes the `skip` route to END — `a` is never + // reached via a real edge. + .add_conditional_edges( + "condition", + |_s: &Vec| "skip".to_string(), + [("skip", END)], + ) + .add_waiting_edge("a", "m") + .add_waiting_edge("c", "m") + .add_barrier_relief("condition", "a", "m") + .set_finish("m") + .compile() + .unwrap(); + + let run = graph.run(Vec::new()).await.unwrap(); + + assert!( + run.visited.iter().any(|n| n.as_str() == "m"), + "m must activate via the barrier relief even though `a` never ran" + ); + assert!( + !run.visited.iter().any(|n| n.as_str() == "a"), + "a must never have run (condition always skipped it)" + ); + assert_eq!( + run.state, + vec!["condition".to_string(), "c".to_string(), "m".to_string()], + "m fires off condition+c's real contributions, with no phantom `a` update" + ); +} + +#[tokio::test] +async fn reducer_error_at_boundary_transitions_run_to_failed() { + // A reducer error raised at the step boundary (after the node ran) must + // still fail the run — emit RunFailed / a Failed status — rather than + // unwinding and leaving observers to see the run stuck in Running. + let sink = Arc::new(CollectingSink::new()); + let graph = GraphBuilder::::new() + .set_reducer(ClosureStateReducer::new(|_s: i32, u: i32| { + if u == 999 { + Err(TinyAgentsError::Graph("reducer boom".to_string())) + } else { + Ok(u) + } + })) + .add_node("boom", |_s, _c: NodeContext| async move { + Ok(NodeResult::Update(999)) + }) + .set_entry("boom") + .set_finish("boom") + .compile() + .unwrap() + .with_event_sink(sink.clone()); + + let err = graph.run(0).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Graph(_)), "got {err:?}"); + assert!( + sink.events() + .iter() + .any(|e| matches!(e, GraphEvent::RunFailed { .. })), + "a boundary reducer error must transition the run to Failed (RunFailed emitted)" + ); +} + +#[tokio::test] +async fn status_snapshot_reports_run() { + let graph = adding_graph(); + let run = graph + .run(Counter { + value: 0, + log: vec![], + }) + .await + .unwrap(); + let status = &run.status; + assert_eq!(status.status, ExecutionStatus::Completed); + assert_eq!(status.current_step, 2); + assert!(status.ended_at.is_some()); + assert!(status.error.is_none()); + assert_eq!(status.graph_id, *graph.graph_id()); +} + +/// A `Send` fan-out delivers a distinct per-branch argument to N parallel +/// activations of the *same* node, and the reducer merges their results. +#[tokio::test] +async fn send_fanout_delivers_distinct_args_to_parallel_branches() { + let graph = GraphBuilder::::new() + .set_reducer(ClosureStateReducer::new(|mut s: Counter, u: i32| { + s.value += u; + s.log.push(format!("worker:{u}")); + Ok(s) + })) + .with_parallel(true) + // dispatch fans out three custom inputs to the same worker node. + .add_node("dispatch", |_s: Counter, _c: NodeContext| async move { + Ok(NodeResult::Command(Command::send([ + Send::new("worker", json!(10)), + Send::new("worker", json!(20)), + Send::new("worker", json!(30)), + ]))) + }) + // each worker invocation consumes its own send arg as the update. + .add_node("worker", |_s: Counter, c: NodeContext| async move { + let arg = c + .send_arg + .expect("worker scheduled via Send carries an arg"); + let v = arg.as_i64().unwrap() as i32; + Ok(NodeResult::Update(v)) + }) + .mark_command_routing("dispatch") + .set_entry("dispatch") + .set_finish("worker") + .compile() + .unwrap(); + + let run = graph + .run(Counter { + value: 0, + log: vec![], + }) + .await + .unwrap(); + + // All three distinct args merged: 10 + 20 + 30. assert_eq!(run.state.value, 60); // The worker ran three times (one activation per Send packet). let worker_runs = run @@ -2514,7 +3346,7 @@ impl Checkpointer for FailNonTerminalCheckpointer { &self, checkpoint: crate::checkpoint::Checkpoint, ) -> tinyagents_harness::error::Result { - if !checkpoint.next_nodes.is_empty() { + if !checkpoint.tasks.is_empty() { return Err(tinyagents_harness::error::TinyAgentsError::Checkpoint( "injected background write failure".to_string(), )); @@ -2812,24 +3644,15 @@ async fn attributed_update_does_not_fire_an_unsatisfied_barrier() { .await .unwrap(); let written = cp.get("t-barrier-update", None).await.unwrap().unwrap(); + let written_next_nodes: Vec = written.tasks.iter().map(|t| t.node.clone()).collect(); assert!( - !written.next_nodes.iter().any(|n| n.as_str() == "merge"), + !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"), + 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!( @@ -2891,11 +3714,15 @@ fn forked_interrupt_graph( #[tokio::test] async fn attributed_update_keeps_other_pending_branches_scheduled() { - // Two independent branches are pending (`x` and the interrupted `c`). A - // manual write attributed to `x` schedules x's successor `y`, but it must - // not discard `c`: the attributed node's successors *add to* the schedule - // rather than replacing it, or the untouched branch is silently dropped and - // never runs again. + // `forked_interrupt_graph` runs `super -> [b, c]` in parallel: `b` + // completes (`Update(1)`) while `c` interrupts. Per the C2 fix, `b`'s + // routing is deferred rather than resolved immediately — only `c` (the + // interrupted branch) is in `next_nodes`/pending, and `b` sits in + // `completed_tasks` awaiting a step-finishing routing pass. A manual + // write attributed to `b` (`update_state`'s carried-completion routing — + // see `state_api::update_state`) resolves that deferred routing (`b`'s + // successor `x`), and must not discard `c`: the untouched interrupted + // branch stays pending alongside it. let cp = Arc::new(InMemoryCheckpointer::::new()); let graph = forked_interrupt_graph(cp.clone(), Arc::new(AtomicBool::new(false))); @@ -2912,42 +3739,43 @@ async fn attributed_update_keeps_other_pending_branches_scheduled() { assert!(paused.is_interrupted()); let before = cp.get("t-fork-update", None).await.unwrap().unwrap(); - assert!( - before.next_nodes.iter().any(|n| n.as_str() == "x") - && before.next_nodes.iter().any(|n| n.as_str() == "c"), - "precondition: both branches pending, got {:?}", - before.next_nodes + assert_eq!( + before + .tasks + .iter() + .map(|t| t.node.to_string()) + .collect::>(), + vec!["c".to_string()], + "precondition: only the interrupted branch is pending, b's routing is deferred" + ); + assert_eq!( + before + .completed + .iter() + .map(|c| c.node.to_string()) + .collect::>(), + vec!["b".to_string()], + "precondition: b completed this step but its routing was not yet resolved" ); graph - .update_state("t-fork-update", 10, Some(NodeId::from("x"))) + .update_state("t-fork-update", 10, Some(NodeId::from("b"))) .await .unwrap(); let written = cp.get("t-fork-update", None).await.unwrap().unwrap(); + let written_next_nodes: Vec = written.tasks.iter().map(|t| t.node.clone()).collect(); assert!( - written.next_nodes.iter().any(|n| n.as_str() == "y"), - "the attributed node's successor must be scheduled, got {:?}", - written.next_nodes + written_next_nodes.iter().any(|n| n.as_str() == "x"), + "b's deferred successor x must now be scheduled, got {written_next_nodes:?}" ); assert!( - written.next_nodes.iter().any(|n| n.as_str() == "c"), - "the untouched pending branch must stay scheduled, got {:?}", - written.next_nodes + written_next_nodes.iter().any(|n| n.as_str() == "c"), + "the untouched pending branch must stay scheduled, got {written_next_nodes:?}" ); assert!( - !written.next_nodes.iter().any(|n| n.as_str() == "x"), - "the attributed node itself is completed, not pending: {:?}", - written.next_nodes + !written_next_nodes.iter().any(|n| n.as_str() == "b"), + "the attributed node itself is completed, not pending: {written_next_nodes:?}" ); - // Resume prefers `pending_activations` over `next_nodes`, so the two must - // never disagree. - 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-fork-update").await.unwrap(); assert!( @@ -2955,8 +3783,14 @@ async fn attributed_update_keeps_other_pending_branches_scheduled() { "the dropped branch must still run, visited {:?}", done.visited ); - // 1 (b) + 10 (manual write) + 2 (c) + 40 (y). - assert_eq!(done.state.value, 53); + assert!( + done.visited.iter().any(|n| n.as_str() == "x"), + "b's deferred successor must run, visited {:?}", + done.visited + ); + // 1 (b, applied at the original boundary) + 10 (manual write) + 2 (c) + + // 20 (x) + 40 (y). + assert_eq!(done.state.value, 73); } #[tokio::test] @@ -2986,9 +3820,9 @@ async fn attributed_update_to_sink_node_keeps_other_pending_branches() { let written = cp.get("t-fork-sink", None).await.unwrap().unwrap(); assert_eq!( written - .next_nodes + .tasks .iter() - .map(|n| n.to_string()) + .map(|t| t.node.to_string()) .collect::>(), vec!["x".to_string()], "the sibling branch must survive an attributed write to a sink node" @@ -3001,10 +3835,15 @@ async fn attributed_update_to_sink_node_keeps_other_pending_branches() { #[tokio::test] async fn attributed_update_preserves_pending_send_args_of_other_branches() { - // Three `Send` activations of `worker` are pending behind an interrupt. A - // write attributed to an unrelated node must carry them over *with* their - // args — dropping them loses the fanout, and re-scheduling them by node id - // alone loses each packet's payload. + // Three `Send` activations of `worker` are scheduled; the arg-1 worker + // interrupts while arg-2 and arg-3 complete in the same (parallel) step. + // Per the C1 fix, the completed higher-index workers are folded into + // state (not discarded/re-run) and are *not* part of the pending set — + // only the genuinely-interrupted arg-1 worker is. A write attributed to + // an unrelated node (`side`) must carry that one pending Send packet over + // *with* its arg (dropping it loses the fanout, and re-scheduling it by + // node id alone loses its payload) without resurrecting the two + // already-completed workers. let cp = Arc::new(InMemoryCheckpointer::::new()); let graph = GraphBuilder::::new() .with_parallel(true) @@ -3058,17 +3897,41 @@ async fn attributed_update_preserves_pending_send_args_of_other_branches() { .await .unwrap(); assert!(paused.is_interrupted()); + // C1: the two completed higher-index workers (args 2 and 3) are folded + // into state despite the lower-index (arg 1) worker interrupting. + assert_eq!( + paused.state.value, 5, + "arg-2 and arg-3 workers must complete despite arg-1 interrupting" + ); - graph - .update_state("t-send-update", 0, Some(NodeId::from("side"))) - .await + let before = cp.get("t-send-update", None).await.unwrap().unwrap(); + let before_pending = before.tasks.clone(); + assert_eq!( + before_pending + .iter() + .filter(|a| a.node.as_str() == "worker") + .filter_map(|a| a.send_arg.as_ref().and_then(|v| v.as_i64())) + .collect::>(), + vec![1], + "only the genuinely-interrupted arg-1 worker is pending, got {before_pending:?}" + ); + assert_eq!( + before + .completed + .iter() + .filter(|c| c.node.as_str() == "worker") + .count(), + 2, + "the two completed workers are recorded as completed, not pending" + ); + + graph + .update_state("t-send-update", 0, Some(NodeId::from("side"))) + .await .unwrap(); let written = cp.get("t-send-update", None).await.unwrap().unwrap(); - let pending = written - .pending_activations - .clone() - .expect("an attributed write must persist the merged activations"); - let mut args: Vec = pending + let pending = written.tasks.clone(); + let args: Vec = pending .iter() .filter(|a| a.node.as_str() == "worker") .map(|a| { @@ -3079,16 +3942,14 @@ async fn attributed_update_preserves_pending_send_args_of_other_branches() { .unwrap() }) .collect(); - args.sort_unstable(); - assert_eq!(args, vec![1, 2, 3], "every pending Send packet survives"); + assert_eq!( + args, + vec![1], + "the still-pending Send packet survives with its arg, the completed ones are not resurrected" + ); assert!( pending.iter().any(|a| a.node.as_str() == "tail"), - "the attributed node's successor is scheduled alongside them" - ); - assert_eq!( - pending.iter().map(|a| a.node.clone()).collect::>(), - written.next_nodes, - "pending activations and next nodes must describe the same schedule" + "the attributed node's successor is scheduled alongside it" ); } @@ -3396,3 +4257,628 @@ async fn async_durability_skips_a_write_whose_predecessor_failed() { "no orphaned checkpoint may be appended after a broken lineage" ); } + +/// C3/R4 regression: two concurrent `run_with_thread` calls for the SAME +/// thread id must not interleave their node execution. Before the executor +/// held its own per-thread lock, nothing serialized two concurrent +/// entry-point calls at this layer (only `delegation::run` worked around it +/// with a private lock of its own) — see the C3 finding in +/// `docs/runtime-comparison/code-review-graph.md`. +#[tokio::test] +async fn concurrent_run_with_thread_calls_on_one_thread_serialize() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let concurrent = Arc::new(AtomicUsize::new(0)); + let max_concurrent = Arc::new(AtomicUsize::new(0)); + let concurrent_for_node = concurrent.clone(); + let max_concurrent_for_node = max_concurrent.clone(); + 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("inc", move |_s: Counter, _c: NodeContext| { + let concurrent = concurrent_for_node.clone(); + let max_concurrent = max_concurrent_for_node.clone(); + async move { + let now = concurrent.fetch_add(1, AtomicOrdering::SeqCst) + 1; + max_concurrent.fetch_max(now, AtomicOrdering::SeqCst); + tokio::time::sleep(Duration::from_millis(20)).await; + concurrent.fetch_sub(1, AtomicOrdering::SeqCst); + Ok(NodeResult::Update(1)) + } + }) + .set_entry("inc") + .set_finish("inc") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let (r1, r2) = tokio::join!( + graph.run_with_thread( + "t-concurrent-serialize", + Counter { + value: 0, + log: vec![], + }, + ), + graph.run_with_thread( + "t-concurrent-serialize", + Counter { + value: 0, + log: vec![], + }, + ), + ); + r1.expect("first run completes"); + r2.expect("second run completes"); + + assert_eq!( + max_concurrent.load(AtomicOrdering::SeqCst), + 1, + "the executor's per-thread lock must serialize concurrent run_with_thread calls" + ); + + // Both runs wrote a complete, un-torn checkpoint for the thread — no + // interleaved/partial record from one run's boundary landing inside the + // other's. + let listed = cp.list("t-concurrent-serialize").await.unwrap(); + assert_eq!( + listed.len(), + 2, + "each serialized run wrote its own checkpoint" + ); + let run_ids: std::collections::HashSet<_> = listed.iter().map(|m| m.run_id.clone()).collect(); + assert_eq!(run_ids.len(), 2, "the two runs must not share a run id"); +} + +// ---- R5: typed task identity --------------------------------------------- + +#[tokio::test] +async fn node_context_task_id_is_stable_across_retry_attempts() { + // `run_node_with_retry` re-clones the *context* for each attempt rather + // than rebuilding it, so `NodeContext::task_id()` — built once per + // activation before the retry loop starts — must read the same value on + // every attempt of one activation. + let seen_ids = Arc::new(std::sync::Mutex::new(Vec::::new())); + let attempts = Arc::new(AtomicUsize::new(0)); + let seen_for_node = seen_ids.clone(); + let attempts_for_node = attempts.clone(); + let graph = GraphBuilder::::overwrite() + .add_node("flaky", move |s, c: NodeContext| { + let seen_ids = seen_for_node.clone(); + let attempts = attempts_for_node.clone(); + async move { + seen_ids + .lock() + .unwrap() + .push(c.task_id().as_str().to_string()); + let n = attempts.fetch_add(1, AtomicOrdering::SeqCst); + if n < 2 { + Err(TinyAgentsError::Model(format!("transient blip {n}"))) + } else { + Ok(NodeResult::Update(s + 1)) + } + } + }) + .set_entry("flaky") + .set_finish("flaky") + .compile() + .unwrap() + .with_node_retry(RetryPolicy::default().with_max_attempts(4)); + + let run = graph.run(0).await.unwrap(); + assert_eq!(run.state, 1); + + let recorded = seen_ids.lock().unwrap(); + assert_eq!(recorded.len(), 3, "one attempt-observation per try"); + assert!( + !recorded[0].is_empty(), + "a real task id was assigned before the retry loop started" + ); + assert!( + recorded.iter().all(|id| id == &recorded[0]), + "every retry attempt of the same activation sees the same task id: {recorded:?}" + ); +} + +#[tokio::test] +async fn legacy_checkpoint_json_without_task_id_fields_still_resumes() { + // `task_id` was added to `PendingActivation`/`Interrupt` as typed fields + // (R5). A checkpoint written before either field existed carries neither + // key at all (not even as an empty string) — `#[serde(default)]` must + // still decode it, and resume must still work, falling back to + // node-id-keyed resume exactly as it did before R5. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let graph = GraphBuilder::::overwrite() + .add_node("gate", |s: i32, c: NodeContext| async move { + match c.resume { + Some(_) => Ok(NodeResult::Update(s + 1)), + None => Ok(NodeResult::Interrupt(Interrupt::new("gate", json!({})))), + } + }) + .set_entry("gate") + .set_finish("gate") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let paused = graph.run_with_thread("t-legacy-task-id", 0).await.unwrap(); + assert!(paused.is_interrupted()); + + // Round-trip the checkpoint through JSON, stripping every `task_id` key + // to simulate a pre-R5 record (checkpoint format v2's `tasks` field is + // where a task id lives today — see `Checkpoint::tasks`). + let mut raw = + serde_json::to_value(cp.get("t-legacy-task-id", None).await.unwrap().unwrap()).unwrap(); + if let Some(activations) = raw.get_mut("tasks").and_then(|v| v.as_array_mut()) { + for activation in activations { + activation.as_object_mut().unwrap().remove("task_id"); + } + } + if let Some(interrupts) = raw.get_mut("interrupts").and_then(|v| v.as_array_mut()) { + for interrupt in interrupts { + interrupt.as_object_mut().unwrap().remove("task_id"); + } + } + let legacy: Checkpoint = serde_json::from_value(raw) + .expect("a pre-R5 checkpoint with no task_id keys at all must still decode"); + assert!(legacy.tasks[0].task_id.as_str().is_empty()); + assert!(legacy.interrupts[0].task_id.is_none()); + cp.put(legacy).await.unwrap(); + + let done = graph + .resume("t-legacy-task-id", Command::resume(json!("go"))) + .await + .unwrap(); + assert!(!done.is_interrupted()); + assert_eq!(done.state, 1); +} + +#[tokio::test] +async fn resume_from_a_checkpoint_format_v1_json_record() { + // A genuine checkpoint format v1 record — as a build before `version`/ + // `tasks`/`completed` existed would have written: no `version` key at + // all, pending work in `next_nodes`, nothing in `pending_activations`. + // `InMemoryCheckpointer::get`/`put` normalize on every decode path (see + // `Checkpoint::normalize`), so a hand-built v1 record put straight into + // the store must resume exactly like a v2 one. + let cp = Arc::new(InMemoryCheckpointer::::new()); + let graph = GraphBuilder::::overwrite() + .add_node("gate", |s: i32, c: NodeContext| async move { + match c.resume { + Some(_) => Ok(NodeResult::Update(s + 1)), + None => Ok(NodeResult::Interrupt(Interrupt::new("gate", json!({})))), + } + }) + .set_entry("gate") + .set_finish("gate") + .compile() + .unwrap() + .with_checkpointer(cp.clone()); + + let v1_json = json!({ + "thread_id": "t-v1-record", + "checkpoint_id": "c1", + "run_id": null, + "parent_checkpoint_id": null, + "namespace": [], + "state": 0, + "next_nodes": ["gate"], + "completed_tasks": [], + "completed_routes": [], + "pending_writes": [], + "interrupts": [], + "pending_activations": null, + "barrier_arrivals": [], + "metadata": { "source": "loop", "step": 1, "interrupted_nodes": ["gate"] }, + }); + let v1: Checkpoint = serde_json::from_value(v1_json).unwrap(); + assert_eq!(v1.version, 1, "precondition: this is a genuine v1 record"); + cp.put(v1).await.unwrap(); + + // The store normalized it on `get` before this handed it back — confirm + // that directly before exercising resume through it. + let normalized = cp.get("t-v1-record", None).await.unwrap().unwrap(); + assert_eq!( + normalized.version, + crate::checkpoint::CHECKPOINT_FORMAT_VERSION + ); + assert_eq!(normalized.tasks.len(), 1); + assert_eq!(normalized.tasks[0].node, NodeId::from("gate")); + + let done = graph + .resume("t-v1-record", Command::resume(json!("go"))) + .await + .unwrap(); + assert!(!done.is_interrupted()); + assert_eq!(done.state, 1, "resumed and ran gate to completion"); +} + +#[tokio::test] +async fn update_state_and_fork_state_write_checkpoint_format_v2() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let graph = chain_graph(cp.clone()); + graph.run_with_thread("t-v2-writes", 0).await.unwrap(); + + let config = graph + .update_state("t-v2-writes", 5, None) + .await + .expect("update_state"); + let written = cp + .get(&config.thread_id, config.checkpoint_id.as_deref()) + .await + .unwrap() + .unwrap(); + assert_eq!( + written.version, + crate::checkpoint::CHECKPOINT_FORMAT_VERSION, + "update_state writes checkpoint format v2" + ); + assert!(written.next_nodes.is_empty(), "v1 fields left unpopulated"); + + let fork_config = graph + .fork_state("t-v2-writes", None, "t-v2-forked") + .await + .expect("fork_state"); + let forked = cp + .get(&fork_config.thread_id, fork_config.checkpoint_id.as_deref()) + .await + .unwrap() + .unwrap(); + assert_eq!( + forked.version, + crate::checkpoint::CHECKPOINT_FORMAT_VERSION, + "fork_state writes checkpoint format v2" + ); + assert!(forked.next_nodes.is_empty(), "v1 fields left unpopulated"); +} + +// ── I4: panic safety, cooperative cancellation, and the run-drop guard ────── + +/// A single-node graph whose handler panics the first `panic_times` +/// invocations, then succeeds with `+1`. Mirrors [`flaky_graph`] but for a +/// panic instead of a transient `Err`, so the panic-safety tests below can +/// reuse the same failure/retry assertions. +fn panicking_graph(panic_times: usize, attempts: Arc) -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node("flaky", move |s, _c: NodeContext| { + let attempts = attempts.clone(); + async move { + let n = attempts.fetch_add(1, AtomicOrdering::SeqCst); + if n < panic_times { + panic!("synthetic node panic {n}"); + } + Ok(NodeResult::Update(s + 1)) + } + }) + .set_entry("flaky") + .set_finish("flaky") + .compile() + .unwrap() +} + +/// A node handler panic must not poison the whole run future: it becomes an +/// ordinary node failure that flows through the normal failure boundary +/// (checkpoint write, `Failed` status), and the checkpoint it leaves is +/// loadable and resumable via `retry` — exactly like a returned `Err`. +#[tokio::test] +async fn node_panic_is_a_resumable_failure_not_a_lost_run() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let attempts = Arc::new(AtomicUsize::new(0)); + let graph = panicking_graph(1, attempts.clone()).with_checkpointer(cp.clone()); + + let err = graph.run_with_thread("panicky", 5).await.unwrap_err(); + assert!(matches!(err, TinyAgentsError::Graph(_)), "got {err:?}"); + assert!( + err.to_string().contains("panicked"), + "error should describe the panic: {err}" + ); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 1); + + // The failure boundary persisted a loadable, resumable checkpoint. + let status = graph.get_state("panicky", None).await.unwrap().unwrap(); + assert_eq!(status.next_nodes, vec![NodeId::from("flaky")]); + + // The panic does not recur: `retry` re-runs the node to completion. + let resumed = graph.retry("panicky").await.unwrap(); + assert_eq!(resumed.state, 6); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 2); +} + +/// A graph whose entry node cancels `token` as soon as it starts (simulating +/// a caller requesting cancellation while the node is already in flight), +/// then keeps running for a while longer before completing — so a test can +/// assert the cancellation is observed *without* waiting for the slow node. +fn cancel_mid_step_graph(token: tinyagents_harness::CancellationToken) -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node("a", move |s, _c: NodeContext| { + let token = token.clone(); + async move { + token.cancel(); + tokio::time::sleep(Duration::from_millis(200)).await; + Ok(NodeResult::Update(s + 1)) + } + }) + .add_node("b", |s, _c: NodeContext| async move { + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("a") + .add_edge("a", "b") + .set_finish("b") + .compile() + .unwrap() +} + +/// Cancelling a [`RunOptions`] token while a superstep's node handlers are +/// still in flight stops the run without waiting for that node to finish: +/// the still-pending activations are persisted as a resumable checkpoint, +/// the run reports `Cancelled`, and a later `resume` completes it. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancellation_mid_step_is_resumable() { + let cp = Arc::new(InMemoryCheckpointer::::new()); + let token = tinyagents_harness::CancellationToken::new(); + let graph = cancel_mid_step_graph(token.clone()).with_checkpointer(cp.clone()); + + let started = std::time::Instant::now(); + let run = graph + .run_with_thread_options("cancel-me", 0, RunOptions::with_cancellation(token)) + .await + .unwrap(); + assert_eq!(run.status.status, ExecutionStatus::Cancelled); + assert!( + started.elapsed() < Duration::from_millis(150), + "cancellation should not wait out node `a`'s 200ms sleep" + ); + + // The pending activation (node `a`, never having completed) is exactly + // what a resume re-runs. + let status = graph.get_state("cancel-me", None).await.unwrap().unwrap(); + assert_eq!(status.next_nodes, vec![NodeId::from("a")]); + + // Resuming with a fresh (never-cancelled) token completes the run. + let fresh = cancel_mid_step_graph(tinyagents_harness::CancellationToken::new()) + .with_checkpointer(cp.clone()); + let resumed = fresh.resume("cancel-me", Command::new()).await.unwrap(); + assert_eq!(resumed.state, 2); + assert_eq!(resumed.status.status, ExecutionStatus::Completed); +} + +/// A graph with one node that sleeps far longer than the caller is willing +/// to wait, so wrapping the run in a short `tokio::time::timeout` drops the +/// run future mid-flight without the executor's own cancellation/failure +/// paths ever running. +fn slow_node_graph() -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node("slow", |s: i32, _c: NodeContext| async move { + tokio::time::sleep(Duration::from_secs(5)).await; + Ok(NodeResult::Update(s + 1)) + }) + .set_entry("slow") + .set_finish("slow") + .compile() + .unwrap() +} + +/// Dropping the run future before it reaches a terminal state (here, via an +/// external `tokio::time::timeout` that outraces the run) must not leave the +/// run's stored status stuck at `Running` forever: the [`RunDropGuard`] +/// (I4 part 3) spawns a best-effort background write that marks it +/// `Cancelled`. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn dropping_the_run_future_marks_status_cancelled_not_running() { + use crate::observability::GraphStatusStore; + let store = Arc::new(crate::observability::InMemoryGraphStatusStore::default()); + let sink = Arc::new(CollectingSink::new()); + let graph = slow_node_graph() + .with_status_store(store.clone()) + .with_event_sink(sink.clone()); + + let outcome = tokio::time::timeout(Duration::from_millis(20), graph.run(0)).await; + assert!( + outcome.is_err(), + "the 5s sleep must outlast the 20ms timeout, dropping the run future" + ); + + // `RunStarted` is emitted synchronously before the node runs, so the run + // id is known even though the run itself never returned. + let run_id = sink + .events() + .into_iter() + .find_map(|e| match e { + GraphEvent::RunStarted { run_id } => Some(run_id), + _ => None, + }) + .expect("RunStarted was emitted before the timeout fired"); + + // The drop guard's write happens on a detached background task; give it + // a moment to land before asserting on the store. + tokio::time::sleep(Duration::from_millis(200)).await; + + let status = store + .get_status(run_id.as_str()) + .await + .unwrap() + .expect("the drop guard persisted a status for this run"); + assert_ne!( + status.status, + ExecutionStatus::Running, + "the drop guard must not leave the run stuck at Running" + ); + assert_eq!(status.status, ExecutionStatus::Cancelled); +} + +// ── M2: hot-path state cloning ──────────────────────────────────────────── +// +// `docs/runtime-comparison/code-review-graph.md` M2: before this fix, every +// node-handler invocation — each attempt of a retried task, and each branch +// of a parallel `Send` fan-out — cloned the whole `State` value at the call +// site (`compiled/step.rs`'s `handler(state.clone(), ctx.clone())`). A +// 4-way fan-out with 3 retries on one branch is 4 + 3 = 7 invocations, so at +// least 7 `State::clone()` calls for that one superstep alone. The fix +// clones `State` at most once per superstep (into an `Arc`); every +// attempt/branch after that shares the `Arc` via a cheap `Arc::clone` +// instead. These tests instrument `State::clone()` itself to prove the +// bound, using `add_node_shared` (M2's zero-clone handler entry point) so +// the count reflects the executor's own cloning rather than the +// `add_node`/`Arc` compatibility adapter's per-invocation clone. + +/// A state value that counts every `Clone::clone()` call made on it, via a +/// shared atomic counter, so a test can assert exactly how many times the +/// executor cloned the whole state during a run. +#[derive(Debug)] +struct CountingState { + clones: Arc, + value: i64, +} + +impl Clone for CountingState { + fn clone(&self) -> Self { + self.clones.fetch_add(1, AtomicOrdering::SeqCst); + Self { + clones: self.clones.clone(), + value: self.value, + } + } +} + +/// A 4-way parallel `Send` fan-out, with 3 retries on the `arg == 1` branch +/// before it succeeds, clones `State` at most once per superstep (M2): one +/// clone for the `dispatch` step, one for the fan-out step — regardless of +/// the fan-out width or the retried branch's attempt count. Before the fix +/// this was at least 7 clones (4 branches + 3 extra attempts) for the +/// fan-out step alone. +#[tokio::test] +async fn parallel_fanout_with_retries_clones_state_at_most_once_per_step() { + let clones = Arc::new(AtomicUsize::new(0)); + let attempts = Arc::new(AtomicUsize::new(0)); + + let graph = GraphBuilder::::new() + .with_parallel(true) + .set_reducer(ClosureStateReducer::new(|mut s: CountingState, u: i64| { + s.value += u; + Ok(s) + })) + .add_node_shared( + "dispatch", + |_s: Arc, _c: NodeContext| async move { + Ok(NodeResult::Command(Command::send([ + Send::new("worker", json!(1)), + Send::new("worker", json!(2)), + Send::new("worker", json!(3)), + Send::new("worker", json!(4)), + ]))) + }, + ) + .add_node_shared("worker", { + let attempts = attempts.clone(); + move |_s: Arc, c: NodeContext| { + let attempts = attempts.clone(); + async move { + let arg = c + .send_arg + .clone() + .expect("worker scheduled via Send must carry its arg") + .as_i64() + .unwrap(); + if arg == 1 { + let n = attempts.fetch_add(1, AtomicOrdering::SeqCst); + if n < 3 { + return Err(TinyAgentsError::Model(format!("transient blip {n}"))); + } + } + Ok(NodeResult::Update(arg)) + } + } + }) + .with_node_policy( + "worker", + crate::builder::NodePolicy { + retry: Some( + RetryPolicy::default() + .with_max_attempts(5) + .with_backoff_sleep(false), + ), + ..crate::builder::NodePolicy::default() + }, + ) + .set_entry("dispatch") + .mark_command_routing("dispatch") + .set_finish("worker") + .compile() + .unwrap(); + + let state = CountingState { + clones: clones.clone(), + value: 0, + }; + let run = graph.run(state).await.unwrap(); + assert_eq!(run.state.value, 1 + 2 + 3 + 4, "every branch's arg applied"); + assert_eq!( + attempts.load(AtomicOrdering::SeqCst), + 4, + "3 failed attempts + 1 success for the arg==1 branch" + ); + + let total_clones = clones.load(AtomicOrdering::SeqCst); + assert!( + total_clones <= 2, + "expected at most one `State` clone per superstep (2 steps: dispatch, \ + then the 4-way fan-out with 3 retries), got {total_clones}" + ); +} + +/// The sequential (non-parallel) counterpart: a single node retried 3 times +/// before it succeeds clones `State` at most once for its one superstep — +/// the per-attempt clone the M2 finding describes is gone regardless of +/// concurrency mode. +#[tokio::test] +async fn sequential_retries_clone_state_at_most_once_per_step() { + let clones = Arc::new(AtomicUsize::new(0)); + let attempts = Arc::new(AtomicUsize::new(0)); + + let graph = GraphBuilder::::new() + .set_reducer(ClosureStateReducer::new(|mut s: CountingState, u: i64| { + s.value += u; + Ok(s) + })) + .add_node_shared("flaky", { + let attempts = attempts.clone(); + move |_s: Arc, _c: NodeContext| { + let attempts = attempts.clone(); + async move { + let n = attempts.fetch_add(1, AtomicOrdering::SeqCst); + if n < 3 { + Err(TinyAgentsError::Model(format!("transient blip {n}"))) + } else { + Ok(NodeResult::Update(1)) + } + } + } + }) + .set_entry("flaky") + .set_finish("flaky") + .compile() + .unwrap() + .with_node_retry( + RetryPolicy::default() + .with_max_attempts(5) + .with_backoff_sleep(false), + ); + + let state = CountingState { + clones: clones.clone(), + value: 0, + }; + let run = graph.run(state).await.unwrap(); + assert_eq!(run.state.value, 1); + assert_eq!(attempts.load(AtomicOrdering::SeqCst), 4); + + let total_clones = clones.load(AtomicOrdering::SeqCst); + assert!( + total_clones <= 2, + "expected at most one `State` clone for this one superstep's 4 \ + attempts, got {total_clones}" + ); +} diff --git a/crates/tinyagents-graph/src/compiled/types.rs b/crates/tinyagents-graph/src/compiled/types.rs index ed25fa36..cf4bd88c 100644 --- a/crates/tinyagents-graph/src/compiled/types.rs +++ b/crates/tinyagents-graph/src/compiled/types.rs @@ -10,9 +10,10 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use std::sync::atomic::AtomicU64; use crate::builder::START; -use crate::builder::{BarrierRelief, Branch, BuilderNode, NodeMeta}; +use crate::builder::{BarrierRelief, Branch, BuilderNode, NodeMeta, NodePolicy, UpdateCodec}; use crate::checkpoint::{CheckpointConfig, CheckpointMetadata, Checkpointer, DurabilityMode}; use crate::command::Interrupt; use crate::observability::{GraphEventJournal, GraphStatusStore}; @@ -31,9 +32,8 @@ pub struct CompiledGraph { /// Optional human-readable graph name surfaced by the topology export. pub(crate) name: Option, pub(crate) nodes: Arc>>, - pub(crate) edges: Arc>, + pub(crate) edges: Arc>>, pub(crate) branches: Arc>>, - #[allow(dead_code)] pub(crate) command_nodes: Arc>, /// Barrier/waiting edges: target -> the predecessor set that must all /// complete (across steps) before the target activates. @@ -89,6 +89,39 @@ pub struct CompiledGraph { /// abort-on-first-error behavior. Configured via /// [`CompiledGraph::with_node_retry`](crate::CompiledGraph::with_node_retry). pub(crate) node_retry: Option, + /// Monotonic sequence counter for [`crate::stream::GraphEventEnvelope::seq`], + /// shared (via this `Arc`) across clones that only change `event_sink` + /// (journal wrapping) so a run's sequence stays continuous end to end. + /// A subgraph embedded as a node gets its own fresh counter (see + /// [`crate::subgraph`]) — its distinct `namespace` already disambiguates + /// its stream, and note in [`crate::stream::GraphEventEnvelope`] why a + /// shared counter is not needed across that boundary. + pub(crate) sequence: Arc, + /// Per-node execution policies; see [`NodePolicy`]. Resolved per + /// activation against `node_defaults` and the legacy graph-wide + /// `node_retry`/`node_timeout` fields by [`NodePolicy::resolve`]. + pub(crate) node_policies: Arc>>, + /// Graph-wide default execution policy (`GraphBuilder::set_node_defaults`). + pub(crate) node_defaults: Option>>, + /// Backend for opt-in per-node result caching; see + /// [`CompiledGraph::with_task_cache`]. `None` (default) disables caching + /// entirely, even for nodes with a [`NodeCachePolicy`] installed via + /// [`CompiledGraph::with_cached_node`]. + pub(crate) task_cache: Option>, + /// Per-node cache policy plus the type-erased `Update` codec installed by + /// [`CompiledGraph::with_cached_node`] (the entry point that supplies the + /// `Serialize + DeserializeOwned` bound this struct itself is free of). + pub(crate) cached_nodes: Arc>>, + /// Nodes the executor pauses before running + /// ([`crate::GraphBuilder::interrupt_before`]). + pub(crate) interrupt_before: Arc>, + /// Nodes the executor pauses after running, holding their result back + /// from committed state until resume + /// ([`crate::GraphBuilder::interrupt_after`]). + pub(crate) interrupt_after: Arc>, + /// The `Update` codec that persists/replays an `interrupt_after` node's + /// deferred result; `None` when no node uses `interrupt_after`. + pub(crate) update_codec: Option>, } impl std::fmt::Debug for CompiledGraph { @@ -133,6 +166,14 @@ impl Clone for CompiledGraph { run_deadline: self.run_deadline, durability: self.durability, node_retry: self.node_retry.clone(), + sequence: self.sequence.clone(), + node_policies: self.node_policies.clone(), + node_defaults: self.node_defaults.clone(), + task_cache: self.task_cache.clone(), + cached_nodes: self.cached_nodes.clone(), + interrupt_before: self.interrupt_before.clone(), + interrupt_after: self.interrupt_after.clone(), + update_codec: self.update_codec.clone(), } } } @@ -286,6 +327,14 @@ pub struct GraphExecution { pub status: GraphRunStatus, /// The latest persisted checkpoint id, if checkpointing was enabled. pub checkpoint_id: Option, + /// `true` when the run stopped at a superstep boundary because its + /// [`RunOptions::drain`] signal was raised (see [`DrainSignal`]): the + /// step in flight finished, its boundary was committed, and the next + /// step's activations were checkpointed instead of run. `status` is + /// then [`tinyagents_harness::ids::ExecutionStatus::Drained`], and the + /// thread continues with [`CompiledGraph::resume`]/[`CompiledGraph::retry`]. + /// `false` on every other outcome. + pub drained: bool, } /// One external input used to seed a graph run. @@ -379,6 +428,151 @@ pub struct StateSnapshot { pub pending_interrupts: Vec, } +/// The shared flag behind a [`DrainHandle`]/[`DrainSignal`] pair. +/// +/// A plain latching `AtomicBool`: the executor only ever *polls* it at +/// superstep boundaries (drain never interrupts a step in flight, so there +/// is nothing to wake), which keeps the pair as cheap as a +/// [`tinyagents_harness::CancellationToken`] clone. +#[derive(Debug, Default)] +struct DrainState { + requested: std::sync::atomic::AtomicBool, +} + +/// The requesting side of a graceful-drain pair (see [`DrainSignal::new`]). +/// +/// Held by whoever decides the run should stop — a shutdown hook, a +/// supervisor, a node handler that detects it should yield — and signalled +/// with [`DrainHandle::drain`]. Cheap to clone; every clone signals the same +/// [`DrainSignal`]. Draining is latching: once requested it cannot be undone. +#[derive(Clone, Debug)] +pub struct DrainHandle { + state: Arc, +} + +impl DrainHandle { + /// Requests a graceful drain: the run finishes the superstep currently in + /// flight (every handler is awaited to completion, the reducer and + /// boundary checkpoint run as normal) and then stops instead of starting + /// the next one. Idempotent. + pub fn drain(&self) { + self.state + .requested + .store(true, std::sync::atomic::Ordering::SeqCst); + } + + /// Whether a drain has been requested through any clone of this handle. + pub fn is_requested(&self) -> bool { + self.state + .requested + .load(std::sync::atomic::Ordering::SeqCst) + } +} + +/// The watched side of a graceful-drain pair, handed to the executor via +/// [`RunOptions::drain`] (see [`DrainSignal::new`]). +/// +/// The deliberate contrast with [`RunOptions::cancellation`]: a cancellation +/// is raced against the in-flight step and abandons it, so nothing from that +/// step is trusted to have run; a drain is *only* checked between supersteps, +/// so the in-flight step always completes and commits before the run stops. +/// The stop is then reported as [`GraphExecution::drained`] (`true`), status +/// [`tinyagents_harness::ids::ExecutionStatus::Drained`], and +/// [`crate::GraphEvent::RunDrained`], with a resumable checkpoint naming the +/// next step's activations persisted on a checkpointed thread. +#[derive(Clone, Debug)] +pub struct DrainSignal { + state: Arc, +} + +impl DrainSignal { + /// Creates a connected `(DrainHandle, DrainSignal)` pair. + pub fn new() -> (DrainHandle, DrainSignal) { + let state = Arc::new(DrainState::default()); + ( + DrainHandle { + state: state.clone(), + }, + DrainSignal { state }, + ) + } + + /// Whether the paired [`DrainHandle`] has requested a drain. + pub fn is_requested(&self) -> bool { + self.state + .requested + .load(std::sync::atomic::Ordering::SeqCst) + } +} + +/// Per-run options threaded through [`CompiledGraph::run_with_options`], +/// [`CompiledGraph::run_with_thread_options`], and +/// [`CompiledGraph::resume_with_options`] (I4 part 2). +/// +/// A small options bag rather than a combinatorial +/// `run_with_cancel`/`run_with_cancel_and_thread`/... family of entry points. +/// Two stop signals are carried, differing in how they treat the step in +/// flight: +/// +/// - `cancellation` — a [`tinyagents_harness::CancellationToken`] requesting +/// cooperative cancellation. The executor checks it at every superstep +/// boundary (before starting a new step) *and* races it against that step's +/// in-flight node-handler futures, so a long-running node cannot +/// indefinitely block a cancellation request. On cancellation the run's +/// status becomes [`tinyagents_harness::ids::ExecutionStatus::Cancelled`] +/// and, on a checkpointed thread, a resumable checkpoint is persisted +/// naming the still-pending activations. +/// - `drain` — a [`DrainSignal`] requesting a *graceful* stop. It is checked +/// only between supersteps: the step in flight always finishes and commits +/// its boundary, then the next step's activations are checkpointed instead +/// of run. The run reports [`GraphExecution::drained`] with status +/// [`tinyagents_harness::ids::ExecutionStatus::Drained`]. +/// +/// Either way the run can be continued later with +/// [`CompiledGraph::resume`]/[`CompiledGraph::retry`]. +#[derive(Clone, Debug, Default)] +pub struct RunOptions { + /// Optional cooperative-cancellation token for this run. + pub cancellation: Option, + /// Optional graceful-drain signal for this run. + pub drain: Option, +} + +impl RunOptions { + /// Builds empty run options (no cancellation token, no drain signal). + pub fn new() -> Self { + Self::default() + } + + /// Builds run options carrying `token` for cooperative cancellation. + pub fn with_cancellation(token: tinyagents_harness::CancellationToken) -> Self { + Self { + cancellation: Some(token), + drain: None, + } + } + + /// Builds run options carrying `signal` for a graceful drain. + pub fn with_drain(signal: DrainSignal) -> Self { + Self { + cancellation: None, + drain: Some(signal), + } + } + + /// Sets the graceful-drain signal, keeping any cancellation token. + pub fn drain(mut self, signal: DrainSignal) -> Self { + self.drain = Some(signal); + self + } + + /// Sets the cancellation token, keeping any drain signal. + pub fn cancellation(mut self, token: tinyagents_harness::CancellationToken) -> Self { + self.cancellation = Some(token); + self + } +} + /// Selects which checkpoint a time-travel resume starts from. /// /// [`CompiledGraph::resume`](crate::CompiledGraph::resume) is shorthand diff --git a/crates/tinyagents-graph/src/delegation/graph.rs b/crates/tinyagents-graph/src/delegation/graph.rs index 3b722f2b..40bcb8b3 100644 --- a/crates/tinyagents-graph/src/delegation/graph.rs +++ b/crates/tinyagents-graph/src/delegation/graph.rs @@ -208,7 +208,7 @@ where "executions": s.executions_texts(), "revisions": s.revisions, }); - tinyagents_tracing::info!( + tracing::info!( revisions = s.revisions, "[interrupt] delegation review reached durable human-approval gate; pausing" ); @@ -220,7 +220,7 @@ where } Some(decision) => { let approved = decision_is_approve(&decision); - tinyagents_tracing::info!( + tracing::info!( approved, "[interrupt] delegation review resumed with human decision" ); @@ -275,9 +275,16 @@ where .mark_command_routing("finalize"); if require_review_approval { - builder = builder - .mark_command_routing("approval") - .mark_interrupt("approval"); + // `approval` pauses *itself* (its handler returns the interrupt + // carrying the review payload), so it must not also be an + // `interrupt_before` node — `mark_interrupt` is that selector now, + // and would pause a second time ahead of the handler with a bare + // `{"phase": "before"}` payload. Annotate it for the export instead. + builder = builder.mark_command_routing("approval").with_node_metadata( + "approval", + "interrupt", + "node-emitted", + ); } let graph = builder diff --git a/crates/tinyagents-graph/src/delegation/run.rs b/crates/tinyagents-graph/src/delegation/run.rs index e7a79d8a..daf4ded1 100644 --- a/crates/tinyagents-graph/src/delegation/run.rs +++ b/crates/tinyagents-graph/src/delegation/run.rs @@ -117,7 +117,7 @@ where graph = graph.with_checkpointer(cp); } - tinyagents_tracing::info!( + tracing::info!( max_revisions = config.max_revisions, durable = thread_id.is_some(), human_gated = config.require_review_approval, @@ -200,7 +200,7 @@ where } let approved = decision_is_approve(&decision); - tinyagents_tracing::info!( + tracing::info!( approved, "[interrupt] resuming durable delegation graph with approval decision" ); @@ -273,7 +273,7 @@ where // guard must treat any version it does not explicitly recognize as // incompatible, not merely an older one. Ok(Some(checkpoint)) if checkpoint.state.schema_version != CURRENT_SCHEMA_VERSION => { - tinyagents_tracing::warn!( + tracing::warn!( thread_id = %tid, schema_version = checkpoint.state.schema_version, current = CURRENT_SCHEMA_VERSION, @@ -283,7 +283,7 @@ where run_delegation_durable(config, run_stage).await } Ok(Some(checkpoint)) if checkpoint_is_resumable(&checkpoint) => { - tinyagents_tracing::info!( + tracing::info!( thread_id = %tid, "[delegation] resuming durable delegation from its last checkpoint boundary" ); @@ -305,12 +305,12 @@ where thread_id: tid.clone(), }); if pending.is_some() { - tinyagents_tracing::warn!( + tracing::warn!( thread_id = %tid, "[delegation] terminal-classified checkpoint carried a pending interrupt; surfacing it" ); } else { - tinyagents_tracing::info!( + tracing::info!( thread_id = %tid, "[delegation] thread already terminal; returning finalized state without re-running" ); @@ -321,7 +321,7 @@ where }) } Ok(None) => { - tinyagents_tracing::debug!( + tracing::debug!( thread_id = %tid, "[delegation] no checkpoint for thread; starting a fresh durable run" ); @@ -332,7 +332,7 @@ where // must NOT silently restart a valid resumable run — it is propagated so // durable work is retried by the caller, not dropped. Err(e) if is_incompatible_checkpoint_error(&e) => { - tinyagents_tracing::warn!( + tracing::warn!( thread_id = %tid, error = %e, "[delegation] undecodable/incompatible checkpoint; pruning and starting fresh" @@ -341,7 +341,7 @@ where run_delegation_durable(config, run_stage).await } Err(e) => { - tinyagents_tracing::error!( + tracing::error!( thread_id = %tid, error = %e, "[delegation] checkpoint read failed (operational); not restarting — propagating error" @@ -357,7 +357,7 @@ where /// forever. Failure to prune is non-fatal (logged at debug). async fn prune_thread(cp: &dyn Checkpointer, thread_id: &str) { if let Err(e) = cp.delete_thread(thread_id).await { - tinyagents_tracing::debug!( + tracing::debug!( thread_id = %thread_id, error = %e, "[delegation] could not prune checkpoint thread (non-fatal)" @@ -402,7 +402,10 @@ fn checkpoint_is_resumable(checkpoint: &Checkpoint) -> bool { if checkpoint.state.final_output.is_some() { return false; } - checkpoint.next_nodes.iter().any(|n| n.as_str() != END) + // `checkpoint` was already normalized on read (every backend's decode + // path calls `Checkpoint::normalize`), so `tasks` is the single source + // of truth regardless of the stored record's original format version. + checkpoint.tasks.iter().any(|t| t.node.as_str() != END) } /// Rebuild the delegation graph (its node closures are not serializable — only @@ -460,7 +463,7 @@ fn into_outcome( thread_id: Option, ) -> DelegationOutcome { let pending = execution.interrupts.first().map(|i| { - tinyagents_tracing::info!( + tracing::info!( interrupt_id = %i.id, node = %i.node.as_str(), "[interrupt] delegation run parked on durable human-approval interrupt" diff --git a/crates/tinyagents-graph/src/delegation/test.rs b/crates/tinyagents-graph/src/delegation/test.rs index d5e6e7eb..d8980ea3 100644 --- a/crates/tinyagents-graph/src/delegation/test.rs +++ b/crates/tinyagents-graph/src/delegation/test.rs @@ -12,7 +12,7 @@ use serde_json::json; use super::run::{decision_is_approve, is_incompatible_checkpoint_error}; use super::*; use crate::Interrupt; -use crate::checkpoint::{Checkpoint, Checkpointer}; +use crate::checkpoint::{Checkpoint, Checkpointer, PendingActivation}; use tinyagents_harness::cancel::CancellationToken; /// A reviewer that rejects the first `reject_first` executions, then approves, @@ -137,6 +137,10 @@ async fn human_gated_run_parks_on_interrupt_then_resume_approves() { let pending = outcome.pending.expect("parked on the approval interrupt"); assert_eq!(pending.node, "approval"); assert_eq!(pending.thread_id, "hg-approve"); + // The pause is the node's own (carrying the review payload), not an + // executor-injected `interrupt_before` ahead of the handler. + assert_eq!(pending.interrupt_id, "delegation-review-approval"); + assert_eq!(pending.payload["kind"], "delegation_review"); assert!( outcome.state.final_output.is_none(), "must not finalize while paused for human approval" @@ -587,13 +591,8 @@ async fn incompatible_checkpoint_expires_to_a_fresh_run() { // Seed the store as the OLD state type under the thread. let legacy_cp: crate::checkpoint::FileCheckpointer = crate::checkpoint::FileCheckpointer::new(dir.path()); - let legacy = Checkpoint { - thread_id: "legacy-1".to_string(), - checkpoint_id: "cp-legacy".to_string(), - run_id: None, - parent_checkpoint_id: None, - namespace: vec![], - state: LegacyState { + let legacy = Checkpoint::new( + LegacyState { plan: Some("old".to_string()), executions: vec!["a".to_string(), "b".to_string()], reviews: vec![], @@ -602,14 +601,12 @@ async fn incompatible_checkpoint_expires_to_a_fresh_run() { final_output: None, cancelled: false, }, - next_nodes: vec![], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![], - pending_activations: None, - barrier_arrivals: vec![], - metadata: json!({}), - }; + Vec::new(), + ) + .with_thread_id("legacy-1".to_string()) + .with_checkpoint_id("cp-legacy".to_string()) + .with_parent_checkpoint_id(None) + .with_metadata(json!({})); legacy_cp.put(legacy).await.expect("seed legacy checkpoint"); // Reopen the SAME store as the current state type and resume: the @@ -638,24 +635,17 @@ async fn checkpoint_below_current_schema_version_expires_to_fresh_run() { let dir = tempfile::tempdir().unwrap(); let seed: crate::checkpoint::FileCheckpointer = crate::checkpoint::FileCheckpointer::new(dir.path()); - let checkpoint = Checkpoint { - thread_id: "old-schema".to_string(), - checkpoint_id: "cp-old".to_string(), - run_id: None, - parent_checkpoint_id: None, - namespace: vec![], - state: DelegationState { + let checkpoint = Checkpoint::new( + DelegationState { plan: Some("stale".to_string()), ..Default::default() }, - next_nodes: vec![], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![], - pending_activations: None, - barrier_arrivals: vec![], - metadata: json!({}), - }; + Vec::new(), + ) + .with_thread_id("old-schema".to_string()) + .with_checkpoint_id("cp-old".to_string()) + .with_parent_checkpoint_id(None) + .with_metadata(json!({})); assert_eq!( checkpoint.state.schema_version, 0, "an un-stamped record is version 0" @@ -699,25 +689,18 @@ async fn checkpoint_above_current_schema_version_also_expires_to_fresh_run() { let dir = tempfile::tempdir().unwrap(); let seed: crate::checkpoint::FileCheckpointer = crate::checkpoint::FileCheckpointer::new(dir.path()); - let checkpoint = Checkpoint { - thread_id: "future-schema".to_string(), - checkpoint_id: "cp-future".to_string(), - run_id: None, - parent_checkpoint_id: None, - namespace: vec![], - state: DelegationState { + let checkpoint = Checkpoint::new( + DelegationState { plan: Some("from-the-future".to_string()), schema_version: CURRENT_SCHEMA_VERSION + 1, ..Default::default() }, - next_nodes: vec![], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![], - pending_activations: None, - barrier_arrivals: vec![], - metadata: json!({}), - }; + Vec::new(), + ) + .with_thread_id("future-schema".to_string()) + .with_checkpoint_id("cp-future".to_string()) + .with_parent_checkpoint_id(None) + .with_metadata(json!({})); seed.put(checkpoint) .await .expect("seed future-schema checkpoint"); @@ -785,26 +768,23 @@ async fn cancelled_checkpoint_still_scheduling_finalize_is_resumed_not_terminal( let dir = tempfile::tempdir().unwrap(); let seed: crate::checkpoint::FileCheckpointer = crate::checkpoint::FileCheckpointer::new(dir.path()); - let checkpoint = Checkpoint { - thread_id: "cancelled-mid-flight".to_string(), - checkpoint_id: "cp-cancel".to_string(), - run_id: None, - parent_checkpoint_id: None, - namespace: vec![], - state: DelegationState { + let checkpoint = Checkpoint::new( + DelegationState { plan: Some("PLAN".to_string()), cancelled: true, schema_version: CURRENT_SCHEMA_VERSION, ..Default::default() }, - next_nodes: vec![tinyagents_harness::ids::NodeId::from("finalize")], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![], - pending_activations: None, - barrier_arrivals: vec![], - metadata: json!({}), - }; + vec![PendingActivation { + node: tinyagents_harness::ids::NodeId::from("finalize"), + send_arg: None, + task_id: tinyagents_harness::ids::TaskId::from(String::new()), + }], + ) + .with_thread_id("cancelled-mid-flight".to_string()) + .with_checkpoint_id("cp-cancel".to_string()) + .with_parent_checkpoint_id(None) + .with_metadata(json!({})); seed.put(checkpoint) .await .expect("seed cancelled-but-not-finalized checkpoint"); @@ -1003,29 +983,29 @@ async fn resume_delegation_rejects_a_schema_mismatched_checkpoint() { let dir = tempfile::tempdir().unwrap(); let seed: crate::checkpoint::FileCheckpointer = crate::checkpoint::FileCheckpointer::new(dir.path()); - let checkpoint = Checkpoint { - thread_id: "resume-future-schema".to_string(), - checkpoint_id: "cp-future".to_string(), - run_id: None, - parent_checkpoint_id: None, - namespace: vec![], - state: DelegationState { + let checkpoint = Checkpoint::new( + DelegationState { plan: Some("PLAN".to_string()), schema_version: CURRENT_SCHEMA_VERSION + 1, ..Default::default() }, - next_nodes: vec![tinyagents_harness::ids::NodeId::from("approval")], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![Interrupt { - id: "int-1".to_string(), + vec![PendingActivation { node: tinyagents_harness::ids::NodeId::from("approval"), - payload: json!({}), + send_arg: None, + task_id: tinyagents_harness::ids::TaskId::from(String::new()), }], - pending_activations: None, - barrier_arrivals: vec![], - metadata: json!({}), - }; + ) + .with_thread_id("resume-future-schema".to_string()) + .with_checkpoint_id("cp-future".to_string()) + .with_parent_checkpoint_id(None) + .with_interrupts(vec![Interrupt { + id: "int-1".to_string(), + node: tinyagents_harness::ids::NodeId::from("approval"), + payload: json!({}), + task_id: None, + response_schema: None, + }]) + .with_metadata(json!({})); seed.put(checkpoint) .await .expect("seed future-schema checkpoint parked on approval"); @@ -1168,25 +1148,15 @@ async fn terminal_checkpoint_with_a_pending_interrupt_surfaces_it() { crate::checkpoint::FileCheckpointer::new(dir.path()); let mut state = DelegationState::new_run(); state.final_output = Some("done".to_string()); - let checkpoint = Checkpoint { - thread_id: "terminal-interrupt".to_string(), - checkpoint_id: "cp-ti".to_string(), - run_id: None, - parent_checkpoint_id: None, - namespace: vec![], - state, - next_nodes: vec![], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![Interrupt::with_id( + let checkpoint = Checkpoint::new(state, Vec::new()) + .with_thread_id("terminal-interrupt") + .with_checkpoint_id("cp-ti") + .with_interrupts(vec![Interrupt::with_id( "intr-1", "approval", json!({ "kind": "delegation_review" }), - )], - pending_activations: None, - barrier_arrivals: vec![], - metadata: json!({}), - }; + )]) + .with_metadata(json!({})); seed.put(checkpoint).await.expect("seed terminal+interrupt"); let cp: Arc> = Arc::new( diff --git a/crates/tinyagents-graph/src/export/mod.rs b/crates/tinyagents-graph/src/export/mod.rs index ea0f4363..32c18970 100644 --- a/crates/tinyagents-graph/src/export/mod.rs +++ b/crates/tinyagents-graph/src/export/mod.rs @@ -442,7 +442,11 @@ impl CompiledGraph { let edges = self .edges .iter() - .map(|(from, to)| (from.to_string(), to.to_string())) + .flat_map(|(from, targets)| { + targets + .iter() + .map(move |to| (from.to_string(), to.to_string())) + }) .collect(); let conditional = self .branches @@ -497,7 +501,11 @@ impl GraphBuilder { let edges = self .edges .iter() - .map(|(from, to)| (from.to_string(), to.to_string())) + .flat_map(|(from, targets)| { + targets + .iter() + .map(move |to| (from.to_string(), to.to_string())) + }) .collect(); let conditional = self .branches diff --git a/crates/tinyagents-graph/src/language/test.rs b/crates/tinyagents-graph/src/language/test.rs new file mode 100644 index 00000000..3d92d160 --- /dev/null +++ b/crates/tinyagents-graph/src/language/test.rs @@ -0,0 +1,403 @@ +use super::*; +use tinyagents_harness::error::TinyAgentsError; +use tinyagents_language::compiler::compile; +use tinyagents_language::parser::parse_str; + +#[derive(Clone, Debug, Default, PartialEq)] +struct S { + trail: Vec, +} + +struct EchoFactory; + +impl NodeFactory for EchoFactory { + fn make(&self, spec: &NodeSpec) -> Result> { + let name = spec.name.clone(); + Ok(Arc::new(move |state: Arc, _ctx: crate::NodeContext| { + let name = name.clone(); + Box::pin(async move { + let mut state = (*state).clone(); + state.trail.push(name); + Ok(crate::NodeResult::Update(state)) + }) as crate::NodeFuture + })) + } +} + +fn blueprint(src: &str) -> Blueprint { + compile(&parse_str(src).unwrap()).unwrap().remove(0) +} + +fn node_mut<'a>(bp: &'a mut Blueprint, name: &str) -> &'a mut NodeSpec { + bp.nodes.iter_mut().find(|n| n.name == name).unwrap() +} + +#[tokio::test] +async fn build_graph_accepts_a_blueprint_with_no_ignored_fields() { + let bp = blueprint( + "graph g { start a node a { kind model next b } node b { kind model next END } }", + ); + assert_eq!(bp.start, "a"); + + let graph = build_graph::(&bp, &EchoFactory).expect("no ignored fields, graph builds"); + let run = graph.run(S::default()).await.expect("graph runs to end"); + assert_eq!(run.state.trail, vec!["a".to_string(), "b".to_string()]); +} + +#[test] +fn build_graph_lowers_options_to_interrupt_marker_and_metadata() { + let bp = blueprint( + "graph g { start a node a { kind model options [\"approve\", \"reject\"] next END } }", + ); + + let graph = build_graph::(&bp, &EchoFactory).expect("options is lowered, not rejected"); + let topology = graph.topology(); + let node = topology.nodes.iter().find(|n| n.id == "a").unwrap(); + assert!( + node.interrupt, + "options marks the node as an interrupt point" + ); + assert_eq!( + node.metadata.get("options").map(String::as_str), + Some("approve,reject") + ); +} + +#[test] +fn build_graph_lowers_node_metadata() { + let bp = blueprint( + "graph g { start a node a { kind model metadata { owner \"triage\" priority 3 } next END } }", + ); + + let graph = build_graph::(&bp, &EchoFactory).expect("metadata is lowered"); + let topology = graph.topology(); + let node = topology.nodes.iter().find(|n| n.id == "a").unwrap(); + assert_eq!( + node.metadata.get("owner").map(String::as_str), + Some("triage") + ); + assert_eq!(node.metadata.get("priority").map(String::as_str), Some("3")); +} + +#[test] +fn build_graph_lowers_sends_to_metadata_and_validates_targets() { + let bp = blueprint( + "graph g { start a \ + node a { kind model sends [send b, send c] } \ + node b { kind model next END } \ + node c { kind model next END } }", + ); + + let graph = build_graph::(&bp, &EchoFactory).expect("sends is lowered"); + let topology = graph.topology(); + let node = topology.nodes.iter().find(|n| n.id == "a").unwrap(); + assert_eq!(node.metadata.get("sends").map(String::as_str), Some("b,c")); + + // A `sends` target that is not a declared node is rejected even though + // the language compiler already validated it at compile time — this + // guards a hand-built or deserialized `Blueprint` that bypassed that + // check (`Blueprint` is `Deserialize`). + let mut tampered = bp.clone(); + node_mut(&mut tampered, "a").sends[0].target = "ghost".to_string(); + let err = build_graph::(&tampered, &EchoFactory).unwrap_err(); + match err { + TinyAgentsError::Compile(message) => { + assert!(message.contains("ghost"), "got: {message}"); + } + other => panic!("expected Compile, got {other:?}"), + } +} + +#[test] +fn build_graph_lowers_command_update_to_metadata() { + let bp = blueprint( + "graph g { start a node a { kind model command { goto END update { status \"done\" } } } }", + ); + + let graph = build_graph::(&bp, &EchoFactory).expect("command.update is lowered"); + let topology = graph.topology(); + let node = topology.nodes.iter().find(|n| n.id == "a").unwrap(); + assert_eq!( + node.metadata.get("command.update").map(String::as_str), + Some("status=done") + ); +} + +#[test] +fn build_graph_lowers_graph_level_joins_to_waiting_edges() { + let bp = blueprint( + "graph g { start a \ + node a { kind model routes { toB -> b toC -> c } } \ + node b { kind model next d } \ + node c { kind model next d } \ + node d { kind model next END } \ + join [b, c] -> d }", + ); + + let graph = build_graph::(&bp, &EchoFactory).expect("joins is lowered"); + let topology = graph.topology(); + let waiting = topology + .waiting_edges + .iter() + .find(|w| w.target == "d") + .expect("d has a waiting/barrier edge"); + assert_eq!(waiting.predecessors, vec!["b".to_string(), "c".to_string()]); +} + +#[test] +fn build_graph_lowers_node_join_sources_to_waiting_edges() { + let bp = blueprint( + "graph g { start a \ + node a { kind model routes { toB -> b toC -> c } } \ + node b { kind model next d } \ + node c { kind model next d } \ + node d { kind join sources [b, c] next END } }", + ); + + let graph = build_graph::(&bp, &EchoFactory).expect("join_sources is lowered"); + let topology = graph.topology(); + let waiting = topology + .waiting_edges + .iter() + .find(|w| w.target == "d") + .expect("d has a waiting/barrier edge"); + assert_eq!(waiting.predecessors, vec!["b".to_string(), "c".to_string()]); +} + +#[test] +fn build_graph_rejects_undeclared_join_source() { + let mut bp = blueprint( + "graph g { start a node a { kind model next b } node b { kind join sources [a] next END } }", + ); + node_mut(&mut bp, "b").join_sources[0] = "ghost".to_string(); + + let err = build_graph::(&bp, &EchoFactory).unwrap_err(); + match err { + TinyAgentsError::Compile(message) => assert!(message.contains("ghost"), "got: {message}"), + other => panic!("expected Compile, got {other:?}"), + } +} + +#[test] +fn build_graph_accepts_checkpoint_and_interrupt_policy_as_validated_noop() { + let bp = blueprint( + "graph g { start a checkpoint inherit interrupt manual node a { kind model next END } }", + ); + assert_eq!(bp.checkpoint.as_deref(), Some("inherit")); + assert_eq!(bp.interrupt.as_deref(), Some("manual")); + + // No runtime attach point exists for a bare policy name (see the + // `build_graph` docs), so this is accepted without error rather than + // silently dropped or falsely claimed as enforced. + build_graph::(&bp, &EchoFactory).expect("checkpoint/interrupt policy names do not error"); +} + +#[test] +fn build_graph_accepts_input_and_output_shapes() { + let bp = blueprint( + "graph g { start a input { question string } output { answer string } \ + node a { kind model next END } }", + ); + + build_graph::(&bp, &EchoFactory).expect("input/output is a validated no-op"); +} + +#[test] +fn build_graph_rejects_duplicate_io_field_names() { + let mut bp = blueprint("graph g { start a node a { kind model next END } }"); + bp.input.push(IoFieldSpec { + name: "question".to_string(), + ty: "string".to_string(), + }); + bp.input.push(IoFieldSpec { + name: "question".to_string(), + ty: "number".to_string(), + }); + + let err = build_graph::(&bp, &EchoFactory).unwrap_err(); + match err { + TinyAgentsError::Compile(message) => { + assert!(message.contains("question"), "got: {message}"); + } + other => panic!("expected Compile, got {other:?}"), + } +} + +#[tokio::test] +async fn build_graph_lowers_both_timeout_literal_forms() { + // `timeout` accepts a `""` string (`"30s"`) or a bare + // number of seconds (`30`) — both parse and lower onto the node's own + // `NodePolicy` without error, and a graph builds and runs normally with + // every node using a comfortably long timeout. + let bp = blueprint( + "graph g { start a node a { kind model timeout \"30s\" next b } \ + node b { kind model timeout 30 next END } }", + ); + + let graph = + build_graph::(&bp, &EchoFactory).expect("both timeout literal forms are lowered"); + let run = graph.run(S::default()).await.expect("graph runs to end"); + assert_eq!(run.state.trail, vec!["a".to_string(), "b".to_string()]); +} + +#[tokio::test] +async fn build_graph_lowers_independent_per_node_timeouts() { + // `a` declares a short timeout but never sleeps (trivially satisfies + // it); `b` declares a much longer one and sleeps just under it. Two + // nodes disagreeing on `timeout` used to be a compile-time rejection + // (no per-node timeout API); each node now gets its own `NodePolicy`, so + // this builds and `b`'s generous timeout — not `a`'s tiny one — governs + // `b`'s attempt. + let bp = blueprint( + "graph g { start a node a { kind model timeout \"5ms\" next b } \ + node b { kind model timeout \"200ms\" next END } }", + ); + + struct SleepFactory; + impl NodeFactory for SleepFactory { + fn make(&self, spec: &NodeSpec) -> Result> { + let name = spec.name.clone(); + Ok(Arc::new(move |state: Arc, _ctx: crate::NodeContext| { + let name = name.clone(); + Box::pin(async move { + if name == "b" { + tokio::time::sleep(std::time::Duration::from_millis(30)).await; + } + let mut state = (*state).clone(); + state.trail.push(name); + Ok(crate::NodeResult::Update(state)) + }) as crate::NodeFuture + })) + } + } + + let graph = build_graph::(&bp, &SleepFactory) + .expect("per-node timeouts are independent, not required to agree"); + let run = graph + .run(S::default()) + .await + .expect("b's own 200ms timeout comfortably covers its 30ms sleep"); + assert_eq!(run.state.trail, vec!["a".to_string(), "b".to_string()]); +} + +#[test] +fn build_graph_rejects_an_unsupported_retry_key() { + let bp = blueprint( + "graph g { start a node a { kind model retry { backoff \"exponential\" } next END } }", + ); + + let err = build_graph::(&bp, &EchoFactory).unwrap_err(); + match err { + TinyAgentsError::Compile(message) => { + assert!(message.contains("backoff"), "got: {message}"); + } + other => panic!("expected Compile, got {other:?}"), + } +} + +#[tokio::test] +async fn build_graph_lowers_independent_per_node_retry() { + // `a` needs exactly 2 attempts and declares `max_attempts 2` (just + // enough); `b` needs 4 attempts and declares `max_attempts 6`. Two + // nodes disagreeing on `retry` used to be a compile-time rejection (no + // per-node retry API); each node now gets its own `NodePolicy`, so this + // builds, and if `b` were incorrectly bound to `a`'s smaller cap instead + // of its own, `b` would exhaust its attempts and the run would fail. + use std::sync::atomic::{AtomicUsize, Ordering}; + + let bp = blueprint( + "graph g { start a node a { kind model retry { max_attempts 2 } next b } \ + node b { kind model retry { max_attempts 6 } next END } }", + ); + + struct FlakyFactory { + a_attempts: Arc, + b_attempts: Arc, + } + + impl NodeFactory for FlakyFactory { + fn make(&self, spec: &NodeSpec) -> Result> { + let name = spec.name.clone(); + let counter = if spec.name == "a" { + self.a_attempts.clone() + } else { + self.b_attempts.clone() + }; + // `a` succeeds on its 2nd attempt (1 failure); `b` succeeds on + // its 4th (3 failures) — only possible under `b`'s own, + // larger `max_attempts`. + let needed_failures = if spec.name == "a" { 1 } else { 3 }; + Ok(Arc::new(move |state: Arc, _ctx: crate::NodeContext| { + let name = name.clone(); + let counter = counter.clone(); + Box::pin(async move { + let n = counter.fetch_add(1, Ordering::SeqCst); + if n < needed_failures { + Err(TinyAgentsError::Model(format!("transient blip {n}"))) + } else { + let mut state = (*state).clone(); + state.trail.push(name); + Ok(crate::NodeResult::Update(state)) + } + }) as crate::NodeFuture + })) + } + } + + let factory = FlakyFactory { + a_attempts: Arc::new(AtomicUsize::new(0)), + b_attempts: Arc::new(AtomicUsize::new(0)), + }; + let graph = build_graph::(&bp, &factory) + .expect("per-node retry policies are independent, not required to agree"); + let run = graph + .run(S::default()) + .await + .expect("b's own max_attempts=6 covers the 4 attempts it needs"); + assert_eq!(run.state.trail, vec!["a".to_string(), "b".to_string()]); +} + +#[tokio::test] +async fn build_graph_lowers_uniform_node_retry_and_recovers_transient_failure() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let bp = blueprint( + "graph g { start flaky node flaky { kind model retry { max_attempts 4 } next END } }", + ); + + struct FlakyFactory { + attempts: Arc, + } + + impl NodeFactory for FlakyFactory { + fn make(&self, _spec: &NodeSpec) -> Result> { + let attempts = self.attempts.clone(); + Ok(Arc::new(move |state: Arc, _ctx: crate::NodeContext| { + let attempts = attempts.clone(); + Box::pin(async move { + let n = attempts.fetch_add(1, Ordering::SeqCst); + if n < 2 { + Err(TinyAgentsError::Model(format!("transient blip {n}"))) + } else { + let mut state = (*state).clone(); + state.trail.push("flaky".to_string()); + Ok(crate::NodeResult::Update(state)) + } + }) as crate::NodeFuture + })) + } + } + + let attempts = Arc::new(AtomicUsize::new(0)); + let factory = FlakyFactory { + attempts: attempts.clone(), + }; + let graph = + build_graph::(&bp, &factory).expect("uniform retry is lowered onto with_node_retry"); + let run = graph + .run(S::default()) + .await + .expect("the retry policy recovers the transient failure"); + assert_eq!(run.state.trail, vec!["flaky".to_string()]); + assert_eq!(attempts.load(Ordering::SeqCst), 3); +} diff --git a/crates/tinyagents-graph/src/lib.rs b/crates/tinyagents-graph/src/lib.rs index 192e6861..e2515a82 100644 --- a/crates/tinyagents-graph/src/lib.rs +++ b/crates/tinyagents-graph/src/lib.rs @@ -21,12 +21,9 @@ //! Each concern lives in its own submodule with `types.rs` (definitions), //! `mod.rs` (implementations), and `test.rs` (unit tests). -#![cfg_attr( - not(feature = "tracing"), - allow(dead_code, unused_imports, unused_variables) -)] - +pub mod agent_loop; pub mod builder; +pub mod cache; pub mod channel; pub mod checkpoint; pub mod command; @@ -52,22 +49,28 @@ pub use tinyagents_harness::error::{Result, TinyAgentsError}; // --- Durable execution model --- pub use builder::{ - END, ForkId, GraphBuilder, GraphDefaults, NodeContext, NodeFuture, NodeHandler, Route, - RouterFn, START, + END, ForkId, GraphBuilder, GraphDefaults, IdleClock, NodeCachePolicy, NodeContext, NodeFuture, + NodeHandler, NodePolicy, Route, RouterFn, START, }; +#[cfg(feature = "sqlite")] +pub use cache::SqliteTaskCache; +pub use cache::{InMemoryTaskCache, TaskCache, TaskCacheKey}; pub use channel::{ - Barrier, BinaryAggregate, Channel, ChannelSet, ChannelState, ChannelUpdate, Delta, Ephemeral, - LastValue, Messages, NamedBarrier, Topic, Untracked, + Barrier, BinaryAggregate, Channel, ChannelSet, ChannelState, ChannelUpdate, ChannelWrite, + Delta, Ephemeral, LastValue, Messages, NamedBarrier, ReducerRegistry, Topic, Untracked, }; #[cfg(feature = "sqlite")] pub use checkpoint::SqliteCheckpointer; pub use checkpoint::{ - BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, - CheckpointTuple, Checkpointer, DurabilityMode, FileCheckpointer, InMemoryCheckpointer, - PendingActivation, PendingWrite, + BarrierArrivals, CHECKPOINT_FORMAT_VERSION, Checkpoint, CheckpointConfig, CheckpointMetadata, + CheckpointSource, CheckpointTuple, Checkpointer, CompletedTask, DurabilityMode, + FileCheckpointer, InMemoryCheckpointer, PendingActivation, PendingWrite, }; pub use command::{Command, Interrupt, NodeResult, RouteTarget, Send}; -pub use compiled::{CompiledGraph, GraphExecution, GraphInput, ResumeTarget, StateSnapshot}; +pub use compiled::{ + CompiledGraph, DrainHandle, DrainSignal, GraphExecution, GraphInput, ResumeTarget, RunOptions, + StateSnapshot, +}; pub use dag::{DagIssue, DagNode}; pub use delegation::{ CURRENT_SCHEMA_VERSION as DELEGATION_SCHEMA_VERSION, DelegationConfig, DelegationOutcome, @@ -111,7 +114,10 @@ pub use reducer::{ OverwriteStateReducer, Reducer, SetUnionReducer, StateReducer, }; pub use status::GraphRunStatus; -pub use stream::{CollectingSink, GraphEvent, GraphEventSink, NoopSink, StreamMode}; +pub use stream::{ + CollectingSink, GraphEvent, GraphEventEnvelope, GraphEventSink, NoopSink, StreamMode, + StreamProjection, project_graph_event, +}; pub use subagent_node::{ AgentInvocation, AgentInvocationBinding, AgentInvoker, InputMapper, OutputMapper, SubAgentBudget, SubAgentInput, SubAgentNode, SubAgentOutput, SubAgentPolicy, subagent_node, diff --git a/crates/tinyagents-graph/src/observability/langfuse/test.rs b/crates/tinyagents-graph/src/observability/langfuse/test.rs index 5d4b6dda..ffe674c8 100644 --- a/crates/tinyagents-graph/src/observability/langfuse/test.rs +++ b/crates/tinyagents-graph/src/observability/langfuse/test.rs @@ -388,6 +388,7 @@ fn checkpoint_events_carry_coordinates_in_metadata() { 1_110, GraphEvent::CheckpointSaved { checkpoint_id: CheckpointId::new("ckpt-7"), + step: Some(8), }, ) }); diff --git a/crates/tinyagents-graph/src/observability/mod.rs b/crates/tinyagents-graph/src/observability/mod.rs index 0b8e710a..e2bb6372 100644 --- a/crates/tinyagents-graph/src/observability/mod.rs +++ b/crates/tinyagents-graph/src/observability/mod.rs @@ -55,7 +55,7 @@ use std::time::SystemTime; use async_trait::async_trait; use crate::status::GraphRunStatus; -use crate::stream::{GraphEvent, GraphEventSink}; +use crate::stream::{GraphEvent, GraphEventEnvelope, GraphEventSink}; use tinyagents_harness::error::Result; use tinyagents_harness::ids::{CheckpointId, EventId, GraphId, NodeId, RunId, ThreadId, now_ms}; use tinyagents_harness::observability::{AppendWorker, DEFAULT_DRAIN_CAPACITY}; @@ -543,12 +543,12 @@ impl JournalGraphSink { } impl GraphEventSink for JournalGraphSink { - fn emit(&self, event: GraphEvent) { - let obs = self.observe(&event); + fn emit(&self, envelope: GraphEventEnvelope) { + let obs = self.observe(&envelope.event); // Hand off to the background drain; never block the executor on I/O. self.worker.submit(obs); if let Some(inner) = &self.inner { - inner.emit(event); + inner.emit(envelope); } } @@ -560,11 +560,27 @@ impl GraphEventSink for JournalGraphSink { } } +impl JournalGraphSink { + /// Number of observations dropped because the background drain's bounded + /// queue was full when they were submitted (G-M7). + /// + /// Journaling is deliberately lossy under load — [`Self::emit`] never + /// blocks the executor waiting for durable I/O, so a burst that outpaces + /// the drain worker drops the observation rather than stalling the run. + /// A non-zero value here means the journal is an incomplete record of + /// what happened during that burst; a caller that needs a complete log + /// should watch this counter (or size the drain capacity generously for + /// its workload) rather than assume every emitted event was persisted. + pub fn dropped(&self) -> u64 { + self.worker.dropped() + } +} + /// Extracts the checkpoint id a [`GraphEvent::CheckpointSaved`] carries, so the /// observation envelope can record it directly. fn checkpoint_of(event: &GraphEvent) -> Option { match event { - GraphEvent::CheckpointSaved { checkpoint_id } => Some(checkpoint_id.clone()), + GraphEvent::CheckpointSaved { checkpoint_id, .. } => Some(checkpoint_id.clone()), _ => None, } } diff --git a/crates/tinyagents-graph/src/observability/test.rs b/crates/tinyagents-graph/src/observability/test.rs index 8cc8f484..7220e19c 100644 --- a/crates/tinyagents-graph/src/observability/test.rs +++ b/crates/tinyagents-graph/src/observability/test.rs @@ -407,15 +407,17 @@ async fn journal_sink_used_directly_forwards_to_inner() { let sink = JournalGraphSink::new(journal.clone(), RunId::new("fixed-run"), GraphId::new("g")) .with_inner(collector.clone()); - sink.emit(GraphEvent::StepStarted { + sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepStarted { step: 1, active: Vec::new(), - }); - sink.emit(GraphEvent::RouteSelected { + })); + sink.emit(GraphEventEnvelope::for_test(GraphEvent::RouteSelected { node: "a".into(), target: "b".into(), - }); - sink.emit(GraphEvent::StepCompleted { step: 1 }); + })); + sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepCompleted { + step: 1, + })); // Forwarded to the live sink. assert_eq!(collector.len(), 3); diff --git a/crates/tinyagents-graph/src/orchestration/reconcile.rs b/crates/tinyagents-graph/src/orchestration/reconcile.rs index a29f2be3..4bb1ad45 100644 --- a/crates/tinyagents-graph/src/orchestration/reconcile.rs +++ b/crates/tinyagents-graph/src/orchestration/reconcile.rs @@ -141,7 +141,7 @@ pub fn reconcile_orphaned_tasks( }; if let ReconcileOutcome::Error(detail) = &outcome { - tinyagents_tracing::warn!( + tracing::warn!( task_id = %task_id.as_str(), prior_status = task_status_label(prior_status), error = %detail, diff --git a/crates/tinyagents-graph/src/orchestration/store_registry.rs b/crates/tinyagents-graph/src/orchestration/store_registry.rs index 180a6cb1..638c8083 100644 --- a/crates/tinyagents-graph/src/orchestration/store_registry.rs +++ b/crates/tinyagents-graph/src/orchestration/store_registry.rs @@ -158,7 +158,7 @@ pub fn open_jsonl_task_store_or_memory(path: &Path) -> Arc { if let Some(parent) = path.parent() && let Err(err) = std::fs::create_dir_all(parent) { - tinyagents_tracing::warn!( + tracing::warn!( dir = %parent.display(), error = %err, "[orchestration] task store directory unavailable; falling back to memory" @@ -168,14 +168,14 @@ pub fn open_jsonl_task_store_or_memory(path: &Path) -> Arc { match JsonlTaskStore::open(path) { Ok(store) => { - tinyagents_tracing::debug!( + tracing::debug!( path = %path.display(), "[orchestration] opened durable task store" ); Arc::new(store) } Err(err) => { - tinyagents_tracing::warn!( + tracing::warn!( path = %path.display(), error = %err, "[orchestration] durable task store unavailable; falling back to memory" diff --git a/crates/tinyagents-graph/src/orchestration/test.rs b/crates/tinyagents-graph/src/orchestration/test.rs index 164824bf..7921ec15 100644 --- a/crates/tinyagents-graph/src/orchestration/test.rs +++ b/crates/tinyagents-graph/src/orchestration/test.rs @@ -40,7 +40,7 @@ fn raw(result: &ToolResult) -> &serde_json::Value { .iter() .find_map(|content| match content { ToolContent::Json { data } => Some(data), - ToolContent::Text { .. } => None, + ToolContent::Text { .. } | ToolContent::Image { .. } | ToolContent::File { .. } => None, }) .expect("orchestration tool returns a JSON payload") } diff --git a/crates/tinyagents-graph/src/recursion/types.rs b/crates/tinyagents-graph/src/recursion/types.rs index a97e31a7..5c8fb894 100644 --- a/crates/tinyagents-graph/src/recursion/types.rs +++ b/crates/tinyagents-graph/src/recursion/types.rs @@ -10,7 +10,7 @@ use std::sync::{Arc, Mutex}; use serde::{Deserialize, Serialize}; use crate::{Result, TinyAgentsError}; -use tinyagents_harness::ids::{GraphId, NodeId, RunId, TaskId}; +use tinyagents_harness::ids::{CheckpointId, GraphId, NodeId, RunId, TaskId}; /// One level of the graph/subgraph/sub-agent recursion tree. /// @@ -201,6 +201,14 @@ pub struct ChildRun { /// on the parent [`GraphExecution`](crate::GraphExecution) rollup. #[serde(default)] pub usage: tinyinference_llm::usage::UsageTotals, + /// The child run's latest persisted checkpoint id, when checkpointing was + /// enabled (C4). Recorded so the parent's own checkpoint metadata + /// (`child_runs`) carries an explicit pointer to the exact child + /// checkpoint a subsequent `drive_child` continuation + /// (retry/resume) would act on, rather than leaving the association + /// implicit in the shared thread id + namespace. + #[serde(default)] + pub checkpoint_id: Option, } /// A thread-safe collector the executor hands to node contexts so that a diff --git a/crates/tinyagents-graph/src/status/mod.rs b/crates/tinyagents-graph/src/status/mod.rs index c327c0ee..6d1d73dc 100644 --- a/crates/tinyagents-graph/src/status/mod.rs +++ b/crates/tinyagents-graph/src/status/mod.rs @@ -42,11 +42,17 @@ impl GraphRunStatus { } } - /// Returns true when the run is in a terminal state. + /// Returns true when the run is in a terminal state — one this run id + /// will never advance from on its own. `Drained` counts (like + /// `Cancelled`): the process stopped the run; continuing the thread is a + /// new run started by `resume`/`retry`. `Interrupted` does not. pub fn is_terminal(&self) -> bool { matches!( self.status, - ExecutionStatus::Completed | ExecutionStatus::Failed | ExecutionStatus::Cancelled + ExecutionStatus::Completed + | ExecutionStatus::Failed + | ExecutionStatus::Cancelled + | ExecutionStatus::Drained ) } } diff --git a/crates/tinyagents-graph/src/stream/mod.rs b/crates/tinyagents-graph/src/stream/mod.rs index 8c95c102..e38ca9b1 100644 --- a/crates/tinyagents-graph/src/stream/mod.rs +++ b/crates/tinyagents-graph/src/stream/mod.rs @@ -13,16 +13,27 @@ //! [`GraphEvent`]s into an optional [`GraphEventSink`]; callers can plug in a //! [`NoopSink`], a test-friendly [`CollectingSink`], or any custom transport. +pub mod project; mod types; -pub use types::{GraphEvent, StreamMode}; +pub use project::{ + Cursored, MessageEntry, ProjectedSince, StreamProjection, SubagentEntry, SubagentPhase, + ToolCallEntry, ToolCallPhase, project_graph_event, +}; +pub use types::{GraphEvent, GraphEventEnvelope, StreamMode}; use std::sync::{Arc, Mutex}; /// A pluggable target for low-level graph events. +/// +/// Every event is delivered wrapped in a [`GraphEventEnvelope`], which +/// carries the run id, checkpoint namespace, and a monotonic sequence number +/// alongside the [`GraphEvent`] itself — see [`GraphEventEnvelope`] for what +/// each field means and how it is scoped. pub trait GraphEventSink: Send + Sync { - /// Receives one graph event. Implementations must not block the executor. - fn emit(&self, event: GraphEvent); + /// Receives one enveloped graph event. Implementations must not block the + /// executor. + fn emit(&self, envelope: GraphEventEnvelope); /// Blocks until every event emitted so far has been durably handled. /// @@ -38,13 +49,13 @@ pub trait GraphEventSink: Send + Sync { pub struct NoopSink; impl GraphEventSink for NoopSink { - fn emit(&self, _event: GraphEvent) {} + fn emit(&self, _envelope: GraphEventEnvelope) {} } /// A sink that records every event for inspection in tests and UIs. #[derive(Clone, Default)] pub struct CollectingSink { - events: Arc>>, + events: Arc>>, } impl CollectingSink { @@ -53,8 +64,21 @@ impl CollectingSink { Self::default() } - /// Returns a clone of the recorded events. + /// Returns a clone of the recorded events, discarding their envelopes. + /// + /// The pre-C3 shape most callers (mostly tests) still want: event kind + /// and payload only, with no run/namespace/sequence attribution. Use + /// [`Self::envelopes`] when that attribution matters. pub fn events(&self) -> Vec { + self.envelopes() + .into_iter() + .map(|envelope| envelope.event) + .collect() + } + + /// Returns a clone of the recorded envelopes (event plus run/namespace/ + /// sequence attribution). + pub fn envelopes(&self) -> Vec { self.events.lock().map(|g| g.clone()).unwrap_or_default() } @@ -70,9 +94,9 @@ impl CollectingSink { } impl GraphEventSink for CollectingSink { - fn emit(&self, event: GraphEvent) { + fn emit(&self, envelope: GraphEventEnvelope) { if let Ok(mut guard) = self.events.lock() { - guard.push(event); + guard.push(envelope); } } } diff --git a/crates/tinyagents-graph/src/stream/project.rs b/crates/tinyagents-graph/src/stream/project.rs new file mode 100644 index 00000000..4988795d --- /dev/null +++ b/crates/tinyagents-graph/src/stream/project.rs @@ -0,0 +1,300 @@ +//! [`StreamMode`] filtering for [`GraphEvent`]s, and [`StreamProjection`] — a +//! cursor-ordered fold of both [`GraphEventEnvelope`]s and harness +//! [`AgentEvent`]s into the three views a UI actually renders: messages, tool +//! calls, and subagent activity. +//! +//! Graph events narrate *structure* (which node/task ran, which checkpoint +//! saved); harness events narrate *content* (what the model said, which tool +//! ran, which subagent was invoked). A consumer watching a recursive run — +//! a graph whose nodes drive harness agent loops, some of which spawn +//! subagents or embed subgraphs — wants both folded into one ordered view. +//! [`StreamProjection`] is that fold. Its [`StreamProjection::cursor`] is a +//! single monotonic counter shared by every view, so a consumer that +//! attaches after a run has already produced output can request only what it +//! missed with [`StreamProjection::since`] instead of re-reading everything. + +use tinyagents_harness::events::AgentEvent; +use tinyagents_harness::ids::{CallId, RunId}; +use tinyinference_llm::message::MessageDelta; + +use super::{GraphEvent, GraphEventEnvelope, StreamMode}; + +/// Returns `true` when `event` should be delivered to a consumer subscribed +/// to `modes`. +/// +/// [`GraphEvent::mode`] gives the single narrow mode most event kinds belong +/// to; the run/step lifecycle events that have none (`RunStarted`, +/// `StepStarted`, …) are debug-only detail and pass only when +/// [`StreamMode::Debug`] is active — mirroring +/// [`tinyagents_harness::stream::project_event_for_modes`]'s treatment of its +/// own lifecycle events. +pub fn project_graph_event(event: &GraphEvent, modes: &[StreamMode]) -> bool { + match event.mode() { + Some(mode) => modes.contains(&mode) || modes.contains(&StreamMode::Debug), + None => modes.contains(&StreamMode::Debug), + } +} + +// --------------------------------------------------------------------------- +// StreamProjection +// --------------------------------------------------------------------------- + +/// One item in a [`StreamProjection`] view, tagged with the projection's +/// monotonic [`StreamProjection::cursor`] value at the moment it was folded +/// in. +#[derive(Clone, Debug, PartialEq)] +pub struct Cursored { + /// This item's position in the projection's global fold order. + pub cursor: u64, + /// The projected value. + pub value: T, +} + +/// One entry in [`StreamProjection::messages`]: an assistant message +/// fragment attributed to its run and model call. +#[derive(Clone, Debug, PartialEq)] +pub struct MessageEntry { + /// The run that produced this fragment. + pub run_id: RunId, + /// The model call this fragment belongs to. + pub call_id: CallId, + /// The incremental text/reasoning/tool-call fragment. + pub delta: MessageDelta, +} + +/// Lifecycle phase of a tool call, folded from [`AgentEvent::ToolStarted`] / +/// [`AgentEvent::ToolCompleted`] / [`AgentEvent::ToolFailed`]. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ToolCallPhase { + /// [`AgentEvent::ToolStarted`]. + Started, + /// [`AgentEvent::ToolCompleted`], successful (`error` was `None`). + Completed, + /// [`AgentEvent::ToolCompleted`] with `error: Some(_)`, or + /// [`AgentEvent::ToolFailed`]. + Failed { + /// The failure message. + error: String, + }, +} + +/// One entry in [`StreamProjection::tool_calls`]. +#[derive(Clone, Debug, PartialEq)] +pub struct ToolCallEntry { + /// Correlates with the call's `Started`/terminal pair. + pub call_id: CallId, + /// The tool's name. + pub tool_name: String, + /// The call's current lifecycle phase. + pub phase: ToolCallPhase, +} + +/// Lifecycle phase of a subagent or subgraph activation. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SubagentPhase { + /// Started (harness [`AgentEvent::SubAgentStarted`] or graph + /// [`GraphEvent::SubgraphStarted`]). + Started, + /// Finished (harness [`AgentEvent::SubAgentCompleted`] or graph + /// [`GraphEvent::SubgraphCompleted`]). + Completed, +} + +/// One entry in [`StreamProjection::subagents`]. +#[derive(Clone, Debug, PartialEq)] +pub struct SubagentEntry { + /// The sub-agent's name (harness activations) or hosting node id (graph + /// subgraph activations). + pub name: String, + /// The activation's current lifecycle phase. + pub phase: SubagentPhase, +} + +/// Folds a run's [`GraphEventEnvelope`]s and harness [`AgentEvent`]s into +/// three consumer-facing views (`messages`, `tool_calls`, `subagents`) under +/// one monotonic cursor. +/// +/// Feed events as they arrive with [`Self::fold_graph_event`] / +/// [`Self::fold_agent_event`], in the order they were emitted (across both +/// sources merged by real time — the projection does not reorder). A +/// consumer that attaches late replays with [`Self::since`] instead of +/// re-reading the full history. +#[derive(Clone, Debug, Default)] +pub struct StreamProjection { + next_cursor: u64, + /// Assistant message fragments, in fold order. + pub messages: Vec>, + /// Tool-call lifecycle entries, in fold order. A call's `Started` and + /// terminal phase are two separate entries sharing `call_id`, not one + /// mutated in place, so [`Self::since`] replay never has to reconstruct + /// history a consumer already saw. + pub tool_calls: Vec>, + /// Subagent/subgraph lifecycle entries, in fold order (same + /// two-entries-per-activation shape as `tool_calls`). + pub subagents: Vec>, +} + +impl StreamProjection { + /// Creates an empty projection. + pub fn new() -> Self { + Self::default() + } + + /// The cursor of the most recently folded item, or zero before any item. + pub fn cursor(&self) -> u64 { + self.next_cursor + } + + fn next(&mut self) -> u64 { + self.next_cursor += 1; + self.next_cursor + } + + /// Folds one graph event. Only [`GraphEvent::SubgraphStarted`] / + /// [`GraphEvent::SubgraphCompleted`] currently project onto a view (as + /// `subagents`); every other kind is structural and is not part of the + /// three content views this projection exposes (subscribe to the raw + /// envelope stream directly for those). + pub fn fold_graph_event(&mut self, envelope: &GraphEventEnvelope) { + match &envelope.event { + GraphEvent::SubgraphStarted { node, .. } => { + self.push_subagent(node.to_string(), SubagentPhase::Started); + } + GraphEvent::SubgraphCompleted { node, .. } => { + self.push_subagent(node.to_string(), SubagentPhase::Completed); + } + _ => {} + } + } + + /// Folds one harness agent event. + pub fn fold_agent_event(&mut self, event: &AgentEvent) { + match event { + AgentEvent::ModelDelta { + run_id, + call_id, + delta, + } => { + let cursor = self.next(); + self.messages.push(Cursored { + cursor, + value: MessageEntry { + run_id: run_id.clone(), + call_id: call_id.clone(), + delta: delta.clone(), + }, + }); + } + AgentEvent::ToolStarted { call_id, tool_name } => { + self.push_tool_call(call_id.clone(), tool_name.clone(), ToolCallPhase::Started); + } + AgentEvent::ToolCompleted { + call_id, + tool_name, + error, + .. + } => { + let phase = match error { + Some(error) => ToolCallPhase::Failed { + error: error.clone(), + }, + None => ToolCallPhase::Completed, + }; + self.push_tool_call(call_id.clone(), tool_name.clone(), phase); + } + AgentEvent::ToolFailed { + call_id, + tool_name, + error, + .. + } => { + self.push_tool_call( + call_id.clone(), + tool_name.clone(), + ToolCallPhase::Failed { + error: error.clone(), + }, + ); + } + AgentEvent::SubAgentStarted { name, .. } => { + self.push_subagent(name.clone(), SubagentPhase::Started); + } + AgentEvent::SubAgentCompleted { name, .. } => { + self.push_subagent(name.clone(), SubagentPhase::Completed); + } + _ => {} + } + } + + fn push_tool_call(&mut self, call_id: CallId, tool_name: String, phase: ToolCallPhase) { + let cursor = self.next(); + self.tool_calls.push(Cursored { + cursor, + value: ToolCallEntry { + call_id, + tool_name, + phase, + }, + }); + } + + fn push_subagent(&mut self, name: String, phase: SubagentPhase) { + let cursor = self.next(); + self.subagents.push(Cursored { + cursor, + value: SubagentEntry { name, phase }, + }); + } + + /// Returns every item across all three views with `cursor > since`, each + /// still tagged with its view, in cursor order — what a late-attaching + /// consumer replays instead of re-reading the full projection. + pub fn since(&self, since: u64) -> Vec { + let mut items: Vec = self + .messages + .iter() + .filter(|item| item.cursor > since) + .map(|item| ProjectedSince::Message(item.clone())) + .chain( + self.tool_calls + .iter() + .filter(|item| item.cursor > since) + .map(|item| ProjectedSince::ToolCall(item.clone())), + ) + .chain( + self.subagents + .iter() + .filter(|item| item.cursor > since) + .map(|item| ProjectedSince::Subagent(item.clone())), + ) + .collect(); + items.sort_by_key(ProjectedSince::cursor); + items + } +} + +/// One replayed item from [`StreamProjection::since`], tagged by which view +/// it belongs to. +#[derive(Clone, Debug, PartialEq)] +pub enum ProjectedSince { + /// A [`StreamProjection::messages`] entry. + Message(Cursored), + /// A [`StreamProjection::tool_calls`] entry. + ToolCall(Cursored), + /// A [`StreamProjection::subagents`] entry. + Subagent(Cursored), +} + +impl ProjectedSince { + /// The item's cursor value, regardless of which view it came from. + pub fn cursor(&self) -> u64 { + match self { + ProjectedSince::Message(item) => item.cursor, + ProjectedSince::ToolCall(item) => item.cursor, + ProjectedSince::Subagent(item) => item.cursor, + } + } +} + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-graph/src/stream/project/test.rs b/crates/tinyagents-graph/src/stream/project/test.rs new file mode 100644 index 00000000..4832e5a6 --- /dev/null +++ b/crates/tinyagents-graph/src/stream/project/test.rs @@ -0,0 +1,172 @@ +use tinyagents_harness::events::AgentEvent; +use tinyagents_harness::ids::{CallId, NodeId, RunId}; +use tinyinference_llm::message::MessageDelta; + +use super::*; + +fn envelope(event: GraphEvent) -> GraphEventEnvelope { + GraphEventEnvelope { + run_id: RunId::from("run-1".to_string()), + task_id: None, + ns: Vec::new(), + seq: 0, + event, + } +} + +#[test] +fn project_graph_event_routes_task_events_to_the_tasks_mode() { + let event = GraphEvent::NodeStarted { + node: NodeId::from("n".to_string()), + step: 1, + }; + assert!(project_graph_event(&event, &[StreamMode::Tasks])); + assert!(!project_graph_event(&event, &[StreamMode::Checkpoints])); + // Debug always sees everything, including narrow-mode events. + assert!(project_graph_event(&event, &[StreamMode::Debug])); +} + +#[test] +fn project_graph_event_lifecycle_events_are_debug_only() { + let event = GraphEvent::RunStarted { + run_id: RunId::from("run-1".to_string()), + }; + assert!(!project_graph_event(&event, &[StreamMode::Tasks])); + assert!(!project_graph_event(&event, &[StreamMode::Checkpoints])); + assert!(project_graph_event(&event, &[StreamMode::Debug])); +} + +#[test] +fn project_graph_event_routes_checkpoint_events_to_the_checkpoints_mode() { + let event = GraphEvent::CheckpointSaved { + checkpoint_id: "ckpt-1".to_string().into(), + step: Some(2), + }; + assert!(project_graph_event(&event, &[StreamMode::Checkpoints])); + assert!(!project_graph_event(&event, &[StreamMode::Updates])); +} + +#[test] +fn stream_projection_folds_model_deltas_into_messages_in_order() { + let mut projection = StreamProjection::new(); + projection.fold_agent_event(&AgentEvent::ModelDelta { + run_id: RunId::from("run-1".to_string()), + call_id: CallId::from("call-1".to_string()), + delta: MessageDelta::text("hel"), + }); + projection.fold_agent_event(&AgentEvent::ModelDelta { + run_id: RunId::from("run-1".to_string()), + call_id: CallId::from("call-1".to_string()), + delta: MessageDelta::text("lo"), + }); + + assert_eq!(projection.messages.len(), 2); + assert_eq!(projection.messages[0].cursor, 1); + assert_eq!(projection.messages[1].cursor, 2); + assert_eq!(projection.messages[0].value.delta.text, "hel"); + assert_eq!(projection.messages[1].value.delta.text, "lo"); + assert_eq!(projection.cursor(), 2); +} + +#[test] +fn stream_projection_folds_tool_lifecycle_as_two_entries() { + let mut projection = StreamProjection::new(); + projection.fold_agent_event(&AgentEvent::ToolStarted { + call_id: CallId::from("call-1".to_string()), + tool_name: "search".into(), + }); + projection.fold_agent_event(&AgentEvent::ToolCompleted { + call_id: CallId::from("call-1".to_string()), + tool_name: "search".into(), + started_at_ms: None, + input: None, + output: None, + duration_ms: None, + output_bytes: None, + error: None, + metadata: None, + }); + + assert_eq!(projection.tool_calls.len(), 2); + assert_eq!(projection.tool_calls[0].value.phase, ToolCallPhase::Started); + assert_eq!( + projection.tool_calls[1].value.phase, + ToolCallPhase::Completed + ); + assert_eq!(projection.tool_calls[0].value.call_id.as_str(), "call-1"); +} + +#[test] +fn stream_projection_folds_failed_tool_completion_as_failed_phase() { + let mut projection = StreamProjection::new(); + projection.fold_agent_event(&AgentEvent::ToolCompleted { + call_id: CallId::from("call-1".to_string()), + tool_name: "search".into(), + started_at_ms: None, + input: None, + output: None, + duration_ms: None, + output_bytes: None, + error: Some("boom".into()), + metadata: None, + }); + assert_eq!( + projection.tool_calls[0].value.phase, + ToolCallPhase::Failed { + error: "boom".into() + } + ); +} + +#[test] +fn stream_projection_folds_subgraph_events_as_subagents() { + let mut projection = StreamProjection::new(); + projection.fold_graph_event(&envelope(GraphEvent::SubgraphStarted { + node: NodeId::from("researcher".to_string()), + namespace: vec!["researcher".into()], + })); + projection.fold_graph_event(&envelope(GraphEvent::SubgraphCompleted { + node: NodeId::from("researcher".to_string()), + namespace: vec!["researcher".into()], + })); + + assert_eq!(projection.subagents.len(), 2); + assert_eq!(projection.subagents[0].value.name, "researcher"); + assert_eq!(projection.subagents[0].value.phase, SubagentPhase::Started); + assert_eq!( + projection.subagents[1].value.phase, + SubagentPhase::Completed + ); +} + +#[test] +fn stream_projection_since_replays_only_items_after_the_given_cursor() { + let mut projection = StreamProjection::new(); + // Cursors start at one, while zero is the empty-projection sentinel. A + // snapshot can therefore be passed directly to `since` without dropping + // the first event subsequently folded. + projection.fold_agent_event(&AgentEvent::ToolStarted { + call_id: CallId::from("call-1".to_string()), + tool_name: "search".into(), + }); + let cursor_after_first = projection.cursor(); + projection.fold_agent_event(&AgentEvent::ModelDelta { + run_id: RunId::from("run-1".to_string()), + call_id: CallId::from("call-2".to_string()), + delta: MessageDelta::text("hi"), + }); + projection.fold_graph_event(&envelope(GraphEvent::SubgraphStarted { + node: NodeId::from("n".to_string()), + namespace: vec!["n".into()], + })); + + let replay = projection.since(0); + assert_eq!(replay.len(), 3, "every item after the empty cursor"); + + let replay = projection.since(cursor_after_first); + assert_eq!(replay.len(), 2, "everything after the first item"); + assert!(matches!(replay[1], ProjectedSince::Subagent(_))); + + // Nothing new since the last item's own cursor. + assert!(projection.since(projection.cursor()).is_empty()); +} diff --git a/crates/tinyagents-graph/src/stream/test.rs b/crates/tinyagents-graph/src/stream/test.rs index 25b43a94..847434fe 100644 --- a/crates/tinyagents-graph/src/stream/test.rs +++ b/crates/tinyagents-graph/src/stream/test.rs @@ -8,11 +8,13 @@ use tinyagents_harness::ids::NodeId; fn collecting_sink_records_events() { let sink = CollectingSink::new(); assert!(sink.is_empty()); - sink.emit(GraphEvent::StepStarted { + sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepStarted { step: 1, active: vec![NodeId::from("a")], - }); - sink.emit(GraphEvent::StepCompleted { step: 1 }); + })); + sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepCompleted { + step: 1, + })); assert_eq!(sink.len(), 2); assert!(matches!( sink.events()[0], @@ -23,5 +25,7 @@ fn collecting_sink_records_events() { #[test] fn noop_sink_drops_events() { let sink = NoopSink; - sink.emit(GraphEvent::StepCompleted { step: 1 }); + sink.emit(GraphEventEnvelope::for_test(GraphEvent::StepCompleted { + step: 1, + })); } diff --git a/crates/tinyagents-graph/src/stream/types.rs b/crates/tinyagents-graph/src/stream/types.rs index ea41994e..c2745679 100644 --- a/crates/tinyagents-graph/src/stream/types.rs +++ b/crates/tinyagents-graph/src/stream/types.rs @@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize}; use crate::command::Interrupt; -use tinyagents_harness::ids::{CheckpointId, NodeId, RunId}; +use tinyagents_harness::ids::{CheckpointId, NodeId, RunId, TaskId}; /// A low-level graph lifecycle event emitted through a [`super::GraphEventSink`]. /// @@ -42,6 +42,23 @@ pub enum GraphEvent { /// Rendered error. error: String, }, + /// The run was cooperatively cancelled via a [`tinyagents_harness::CancellationToken`] + /// (I4 part 2), either between supersteps or while a superstep's node + /// handlers were still in flight. + RunCancelled { + /// The run that was cancelled. + run_id: RunId, + }, + /// The run stopped gracefully at a superstep boundary because its + /// [`crate::DrainSignal`] was raised: the step in flight finished and + /// committed, and the next step's activations were checkpointed instead + /// of run (see [`crate::GraphExecution::drained`]). + RunDrained { + /// The run that drained. + run_id: RunId, + /// The superstep count at which it stopped (the last completed step). + steps: usize, + }, /// A superstep started with the given active node set. StepStarted { /// 1-based step number. @@ -61,6 +78,34 @@ pub enum GraphEvent { /// Step number. step: usize, }, + /// A task began executing (the [`StreamMode::Tasks`] counterpart of + /// [`GraphEvent::NodeStarted`], emitted alongside it at the same + /// boundary). + TaskStarted { + /// Target node. + node: NodeId, + /// Step number. + step: usize, + }, + /// A task finished, successfully or not (the [`StreamMode::Tasks`] + /// counterpart of [`GraphEvent::NodeCompleted`]/[`GraphEvent::NodeFailed`], + /// emitted alongside them at the same boundary). Also the cache-aware + /// signal for nodes with an opt-in [`crate::NodeCachePolicy`] (see + /// [`crate::CompiledGraph::with_cached_node`]): `cached: true` means the + /// handler was skipped and a stored `Update` was replayed in its place, + /// substituting for the handler's normal `NodeStarted`/`NodeCompleted` + /// pair; `cached: false` means the handler ran (and, on success, its + /// result was written back to the [`crate::cache::TaskCache`] when one is + /// attached). + TaskCompleted { + /// Target node. + node: NodeId, + /// Step number. + step: usize, + /// Whether this task's result came from the cache rather than + /// executing the handler. + cached: bool, + }, /// A node handler began executing. NodeStarted { /// Node id. @@ -113,6 +158,11 @@ pub enum GraphEvent { CheckpointSaved { /// Persisted checkpoint id. checkpoint_id: CheckpointId, + /// The superstep this checkpoint was saved at, when the save site + /// knows it (`None` for saves outside the ordinary superstep boundary, + /// such as a resume-time bootstrap checkpoint). + #[serde(default, skip_serializing_if = "Option::is_none")] + step: Option, }, /// A checkpoint was loaded to resume/replay a run (a read, not a write). CheckpointRestored { @@ -172,9 +222,13 @@ impl GraphEvent { GraphEvent::RunStarted { .. } => "run.started", GraphEvent::RunCompleted { .. } => "run.completed", GraphEvent::RunFailed { .. } => "run.failed", + GraphEvent::RunCancelled { .. } => "run.cancelled", + GraphEvent::RunDrained { .. } => "run.drained", GraphEvent::StepStarted { .. } => "step.started", GraphEvent::StepCompleted { .. } => "step.completed", GraphEvent::TaskScheduled { .. } => "task.scheduled", + GraphEvent::TaskStarted { .. } => "task.started", + GraphEvent::TaskCompleted { .. } => "task.completed", GraphEvent::NodeStarted { .. } => "node.started", GraphEvent::NodeCompleted { .. } => "node.completed", GraphEvent::NodeFailed { .. } => "node.failed", @@ -200,13 +254,17 @@ impl GraphEvent { GraphEvent::StepStarted { step, .. } | GraphEvent::StepCompleted { step } | GraphEvent::TaskScheduled { step, .. } + | GraphEvent::TaskStarted { step, .. } + | GraphEvent::TaskCompleted { step, .. } | GraphEvent::NodeStarted { step, .. } | GraphEvent::NodeCompleted { step, .. } | GraphEvent::NodeFailed { step, .. } | GraphEvent::NodeRetryScheduled { step, .. } | GraphEvent::StateUpdated { step, .. } | GraphEvent::ContextForked { step, .. } => Some(*step), - GraphEvent::RunCompleted { steps, .. } => Some(*steps), + GraphEvent::RunCompleted { steps, .. } | GraphEvent::RunDrained { steps, .. } => { + Some(*steps) + } _ => None, } } @@ -214,9 +272,13 @@ impl GraphEvent { /// High-level projection modes for a graph run stream. /// -/// These mirror the LangGraph stream modes. The milestone executor exposes them -/// as a selection enum; richer typed `StreamPart` projection is future work. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +/// These mirror the LangGraph stream modes. [`GraphEvent::mode`] maps every +/// event kind onto one of these (or `None` for the lifecycle events every +/// mode should still see); [`super::project::project_graph_event`] applies +/// that mapping to filter a raw [`GraphEventEnvelope`] stream the way +/// [`tinyagents_harness::stream::project_event_for_modes`] does for +/// [`tinyagents_harness::events::AgentEvent`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum StreamMode { /// Full state values after each step. Values, @@ -230,4 +292,95 @@ pub enum StreamMode { Interrupts, /// Arbitrary user stream writes from inside nodes. Custom, + /// Task-level lifecycle: [`GraphEvent::TaskScheduled`], + /// [`GraphEvent::TaskStarted`], and [`GraphEvent::TaskCompleted`] — the + /// LangGraph `"tasks"` mode, narrower than [`StreamMode::Debug`] (no + /// step/checkpoint/routing internals, just task start/end). + Tasks, + /// Checkpoint lifecycle only: [`GraphEvent::CheckpointSaved`] and + /// [`GraphEvent::CheckpointRestored`] — the LangGraph `"checkpoints"` + /// mode. + Checkpoints, +} + +impl GraphEvent { + /// Returns the [`StreamMode`] this event projects onto, when it belongs + /// to a narrower mode than [`StreamMode::Debug`] (which every event kind + /// still counts toward — see + /// [`super::project::project_graph_event`]). + /// + /// Run/step lifecycle events (`RunStarted`, `StepStarted`, …) have no + /// narrower home and return `None`: they surface only under + /// [`StreamMode::Debug`]. + pub fn mode(&self) -> Option { + match self { + GraphEvent::TaskScheduled { .. } + | GraphEvent::TaskStarted { .. } + | GraphEvent::TaskCompleted { .. } + | GraphEvent::NodeStarted { .. } + | GraphEvent::NodeCompleted { .. } + | GraphEvent::NodeFailed { .. } + | GraphEvent::NodeRetryScheduled { .. } => Some(StreamMode::Tasks), + GraphEvent::StateUpdated { .. } => Some(StreamMode::Updates), + GraphEvent::CheckpointSaved { .. } | GraphEvent::CheckpointRestored { .. } => { + Some(StreamMode::Checkpoints) + } + GraphEvent::InterruptEmitted { .. } => Some(StreamMode::Interrupts), + GraphEvent::Custom { .. } => Some(StreamMode::Custom), + _ => None, + } + } +} + +// --------------------------------------------------------------------------- +// GraphEventEnvelope +// --------------------------------------------------------------------------- + +/// A [`GraphEvent`] wrapped with the run/task correlation and ordering +/// metadata every emission site needs to be attributable in a merged, +/// multi-run stream. +/// +/// `run_id` and `ns` (the checkpoint namespace) identify which run — and +/// which level of subgraph nesting within it — emitted the event, so a +/// parent run's observer can tell its own events apart from a nested +/// subgraph's. `seq` is a monotonic counter scoped to the emitting +/// [`crate::compiled::CompiledGraph`] instance (shared across a clone that +/// only changes `event_sink`, such as journal wrapping, but **not** shared +/// between a parent graph and a subgraph embedded as a node — the subgraph's +/// [`Self::ns`] already distinguishes its stream). `task_id` is `None` until +/// per-task correlation ids land end-to-end +/// (`docs/runtime-comparison/feature-gaps.md` D4); the field exists now so +/// adding that id later is additive. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct GraphEventEnvelope { + /// The run that emitted this event. + pub run_id: RunId, + /// Correlation id for the task this event belongs to, when task ids are + /// wired end to end. `None` today. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub task_id: Option, + /// Checkpoint namespace of the emitting graph instance (empty for a + /// top-level run; one segment deeper per level of subgraph nesting). + pub ns: Vec, + /// Monotonically increasing sequence number, scoped as described on + /// [`Self`]. + pub seq: u64, + /// The wrapped event. + pub event: GraphEvent, +} + +impl GraphEventEnvelope { + /// Wraps `event` in a minimal envelope (empty run id/namespace, `seq: + /// 0`, no task id) for tests that only care about the event payload + /// reaching a sink, not its attribution. + #[cfg(test)] + pub(crate) fn for_test(event: GraphEvent) -> Self { + Self { + run_id: RunId::from(String::new()), + task_id: None, + ns: Vec::new(), + seq: 0, + event, + } + } } diff --git a/crates/tinyagents-graph/src/subagent_node/mod.rs b/crates/tinyagents-graph/src/subagent_node/mod.rs index d0b7d09b..b870ef0a 100644 --- a/crates/tinyagents-graph/src/subagent_node/mod.rs +++ b/crates/tinyagents-graph/src/subagent_node/mod.rs @@ -188,6 +188,7 @@ fn record_child_run(ctx: &NodeContext, agent: &str, output: &SubAgentOutput) { run_id: RunId::new(format!("subagent-{}", next_seq())), root_run_id, usage: output.usage, + checkpoint_id: None, }); } diff --git a/crates/tinyagents-graph/src/subgraph/mod.rs b/crates/tinyagents-graph/src/subgraph/mod.rs index 14325f2b..17c45983 100644 --- a/crates/tinyagents-graph/src/subgraph/mod.rs +++ b/crates/tinyagents-graph/src/subgraph/mod.rs @@ -106,10 +106,32 @@ where /// Clones `child` and extends its checkpoint namespace with the embedding node /// id, preventing parent/child checkpoint collisions. +/// +/// I1: when this node is activated more than once in the same superstep (a +/// `Send` fan-out of a subgraph node — map-reduce over a subgraph), each +/// concurrent activation gets its own namespace (`[node_id, task_id]`) +/// instead of sharing one (`[node_id]`) across every fan-out branch. Sharing +/// one namespace is what let N concurrent activations write interleaved +/// lineages under the same key and made every fan-out branch's `resume` +/// non-deterministically pick up whichever child checkpoint was written +/// last. With exactly one activation (`ctx.siblings <= 1`, the overwhelmingly +/// common case) the namespace stays `[node_id]` so existing checkpoints +/// remain readable — this is purely additive. fn namespaced(child: &CompiledGraph, ctx: &NodeContext) -> CompiledGraph { let mut namespace = child.namespace().to_vec(); namespace.push(ctx.node_id.to_string()); - child.clone().with_namespace(namespace) + if ctx.siblings > 1 { + namespace.push(ctx.task_id().as_str().to_string()); + } + // A fresh sequence counter: the embedded instance's own namespace already + // disambiguates its event stream from the parent's (and from any sibling + // fan-out activation of this same node), so there is no benefit — only + // entanglement — in sharing the parent's `seq` counter. See + // `crate::stream::GraphEventEnvelope`. + child + .clone() + .with_namespace(namespace) + .with_fresh_sequence() } /// Prepares an embedded `child` graph for a subgraph run: extends its checkpoint @@ -123,6 +145,69 @@ fn child_for(child: &CompiledGraph, ctx: &NodeContext) -> CompiledGr .with_recursion_node(ctx.node_id.clone()) } +/// What an existing child checkpoint (if any) says to do instead of running +/// fresh — the C4 fix. +enum ChildContinuation { + /// A prior activation left a resumable failure-boundary checkpoint: + /// `retry` it so the child's already-completed nodes (and their side + /// effects) do not re-run. + Retry, + /// A prior activation left an interrupted checkpoint and the parent was + /// handed a resume value for this activation: resume the child with it. + Resume(serde_json::Value), +} + +/// Checks the child's own checkpoint namespace for a record left by an +/// earlier activation of this same subgraph node, before a fresh run would +/// otherwise discard it (C4). +/// +/// `child.run_with_thread(..)` on a thread that already has a child +/// checkpoint does not resume it — it starts a second, root-less lineage +/// under the same namespace, silently re-running the child's completed nodes +/// (and their side effects) and orphaning its partial progress. This is what +/// let a parent `retry()` after a subgraph node failure restart the child +/// from scratch instead of continuing it. Consulted by [`drive_child`] +/// before every fresh-run path (not only after a failure): a checkpoint +/// stamped `failed_node` always retries; one stamped `interrupted_nodes` +/// only resumes when the caller supplied a resume value — otherwise (no +/// checkpoint, or an interrupted one with no resume value) the caller's +/// original fresh/resume decision stands. +async fn child_continuation( + child: &CompiledGraph, + thread_id: &tinyagents_harness::ids::ThreadId, + resume: Option<&serde_json::Value>, +) -> Result> +where + S: Clone + Send + Sync + 'static, + U: Send + 'static, +{ + let Some(checkpointer) = child.checkpointer.as_ref() else { + return Ok(None); + }; + let Some(checkpoint) = checkpointer + .get_scoped(thread_id.as_str(), None, child.namespace()) + .await? + else { + return Ok(None); + }; + // `checkpoint` was already normalized on read (every backend's decode + // path calls `Checkpoint::normalize`), so `tasks` is the single source + // of truth regardless of the stored record's original format version. + let has_pending = !checkpoint.tasks.is_empty(); + if !has_pending { + return Ok(None); + } + if checkpoint.metadata.get("failed_node").is_some() { + return Ok(Some(ChildContinuation::Retry)); + } + if checkpoint.metadata.get("interrupted_nodes").is_some() + && let Some(value) = resume + { + return Ok(Some(ChildContinuation::Resume(value.clone()))); + } + Ok(None) +} + /// Drives an embedded child graph for one parent-node activation. /// /// On a fresh activation (`resume == None`) the child runs from `state`. On a @@ -132,6 +217,11 @@ fn child_for(child: &CompiledGraph, ctx: &NodeContext) -> CompiledGr /// re-running the child (which would just re-interrupt forever). Resuming /// requires the child to have run under a thread; without one, a paused child /// could not have persisted, so we fall back to a fresh run. +/// +/// C4: on a threaded child, [`child_continuation`] is consulted first — a +/// failed child is retried (not restarted) and an interrupted child with a +/// resume value in hand is resumed, regardless of which of the branches below +/// the caller's own `(resume, binding)` shape would otherwise have taken. async fn drive_child( child: CompiledGraph, thread_id: Option, @@ -143,6 +233,25 @@ where S: Clone + Send + Sync + 'static, U: Send + 'static, { + if let Some(thread_id) = &thread_id + && let Some(continuation) = child_continuation(&child, thread_id, resume.as_ref()).await? + { + let thread_id = thread_id.clone(); + return match (continuation, binding) { + (ChildContinuation::Retry, Some(binding)) => { + child.retry_with_agent_binding(thread_id, binding).await + } + (ChildContinuation::Retry, None) => child.retry(thread_id).await, + (ChildContinuation::Resume(value), Some(binding)) => { + child + .resume_with_agent_binding(thread_id, Command::resume(value), binding) + .await + } + (ChildContinuation::Resume(value), None) => { + child.resume(thread_id, Command::resume(value)).await + } + }; + } match (thread_id, resume, binding) { (Some(thread_id), None, Some(binding)) => { child @@ -198,6 +307,11 @@ impl ChildRunRecorder { run_id: execution.run_id.clone(), root_run_id: execution.root_run_id.clone(), usage: tinyinference_llm::usage::UsageTotals::default(), + // C4: the child's latest checkpoint, so the parent's own + // checkpoint metadata (`child_runs`) carries an explicit + // pointer to the exact record a later `retry`/`resume` + // continuation would act on. + checkpoint_id: execution.checkpoint_id.clone(), }); } } diff --git a/crates/tinyagents-graph/src/subgraph/test.rs b/crates/tinyagents-graph/src/subgraph/test.rs index 3ac0a04e..fd69a350 100644 --- a/crates/tinyagents-graph/src/subgraph/test.rs +++ b/crates/tinyagents-graph/src/subgraph/test.rs @@ -80,6 +80,14 @@ fn nested_binding( /// Builds a minimal [`NodeContext`] standing in for the embedding node `id`. fn ctx_for(id: &str) -> NodeContext { + ctx_for_task(id, "task-test", 1) +} + +/// Builds a minimal [`NodeContext`] standing in for a `Send` fan-out +/// activation of embedding node `id`: `task_id` names this activation and +/// `siblings` is the fan-out width (I1), which is what a subgraph node +/// consults to decide whether to namespace its child checkpoint by task id. +fn ctx_for_task(id: &str, task_id: &str, siblings: usize) -> NodeContext { NodeContext { graph_id: tinyagents_harness::ids::GraphId::new("graph-test"), node_id: NodeId::from(id), @@ -93,6 +101,12 @@ fn ctx_for(id: &str) -> NodeContext { recursion_frames: Vec::new(), child_runs: None, agent_binding: None, + task_id: tinyagents_harness::ids::TaskId::from(task_id), + siblings, + channel_versions: Default::default(), + versions_seen: Default::default(), + idle_clock: Default::default(), + durable_writes: Default::default(), } } @@ -939,3 +953,256 @@ async fn child_runs_recorded_in_checkpoint_metadata() { } assert!(found, "child_runs not found in any checkpoint metadata"); } + +// ---- I1: Send fan-out of a subgraph node gets its own checkpoint namespace, +// and R5: task-scoped interrupt/resume identity ----------------------- + +#[tokio::test] +async fn send_fanout_of_subgraph_node_gets_per_task_namespaces_and_resume() { + // A `Send` fan-out of three activations of one subgraph node, each of + // whose children interrupts: the parent must surface all three as + // distinct task-scoped interrupts (not just the lowest-index one), their + // children must persist under three distinct checkpoint namespaces (not + // one shared `["child"]` namespace all three interleave into), and a + // per-task resume map must deliver each activation its own value. + let ckpt = Arc::new(InMemoryCheckpointer::::new()); + let seen = Arc::new(std::sync::Mutex::new(Vec::::new())); + let seen_child = seen.clone(); + let child = GraphBuilder::::overwrite() + .add_node("gate", move |s: i32, c: NodeContext| { + let seen_child = seen_child.clone(); + async move { + match c.resume { + Some(v) => { + seen_child.lock().unwrap().push(v.as_i64().unwrap()); + Ok(NodeResult::Update(s)) + } + None => Ok(NodeResult::Interrupt(crate::command::Interrupt::new( + "gate", + serde_json::json!({ "ask": "ok?" }), + ))), + } + } + }) + .set_entry("gate") + .set_finish("gate") + .compile() + .unwrap() + .with_checkpointer(ckpt.clone()); + + let parent = GraphBuilder::::overwrite() + .with_parallel(true) + .add_node("dispatch", |_s: i32, _c: NodeContext| async move { + Ok(NodeResult::Command(crate::command::Command::send([ + crate::command::Send::new("child", serde_json::json!(0)), + crate::command::Send::new("child", serde_json::json!(1)), + crate::command::Send::new("child", serde_json::json!(2)), + ]))) + }) + .add_node("child", shared_subgraph_node(child)) + .set_entry("dispatch") + .mark_command_routing("dispatch") + .set_finish("child") + .compile() + .unwrap() + .with_checkpointer(ckpt.clone()); + + let paused = parent.run_with_thread("t", 0).await.unwrap(); + assert!( + paused.is_interrupted(), + "every fan-out branch's child interrupted" + ); + assert_eq!( + paused.interrupts.len(), + 3, + "all three fan-out branches are surfaced, not only the lowest-index one" + ); + let task_ids: HashSet = paused + .interrupts + .iter() + .map(|i| { + i.task_id + .clone() + .expect("stamped with its own branch's task id (R5)") + .as_str() + .to_string() + }) + .collect(); + assert_eq!( + task_ids.len(), + 3, + "each fan-out branch's interrupt carries a distinct task id" + ); + + // Each branch's child persisted under its own namespace (I1): three + // distinct `["child", task_id]` namespaces, not one shared `["child"]`. + let list = ckpt.list("t").await.unwrap(); + let child_namespaces: HashSet> = list + .iter() + .filter(|m| m.namespace.first().map(String::as_str) == Some("child")) + .map(|m| m.namespace.clone()) + .collect(); + assert_eq!( + child_namespaces.len(), + 3, + "each fan-out branch's child checkpoints live under a distinct namespace" + ); + for ns in &child_namespaces { + assert_eq!( + ns.len(), + 2, + "namespace is [node_id, task_id] once the node fans out (I1): {ns:?}" + ); + } + + // Resume every branch with its own value in one call (I1). + let pairs: Vec<(tinyagents_harness::ids::TaskId, serde_json::Value)> = paused + .interrupts + .iter() + .enumerate() + .map(|(i, interrupt)| { + ( + interrupt.task_id.clone().unwrap(), + serde_json::json!(100 + i as i64), + ) + }) + .collect(); + let done = parent + .resume("t", crate::command::Command::resume_tasks(pairs)) + .await + .unwrap(); + assert!(!done.is_interrupted()); + + let mut delivered = seen.lock().unwrap().clone(); + delivered.sort_unstable(); + assert_eq!( + delivered, + vec![100, 101, 102], + "each activation's child received its own resume value" + ); +} + +// ---- C4: a subgraph child failure is resumable through the parent -------- + +#[tokio::test] +async fn parent_retry_after_subgraph_child_failure_resumes_not_restarts() { + // The child increments a shared side-effect counter on its first node, + // then fails on its second. A parent `retry()` must continue the child + // from its own resumable checkpoint (not restart it from scratch), so + // the first node's side effect never runs twice. + let ckpt = Arc::new(InMemoryCheckpointer::::new()); + let counter = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let should_fail = Arc::new(std::sync::atomic::AtomicBool::new(true)); + let counter_for_node = counter.clone(); + let should_fail_for_node = should_fail.clone(); + let child = GraphBuilder::::overwrite() + .add_node("bump", move |s: i32, _c: NodeContext| { + let counter = counter_for_node.clone(); + async move { + counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(NodeResult::Update(s + 1)) + } + }) + .add_node("maybe_fail", move |s: i32, _c: NodeContext| { + let should_fail = should_fail_for_node.clone(); + async move { + if should_fail.swap(false, std::sync::atomic::Ordering::SeqCst) { + return Err(crate::TinyAgentsError::Graph("boom".to_string())); + } + Ok(NodeResult::Update(s + 1)) + } + }) + .set_entry("bump") + .add_edge("bump", "maybe_fail") + .set_finish("maybe_fail") + .compile() + .unwrap() + .with_checkpointer(ckpt.clone()); + + let parent = GraphBuilder::::overwrite() + .add_node("child", shared_subgraph_node(child)) + .set_entry("child") + .set_finish("child") + .compile() + .unwrap() + .with_checkpointer(ckpt.clone()); + + let failed = parent.run_with_thread("t", 0).await; + assert!( + failed.is_err(), + "the child's node failure aborts the parent run" + ); + assert_eq!( + counter.load(std::sync::atomic::Ordering::SeqCst), + 1, + "the child's first node ran exactly once before its second node failed" + ); + + let done = parent + .retry("t") + .await + .expect("retry must continue the child from its resumable checkpoint"); + assert!(!done.is_interrupted()); + assert_eq!( + counter.load(std::sync::atomic::Ordering::SeqCst), + 1, + "C4: retry must not re-run the child's already-completed node" + ); + // bump(+1) -> maybe_fail(+1) = 2, once retried past the failure. + assert_eq!(done.state, 2); +} + +#[tokio::test] +async fn nested_subgraph_run_yields_envelopes_with_correct_namespace_depth_and_seq() { + // C3: a subgraph run's envelopes must carry the deeper `ns` of the + // embedding node, and each graph instance's own `seq` counter must be + // strictly increasing within its emitted stream. Both the child and the + // parent are wired to the same collecting sink here — as they must be + // today for nested observability, since a subgraph node does not + // automatically inherit the parent's `event_sink` (D4 tracks true + // end-to-end task/observability propagation as future work). + let collector = Arc::new(crate::stream::CollectingSink::new()); + let child = child_add_ten().with_event_sink(collector.clone()); + let parent = GraphBuilder::::overwrite() + .add_node("child", shared_subgraph_node(child)) + .set_entry("child") + .set_finish("child") + .compile() + .unwrap() + .with_event_sink(collector.clone()); + + parent.run_with_thread("t", 0).await.unwrap(); + + let envelopes = collector.envelopes(); + assert!(!envelopes.is_empty()); + + let parent_ns: Vec<_> = envelopes.iter().filter(|e| e.ns.is_empty()).collect(); + let child_ns: Vec<_> = envelopes + .iter() + .filter(|e| e.ns == vec!["child".to_string()]) + .collect(); + assert!(!parent_ns.is_empty(), "some events at the top-level ns"); + assert!(!child_ns.is_empty(), "some events at the child's deeper ns"); + + // Every event belongs to one of exactly these two namespaces — there is + // no third, unexpected depth. + assert_eq!(parent_ns.len() + child_ns.len(), envelopes.len()); + + // Each graph instance's own sequence is strictly increasing. + let parent_seqs: Vec = parent_ns.iter().map(|e| e.seq).collect(); + let mut sorted_parent = parent_seqs.clone(); + sorted_parent.sort_unstable(); + assert_eq!(parent_seqs, sorted_parent); + assert!(parent_seqs.windows(2).all(|w| w[0] < w[1])); + + let child_seqs: Vec = child_ns.iter().map(|e| e.seq).collect(); + let mut sorted_child = child_seqs.clone(); + sorted_child.sort_unstable(); + assert_eq!(child_seqs, sorted_child); + assert!(child_seqs.windows(2).all(|w| w[0] < w[1])); + + // The child's own sequence starts fresh (not continuing the parent's), + // per `namespaced`'s `with_fresh_sequence` — its distinct `ns` already + // disambiguates the stream. + assert_eq!(child_seqs[0], 0); +} diff --git a/crates/tinyagents-graph/src/testkit/conformance.rs b/crates/tinyagents-graph/src/testkit/conformance.rs index 91b18862..1d9b5081 100644 --- a/crates/tinyagents-graph/src/testkit/conformance.rs +++ b/crates/tinyagents-graph/src/testkit/conformance.rs @@ -8,7 +8,7 @@ //! Each function panics with a descriptive message on the first violation, so //! call them from a `#[tokio::test]` / `#[test]`. -use crate::checkpoint::{Checkpoint, Checkpointer}; +use crate::checkpoint::{Checkpoint, Checkpointer, PendingActivation}; use crate::orchestration::{ OrchestrationTaskFilter, OrchestrationTaskKind, OrchestrationTaskResult, OrchestrationTaskStatus, TaskStore, @@ -21,21 +21,18 @@ fn contract_checkpoint( parent: Option<&str>, step: usize, ) -> Checkpoint { - Checkpoint { - thread_id: thread.to_string(), - checkpoint_id: id.to_string(), - run_id: None, - parent_checkpoint_id: parent.map(str::to_string), - namespace: vec![], - state: step as i32, - next_nodes: vec![NodeId::from("n")], - completed_tasks: vec![], - pending_writes: vec![], - interrupts: vec![], - pending_activations: None, - barrier_arrivals: vec![], - metadata: serde_json::json!({ "source": "loop", "step": step }), - } + Checkpoint::new( + step as i32, + vec![PendingActivation { + node: NodeId::from("n"), + send_arg: None, + task_id: TaskId::from(String::new()), + }], + ) + .with_thread_id(thread.to_string()) + .with_checkpoint_id(id.to_string()) + .with_parent_checkpoint_id(parent.map(str::to_string)) + .with_metadata(serde_json::json!({ "source": "loop", "step": step })) } /// Runs the [`Checkpointer`] contract against `cp`. diff --git a/crates/tinyagents-graph/src/testkit/mod.rs b/crates/tinyagents-graph/src/testkit/mod.rs index d9611047..a2ced066 100644 --- a/crates/tinyagents-graph/src/testkit/mod.rs +++ b/crates/tinyagents-graph/src/testkit/mod.rs @@ -270,6 +270,7 @@ where )), root_run_id, usage, + checkpoint_id: None, }); } Ok(NodeResult::Update(update)) diff --git a/crates/tinyagents-graph/src/todos/dispatch/registry.rs b/crates/tinyagents-graph/src/todos/dispatch/registry.rs index e1ad9ff8..2c90c7bf 100644 --- a/crates/tinyagents-graph/src/todos/dispatch/registry.rs +++ b/crates/tinyagents-graph/src/todos/dispatch/registry.rs @@ -90,7 +90,7 @@ impl ActiveRunRegistry { if let Some(run_id) = run_id { match runs.get(thread_id) { None => { - tinyagents_tracing::debug!( + tracing::debug!( thread_id = %thread_id, request_run_id = %run_id, "[graph:todos:dispatch] scoped cancel ignored: no active run on thread" @@ -98,7 +98,7 @@ impl ActiveRunRegistry { return None; } Some(active) if active.run_id != run_id => { - tinyagents_tracing::debug!( + tracing::debug!( thread_id = %thread_id, request_run_id = %run_id, active_run_id = %active.run_id, diff --git a/crates/tinyagents-graph/src/todos/runs/store.rs b/crates/tinyagents-graph/src/todos/runs/store.rs index 61d2ed3b..4db2204d 100644 --- a/crates/tinyagents-graph/src/todos/runs/store.rs +++ b/crates/tinyagents-graph/src/todos/runs/store.rs @@ -122,7 +122,7 @@ pub async fn create_run( runs.push(run.clone()); save(store, &thread_id, &runs).await?; - tinyagents_tracing::info!( + tracing::info!( thread_id = %thread_id, run_id = %run.run_id, card_id = %card_id, @@ -207,7 +207,7 @@ pub async fn complete_run( let completed = run.clone(); save(store, &thread_id, &runs).await?; - tinyagents_tracing::info!( + tracing::info!( thread_id = %thread_id, run_id = %run_id, outcome = ?completed.outcome, @@ -314,7 +314,7 @@ pub async fn reclaim_stale( ) .await { - tinyagents_tracing::warn!( + tracing::warn!( thread_id = %thread_id, run_id = %run.run_id, %error, @@ -356,7 +356,7 @@ pub async fn reclaim_stale( reason: reason.clone(), new_card_status: status.as_str().to_string(), }); - tinyagents_tracing::info!( + tracing::info!( thread_id = %thread_id, run_id = %run.run_id, card_id = %run.card_id, @@ -366,7 +366,7 @@ pub async fn reclaim_stale( "[graph:todos:runs] card reclaimed" ); } - Err(error) => tinyagents_tracing::warn!( + Err(error) => tracing::warn!( thread_id = %thread_id, run_id = %run.run_id, card_id = %run.card_id, @@ -398,7 +398,7 @@ pub fn spawn_heartbeat_task( tokio::select! { _ = ticker.tick() => { if let Err(error) = update_heartbeat(&store, &thread_id, &run_id).await { - tinyagents_tracing::debug!( + tracing::debug!( thread_id = %thread_id, run_id = %run_id, %error, @@ -408,7 +408,7 @@ pub fn spawn_heartbeat_task( } } _ = cancel.changed() => { - tinyagents_tracing::debug!( + tracing::debug!( thread_id = %thread_id, run_id = %run_id, "[graph:todos:runs] heartbeat cancelled" diff --git a/crates/tinyagents-graph/src/todos/test.rs b/crates/tinyagents-graph/src/todos/test.rs index 6056a508..b4c5b6d0 100644 --- a/crates/tinyagents-graph/src/todos/test.rs +++ b/crates/tinyagents-graph/src/todos/test.rs @@ -526,7 +526,9 @@ mod tool_tests { .iter() .find_map(|block| match block { ToolContent::Json { data } => Some(data), - ToolContent::Text { .. } => None, + ToolContent::Text { .. } | ToolContent::Image { .. } | ToolContent::File { .. } => { + None + } }) .expect("successful todo result has a JSON payload") } diff --git a/crates/tinyagents-harness/Cargo.toml b/crates/tinyagents-harness/Cargo.toml index 9fec13f3..918db3b9 100644 --- a/crates/tinyagents-harness/Cargo.toml +++ b/crates/tinyagents-harness/Cargo.toml @@ -5,51 +5,59 @@ version.workspace = true edition.workspace = true license.workspace = true repository.workspace = true +rust-version.workspace = true description = "Provider-neutral model, tool, middleware, and agent-loop runtime." [dependencies] -anyhow = "1" -async-trait = "0.1" +anyhow = { workspace = true } +async-trait = { workspace = true } # Keep this aligned with reqwest's transitive version. Claude Code request # rendering and multimodal data URIs both require it. base64 = "0.22" -chrono = { version = "0.4", features = ["serde"] } +chrono = { workspace = true } chrono-tz = { version = "0.10", optional = true } +dirs = { version = "6", optional = true } flate2 = { version = "1", optional = true } -futures = "0.3" -dirs = "6" -log = "0.4" +futures = { workspace = true } regex = "1" -reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream", "http2"] } -rusqlite = { version = "0.40", features = ["bundled"], optional = true } -serde = { version = "1", features = ["derive"] } -serde_json = "1" +reqwest = { workspace = true, features = ["stream", "http2"], optional = true } +rusqlite = { workspace = true, optional = true } +serde = { workspace = true } +serde_json = { workspace = true } sha2 = "0.11" thiserror = "2" -tinyagents-tracing = { path = "../tinyagents-tracing", version = "2.1.2", default-features = false } +tracing = { workspace = true } tinyagents-definition = { path = "../tinyagents-definition", version = "2.1.2" } tinytools-agent = { path = "../../vendor/tinytools/crates/tinytools-agent", version = "0.3.0", default-features = false } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.3.0" } tinyinference-embeddings = { path = "../../vendor/tinyinference/crates/tinyinference-embeddings", version = "0.3.0" } tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.3.0" } -tokio = { version = "1", default-features = false, features = ["sync", "time", "macros", "rt", "rt-multi-thread", "fs", "io-util", "process"] } -tempfile = "3" -wait-timeout = "0.2" -uuid = { version = "1", features = ["v4"] } - -[target.'cfg(unix)'.dependencies] -libc = "0.2" +tokio = { workspace = true, features = ["sync", "time", "macros", "rt", "rt-multi-thread", "fs", "io-util", "process"] } +tempfile = { workspace = true } +wait-timeout = { version = "0.2", optional = true } +uuid = { workspace = true, optional = true } [features] -default = [] +default = ["claude-code", "langfuse"] sqlite = ["dep:rusqlite"] -tools = ["dep:chrono-tz"] -multimodal = ["dep:flate2"] -tracing = ["tinyagents-tracing/tracing", "tinytools-agent/tracing"] +# `tools` is a deprecated alias kept so downstream feature forwards that +# still spell out the old name keep compiling. +tools = ["builtin-tools"] +builtin-tools = ["dep:chrono-tz"] +multimodal = ["dep:flate2", "dep:reqwest"] +# Gates the Claude Code CLI and Claude Agent SDK provider adapters, which are +# the only consumers of `uuid`, `tempfile` beyond the artifact tests, `dirs`, +# and `wait-timeout` in this crate. +claude-code = ["dep:uuid", "dep:wait-timeout", "dep:dirs"] +# Gates the Langfuse observability exporter and its `reqwest` transport. +langfuse = ["dep:reqwest"] +# Tracing instrumentation is now always compiled in (via the `tracing` crate +# dependency above). This feature is retained as a no-op so downstream +# feature forwards keep compiling. +tracing = ["tinytools-agent/tracing"] [dev-dependencies] -tempfile = "3" -tokio = { version = "1", features = ["macros", "rt-multi-thread", "time", "test-util"] } +tokio = { workspace = true, default-features = true, features = ["macros", "rt-multi-thread", "time", "test-util"] } [lints] workspace = true diff --git a/crates/tinyagents-harness/src/agent_loop/README.md b/crates/tinyagents-harness/src/agent_loop/README.md index 7d3a4585..76203f3a 100644 --- a/crates/tinyagents-harness/src/agent_loop/README.md +++ b/crates/tinyagents-harness/src/agent_loop/README.md @@ -42,18 +42,39 @@ A turn's tool calls are driven in three phases — serial **admission** (cancellation/deadline/limit checks, `before_tool`, unknown-tool policy, schema validation, `ToolStarted`), **execution**, and a serial **fold** in original call order (`after_tool`, `ToolCompleted`, transcript append). - -When a turn requests two or more tools and **no tool-wrap middleware** -(`ToolMiddleware`) is registered, execution runs concurrently (`join_all`), -so turn latency is the slowest tool instead of the sum. Tool-wrap middleware -holds `&mut RunContext` across each wrapped call — part of its public -contract — so its presence keeps the historical serial path. In both modes -results are attached to their original `tool_call_id` in the calls' original -order, every call's `ToolStarted` precedes its `ToolCompleted`, and -`ToolCompleted` events are emitted in call order. The first failing call (in -call order) fails the turn; in concurrent mode already-launched siblings run -to completion before the error surfaces. See `tools.rs` for the full design -notes. +The fold also records a result's host-only `metadata` on the event and on +`AgentRun::tool_metadata`, and hands back its `follow_up` content as a user +message that the batch driver appends only after the batch's last tool row +(B2) — a provider requires every tool row to follow its assistant row +directly, so follow-ups never interleave with tool rows. The same holds for +the deferred-resume batch in `apply_deferred_results`. + +Execution runs concurrently only when *all* of the following hold: the turn +requests two or more tools, zero tool-wrap middleware (`ToolMiddleware`) is +registered, and every call's tool reports `Tool::is_concurrency_safe() == +true` (the trait default is `false`, so a tool must opt in). See +`should_execute_tools_concurrently` and `batch_is_canonical_parallel_safe` in +`tools.rs`. Tool-wrap middleware holds `&mut RunContext` across each wrapped +call — part of its public contract — so its presence keeps the historical +serial path. Lifecycle middleware does **not** force the serial path: every +`before_tool` hook runs during serial admission, which completes in full for +every call in the batch before any concurrent future is built, so there is +nothing left for a lifecycle middleware to mutate once execution starts +(I-8; this used to force serial execution unconditionally). + +When concurrency does trigger, the batch runs via `futures::stream::iter(..) +.buffered(n)` — not an unbounded `join_all` — where `n` is +`RunPolicy::limits.max_tool_concurrency` (unbounded, i.e. every eligible call +starts at once, when `None`, the default). `buffered` yields results in input +order, same as `join_all` did, so downstream folding is unaffected; it just +caps how many calls are in flight simultaneously. Turn latency is then the +slowest *batch* of at most `n` tools instead of the slowest single tool. In +both modes results are attached to their original `tool_call_id` in the +calls' original order, every call's `ToolStarted` precedes its +`ToolCompleted`, and `ToolCompleted` events are emitted in call order. The +first failing call (in call order) fails the turn; in concurrent mode +already-launched siblings run to completion before the error surfaces. See +`tools.rs` for the full design notes. ## Limits @@ -64,6 +85,22 @@ surfaces as `TinyAgentsError::Timeout`. The run context's own `limits::LimitTracker` is also advanced so its counters stay consistent with the enforced caps. +## Cancellation and wall-clock bounding + +Every host/provider I/O boundary on the loop path (model resolution, budget +admission and usage recording, tool authorization, tool-output screening, +host turn preparation, the unary provider call) races cooperative +cancellation against an optional wall-clock deadline through one shared +helper, `context::RunContext::bounded(deadline, fut, timeout_message)`, +instead of each call site copying its own `tokio::select! { biased; _ = +cancelled() => .., _ = timeout(remaining, fut) => .. }` block (R-1 from +`docs/runtime-comparison/code-review-harness.md`). `timeout_message` is a +closure so the call-specific message is only built on the timeout path, not +on every call. The streaming provider loop's per-chunk pull races +cancellation against `stream.next()` directly — its future yields an +`Option`, not a `Result`, so it does not fit `bounded`'s signature and stays +a bespoke `select!`. + ## Backoff Retry backoff durations are *computed* via diff --git a/crates/tinyagents-harness/src/agent_loop/deferred_test.rs b/crates/tinyagents-harness/src/agent_loop/deferred_test.rs new file mode 100644 index 00000000..ecbc7b67 --- /dev/null +++ b/crates/tinyagents-harness/src/agent_loop/deferred_test.rs @@ -0,0 +1,643 @@ +//! Tests for deferred tool calls (A2): the loop exiting with +//! `AgentRun::deferred`, resuming with `DeferredToolResults`, external tools, +//! the inline `DeferredToolHandler` path, and `HumanApprovalMiddleware`'s +//! `ApprovalOutcome::Defer`. + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use serde_json::json; + +use crate::context::{RunConfig, RunContext}; +use crate::error::TinyAgentsError; +use crate::events::AgentEvent; +use crate::ids::CallId; +use crate::runtime::AgentHarness; +use crate::testkit::EventRecorder; +use crate::tool::{DeferredToolRequests, DeferredToolResults}; +use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; +use tinyinference_llm::model::ModelResponse; +use tinyinference_llm::providers::MockModel; +use tinyinference_llm::tool::ToolCall; +use tinyinference_llm::usage::Usage; +use tinytools::{Tool, ToolPolicy, ToolResult}; + +// ── Helpers ───────────────────────────────────────────────────────────────── + +/// A tool that records the arguments it ran with and returns a fixed reply. +/// `policy` lets a test declare `approval_required`. +struct RecordingTool { + name: &'static str, + reply: &'static str, + policy: ToolPolicy, + seen: Mutex>, +} + +impl RecordingTool { + fn plain(name: &'static str, reply: &'static str) -> Arc { + Arc::new(Self { + name, + reply, + policy: ToolPolicy::read_only(), + seen: Mutex::new(Vec::new()), + }) + } + + fn approval_gated(name: &'static str, reply: &'static str) -> Arc { + Arc::new(Self { + name, + reply, + policy: ToolPolicy::classified().requiring_approval(), + seen: Mutex::new(Vec::new()), + }) + } + + fn calls(&self) -> Vec { + self.seen.lock().unwrap().clone() + } +} + +#[async_trait] +impl Tool for RecordingTool { + fn name(&self) -> &str { + self.name + } + fn description(&self) -> &str { + "recording tool" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object"}) + } + fn policy(&self) -> ToolPolicy { + self.policy.clone() + } + async fn execute(&self, arguments: serde_json::Value) -> anyhow::Result { + self.seen.lock().unwrap().push(arguments); + Ok(ToolResult::success(self.reply)) + } +} + +fn response(tool_calls: Vec, text: &str) -> ModelResponse { + let content = if text.is_empty() { + Vec::new() + } else { + vec![ContentBlock::Text(text.to_string())] + }; + ModelResponse { + message: AssistantMessage { + id: None, + content, + tool_calls, + usage: Some(Usage::new(1, 1)), + origin: None, + }, + usage: Some(Usage::new(1, 1)), + finish_reason: Some("stop".to_string()), + raw: None, + resolved_model: None, + continue_turn: None, + served_from_cache: false, + correlation: None, + resolved_route: None, + } +} + +/// The two-call batch every deferral test starts from: `delete` needs +/// approval, `lookup` does not. +fn mixed_batch() -> ModelResponse { + response( + vec![ + ToolCall::new("call-delete", "delete", json!({"path": "/tmp/x"})), + ToolCall::new("call-lookup", "lookup", json!({"q": "x"})), + ], + "", + ) +} + +fn tool_result_text(messages: &[Message], call_id: &str) -> Option { + messages.iter().find_map(|message| match message { + Message::Tool(tool) if tool.tool_call_id == call_id => Some(message.text()), + _ => None, + }) +} + +// ── Deferral ──────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn approval_required_call_defers_the_run_after_its_siblings_execute() { + let recorder = EventRecorder::new(); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + mixed_batch(), + response(Vec::new(), "never reached"), + ])), + ); + let delete = RecordingTool::approval_gated("delete", "deleted"); + let lookup = RecordingTool::plain("lookup", "found"); + harness.register_tool(delete.clone()); + harness.register_tool(lookup.clone()); + + let ctx = RunContext::new(RunConfig::new("defer"), ()).with_events(recorder.sink()); + let result = harness + .invoke_in_context_with_status(&(), ctx, vec![Message::user("go")]) + .await + .expect("a deferral is not an error"); + let run = result.run; + + // The non-deferred sibling ran and its result is on the transcript; the + // assistant's tool-call row is intact and the deferred call is unanswered. + assert_eq!(lookup.calls().len(), 1); + assert!( + delete.calls().is_empty(), + "an approval-gated tool must not run" + ); + assert!(matches!(&run.messages[1], Message::Assistant(a) if a.tool_calls.len() == 2)); + assert_eq!( + tool_result_text(&run.messages, "call-lookup").as_deref(), + Some("found") + ); + assert!(tool_result_text(&run.messages, "call-delete").is_none()); + assert_eq!(run.model_calls, 1, "the loop must not call the model again"); + assert!(run.final_response.is_none()); + + let deferred = run.deferred.expect("run reports the pending approval"); + assert_eq!(deferred.approvals.len(), 1); + assert_eq!(deferred.approvals[0].id, "call-delete"); + assert!(deferred.calls.is_empty()); + assert_eq!(deferred.remaining(&DeferredToolResults::default()).len(), 1); + assert_eq!( + result.status.status, + crate::ids::ExecutionStatus::Interrupted, + "a deferred run is interrupted, not completed" + ); + assert!(recorder.events().iter().any(|event| matches!( + event, + AgentEvent::ToolDeferred { call_id, reason } + if call_id == &CallId::new("call-delete") && reason == "approval_required" + ))); +} + +// ── Resume ────────────────────────────────────────────────────────────────── + +/// Runs the mixed batch to its deferral and returns the harness, the tools, +/// and the deferred run, ready to resume. +async fn deferred_run( + recorder: &EventRecorder, +) -> ( + AgentHarness<()>, + Arc, + Arc, + crate::middleware::AgentRun, +) { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + mixed_batch(), + response(Vec::new(), "all done"), + ])), + ); + let delete = RecordingTool::approval_gated("delete", "deleted"); + let lookup = RecordingTool::plain("lookup", "found"); + harness.register_tool(delete.clone()); + harness.register_tool(lookup.clone()); + let ctx = RunContext::new(RunConfig::new("first"), ()).with_events(recorder.sink()); + let run = harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("first leg defers"); + assert!(run.deferred.is_some()); + (harness, delete, lookup, run) +} + +#[tokio::test] +async fn resume_with_approve_runs_the_tool_and_continues_to_the_model() { + let recorder = EventRecorder::new(); + let (harness, delete, lookup, first) = deferred_run(&recorder).await; + let pending: DeferredToolRequests = first.deferred.clone().unwrap(); + + let results = DeferredToolResults::new().approve("call-delete"); + assert!(pending.remaining(&results).is_empty()); + let ctx = RunContext::new(RunConfig::new("second"), ()).with_events(recorder.sink()); + let run = harness + .resume_deferred(&(), ctx, first.messages.clone(), results) + .await + .expect("resume completes the run"); + + assert_eq!(delete.calls(), vec![json!({"path": "/tmp/x"})]); + assert_eq!( + lookup.calls().len(), + 1, + "the sibling is not re-run on resume" + ); + assert_eq!( + tool_result_text(&run.messages, "call-delete").as_deref(), + Some("deleted") + ); + assert_eq!(run.text().as_deref(), Some("all done")); + assert!(run.deferred.is_none()); + assert_eq!( + run.model_calls, 1, + "resume spends exactly one new model call" + ); + assert!(recorder.events().iter().any(|event| matches!( + event, + AgentEvent::ToolApproved { call_id } if call_id == &CallId::new("call-delete") + ))); +} + +#[tokio::test] +async fn resume_with_approve_with_args_runs_the_tool_with_the_edited_arguments() { + let recorder = EventRecorder::new(); + let (harness, delete, _lookup, first) = deferred_run(&recorder).await; + + let results = + DeferredToolResults::new().approve_with_args("call-delete", json!({"path": "/tmp/safer"})); + let ctx = RunContext::new(RunConfig::new("second"), ()).with_events(recorder.sink()); + let run = harness + .resume_deferred(&(), ctx, first.messages.clone(), results) + .await + .expect("resume completes the run"); + + assert_eq!(delete.calls(), vec![json!({"path": "/tmp/safer"})]); + assert_eq!( + tool_result_text(&run.messages, "call-delete").as_deref(), + Some("deleted") + ); + assert_eq!(run.text().as_deref(), Some("all done")); +} + +#[tokio::test] +async fn resume_with_deny_answers_the_call_with_the_message_and_never_runs_it() { + let recorder = EventRecorder::new(); + let (harness, delete, _lookup, first) = deferred_run(&recorder).await; + + let results = DeferredToolResults::new().deny("call-delete", "operator refused the delete"); + let ctx = RunContext::new(RunConfig::new("second"), ()).with_events(recorder.sink()); + let run = harness + .resume_deferred(&(), ctx, first.messages.clone(), results) + .await + .expect("a denial is not a failure"); + + assert!(delete.calls().is_empty()); + let denial = run + .messages + .iter() + .find_map(|message| match message { + Message::Tool(tool) if tool.tool_call_id == "call-delete" => Some(tool.clone()), + _ => None, + }) + .expect("the denial is a tool-result row"); + assert_eq!( + denial.content, + vec![ContentBlock::Text("operator refused the delete".into())] + ); + assert_eq!(denial.artifact.as_ref().unwrap()["is_error"], true); + assert_eq!(run.text().as_deref(), Some("all done")); + assert!(!run.executed_tools.iter().any(|name| name == "delete")); + assert!(recorder.events().iter().any(|event| matches!( + event, + AgentEvent::ToolDenied { call_id, message } + if call_id == &CallId::new("call-delete") && message == "operator refused the delete" + ))); +} + +#[tokio::test] +async fn resume_refuses_an_incomplete_resolution_and_names_the_missing_ids() { + let recorder = EventRecorder::new(); + let (harness, delete, _lookup, first) = deferred_run(&recorder).await; + let pending = first.deferred.clone().unwrap(); + + let results = DeferredToolResults::new(); + assert_eq!( + pending.remaining(&results), + vec![CallId::new("call-delete")] + ); + let ctx = RunContext::new(RunConfig::new("second"), ()); + let error = harness + .resume_deferred(&(), ctx, first.messages.clone(), results) + .await + .expect_err("nothing was resolved"); + assert!( + matches!(&error, TinyAgentsError::Validation(message) if message.contains("call-delete")), + "{error}" + ); + assert!(delete.calls().is_empty()); +} + +// ── External tools ────────────────────────────────────────────────────────── + +#[tokio::test] +async fn external_tool_call_is_deferred_and_its_host_result_is_injected_on_resume() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + response( + vec![ToolCall::new( + "call-ext", + "browser_click", + json!({"x": 1, "y": 2}), + )], + "", + ), + response(Vec::new(), "clicked"), + ])), + ); + harness.register_external_tool(tinyinference_llm::tool::ToolSchema { + name: "browser_click".into(), + description: "Click at a screen coordinate (runs in the client).".into(), + parameters: json!({"type": "object", "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}}}), + format: tinyinference_llm::tool::ToolFormat::Json, + }); + + let first = harness + .invoke_default(&(), vec![Message::user("click it")]) + .await + .expect("first leg defers"); + let pending = first.deferred.clone().expect("external call is pending"); + assert!(pending.approvals.is_empty()); + assert_eq!(pending.calls.len(), 1); + assert_eq!(pending.calls[0].name, "browser_click"); + assert_eq!(pending.calls[0].arguments, json!({"x": 1, "y": 2})); + assert!(tool_result_text(&first.messages, "call-ext").is_none()); + + let results = + DeferredToolResults::new().respond("call-ext", ToolResult::success("ok: clicked (1,2)")); + let ctx = RunContext::new(RunConfig::new("second"), ()); + let run = harness + .resume_deferred(&(), ctx, first.messages.clone(), results) + .await + .expect("resume completes the run"); + assert_eq!( + tool_result_text(&run.messages, "call-ext").as_deref(), + Some("ok: clicked (1,2)") + ); + assert_eq!(run.text().as_deref(), Some("clicked")); + assert!( + run.executed_tools.is_empty(), + "the harness never ran the external tool" + ); +} + +// ── Inline handler ────────────────────────────────────────────────────────── + +/// A handler that approves everything and records what it was asked. +struct ApproveAllHandler { + asked: Mutex>, +} + +#[async_trait] +impl crate::tool::DeferredToolHandler for ApproveAllHandler { + async fn handle( + &self, + requests: &DeferredToolRequests, + ) -> crate::error::Result { + self.asked.lock().unwrap().push(requests.clone()); + Ok(requests.approve_all()) + } +} + +#[tokio::test] +async fn inline_handler_resolves_deferrals_without_surfacing_them() { + let recorder = EventRecorder::new(); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + mixed_batch(), + response(Vec::new(), "all done"), + ])), + ); + let delete = RecordingTool::approval_gated("delete", "deleted"); + harness.register_tool(delete.clone()); + harness.register_tool(RecordingTool::plain("lookup", "found")); + let handler = Arc::new(ApproveAllHandler { + asked: Mutex::new(Vec::new()), + }); + harness.with_deferred_tool_handler(handler.clone()); + + let ctx = RunContext::new(RunConfig::new("inline"), ()).with_events(recorder.sink()); + let run = harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("the handler settles the batch"); + + assert!(run.deferred.is_none(), "the caller never sees the deferral"); + assert_eq!(run.text().as_deref(), Some("all done")); + assert_eq!(delete.calls(), vec![json!({"path": "/tmp/x"})]); + assert_eq!( + tool_result_text(&run.messages, "call-delete").as_deref(), + Some("deleted") + ); + let asked = handler.asked.lock().unwrap(); + assert_eq!(asked.len(), 1); + assert_eq!(asked[0].approvals[0].id, "call-delete"); + let events = recorder.events(); + assert!( + events + .iter() + .any(|e| matches!(e, AgentEvent::ToolDeferred { .. })) + ); + assert!( + events + .iter() + .any(|e| matches!(e, AgentEvent::ToolApproved { .. })) + ); +} + +/// A handler that leaves the request unresolved. +struct SilentHandler; + +#[async_trait] +impl crate::tool::DeferredToolHandler for SilentHandler { + async fn handle( + &self, + _requests: &DeferredToolRequests, + ) -> crate::error::Result { + Ok(DeferredToolResults::new()) + } +} + +#[tokio::test] +async fn inline_handler_that_leaves_calls_unresolved_fails_the_run() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![mixed_batch()])), + ); + harness.register_tool(RecordingTool::approval_gated("delete", "deleted")); + harness.register_tool(RecordingTool::plain("lookup", "found")); + harness.with_deferred_tool_handler(Arc::new(SilentHandler)); + + let error = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect_err("an incomplete resolution is a validation failure"); + assert!( + matches!(&error, TinyAgentsError::Validation(message) if message.contains("call-delete")), + "{error}" + ); +} + +// ── Execution-time deferral (`Err(ApprovalRequired)` from the tool) ───────── + +struct SelfDeferringTool; + +#[async_trait] +impl Tool for SelfDeferringTool { + fn name(&self) -> &str { + "wire_money" + } + fn description(&self) -> &str { + "asks for approval from inside execute" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object"}) + } + async fn execute(&self, arguments: serde_json::Value) -> anyhow::Result { + Err(TinyAgentsError::ApprovalRequired { + metadata: json!({"amount": arguments["amount"]}), + } + .into()) + } +} + +#[tokio::test] +async fn tool_raising_approval_required_defers_with_its_metadata() { + let recorder = EventRecorder::new(); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![response( + vec![ToolCall::new( + "call-wire", + "wire_money", + json!({"amount": 500}), + )], + "", + )])), + ); + harness.register_tool(Arc::new(SelfDeferringTool)); + + let ctx = RunContext::new(RunConfig::new("exec-defer"), ()).with_events(recorder.sink()); + let run = harness + .invoke_in_context(&(), ctx, vec![Message::user("pay")]) + .await + .expect("a deferral is not an error"); + let pending = run.deferred.expect("pending approval"); + assert_eq!(pending.approvals[0].id, "call-wire"); + assert_eq!( + pending.metadata.get(&CallId::new("call-wire")), + Some(&json!({"amount": 500})) + ); + // The `ToolStarted` emitted before execution has exactly one terminal + // partner, the `ToolDeferred`, and no `ToolFailed`. + let events = recorder.events(); + assert!( + events + .iter() + .any(|e| matches!(e, AgentEvent::ToolStarted { .. })) + ); + assert!( + events + .iter() + .any(|e| matches!(e, AgentEvent::ToolDeferred { .. })) + ); + assert!( + !events + .iter() + .any(|e| matches!(e, AgentEvent::ToolFailed { .. })) + ); + assert_eq!(run.tool_calls, 0); +} + +// ── HumanApprovalMiddleware ───────────────────────────────────────────────── + +#[tokio::test] +async fn human_approval_middleware_defer_outcome_produces_the_deferred_exit() { + use crate::middleware::library::{ApprovalOutcome, HumanApprovalMiddleware}; + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + mixed_batch(), + response(Vec::new(), "all done"), + ])), + ); + // Neither tool declares approval in its policy; the middleware decides. + let delete = RecordingTool::plain("delete", "deleted"); + let lookup = RecordingTool::plain("lookup", "found"); + harness.register_tool(delete.clone()); + harness.register_tool(lookup.clone()); + harness.push_middleware(Arc::new( + HumanApprovalMiddleware::new(["delete"]).with_approval_outcome(Arc::new( + |call: &ToolCall| { + if call.arguments["path"] == "/tmp/x" { + ApprovalOutcome::Defer + } else { + ApprovalOutcome::Allow + } + }, + )), + )); + + let first = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("defer is not an error"); + let pending = first.deferred.clone().expect("the flagged call is pending"); + assert_eq!(pending.approvals[0].id, "call-delete"); + assert!(delete.calls().is_empty()); + assert_eq!(lookup.calls().len(), 1); + + // On resume the same middleware sees the approval and lets it through. + let run = harness + .resume_deferred( + &(), + RunContext::new(RunConfig::new("second"), ()), + first.messages.clone(), + DeferredToolResults::new().approve("call-delete"), + ) + .await + .expect("resume completes"); + assert_eq!(delete.calls().len(), 1); + assert_eq!(run.text().as_deref(), Some("all done")); +} + +#[tokio::test] +async fn human_approval_middleware_deny_outcome_answers_the_model_without_running() { + use crate::middleware::library::{ApprovalOutcome, HumanApprovalMiddleware}; + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + mixed_batch(), + response(Vec::new(), "understood"), + ])), + ); + let delete = RecordingTool::plain("delete", "deleted"); + harness.register_tool(delete.clone()); + harness.register_tool(RecordingTool::plain("lookup", "found")); + harness.push_middleware(Arc::new( + HumanApprovalMiddleware::new(["delete"]).with_approval_outcome(Arc::new( + |_call: &ToolCall| ApprovalOutcome::Deny("policy forbids deletes".into()), + )), + )); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("a denial is answered, not raised"); + assert!(delete.calls().is_empty()); + assert!(run.deferred.is_none()); + assert_eq!( + tool_result_text(&run.messages, "call-delete").as_deref(), + Some("policy forbids deletes") + ); + assert_eq!(run.text().as_deref(), Some("understood")); +} diff --git a/crates/tinyagents-harness/src/agent_loop/dialect.rs b/crates/tinyagents-harness/src/agent_loop/dialect.rs index 6d25bba0..89d5d613 100644 --- a/crates/tinyagents-harness/src/agent_loop/dialect.rs +++ b/crates/tinyagents-harness/src/agent_loop/dialect.rs @@ -210,7 +210,7 @@ pub(super) fn recover_text_calls( return; } for diagnostic in &outcome.diagnostics { - tinyagents_tracing::debug!(?diagnostic, "[agent_loop] text-dialect recovery"); + tracing::debug!(?diagnostic, "[agent_loop] text-dialect recovery"); } // Appended, not assigned: a provider can legitimately return one native // structured call *and* narrate a second one as text in the same diff --git a/crates/tinyagents-harness/src/agent_loop/entry.rs b/crates/tinyagents-harness/src/agent_loop/entry.rs index c3201c1b..eac8a7dd 100644 --- a/crates/tinyagents-harness/src/agent_loop/entry.rs +++ b/crates/tinyagents-harness/src/agent_loop/entry.rs @@ -23,9 +23,33 @@ impl AgentBaseCall ) -> BoxAgentFuture<'a> { Box::pin(async move { ctx.streaming = request.streaming; - self.harness - .run_loop(state, ctx, run, status, request.input, request.streaming) - .await + match self.harness.policy.execution { + crate::runtime::LoopExecution::Graph => match self.harness.loop_driver.clone() { + Some(driver) => { + driver + .drive( + self.harness, + state, + ctx, + run, + status, + request.input, + request.streaming, + ) + .await + } + None => Err(TinyAgentsError::Validation( + "RunPolicy::execution is LoopExecution::Graph but no LoopDriver is \ + installed; call AgentHarness::with_loop_driver first" + .to_string(), + )), + }, + crate::runtime::LoopExecution::Direct => { + self.harness + .run_loop(state, ctx, run, status, request.input, request.streaming) + .await + } + } }) } } @@ -51,7 +75,16 @@ impl TerminalRunGuard { fn complete(mut self, succeeded: bool, error: Option) -> AgentRun { if let Some(observer) = self.observer.take() { - observer(self.run.clone(), succeeded, error); + // A cheap summary (M-6), not a clone of the whole run: the + // observer only ever reads text/usage/executed-tools, and cloning + // `self.run` here duplicated the entire transcript just to throw + // it away after the observer call — `mem::take` below is the only + // place that needs to move the real run out. + observer( + crate::context::TerminalRunSummary::from_run(&self.run), + succeeded, + error, + ); } std::mem::take(&mut self.run) } @@ -61,7 +94,7 @@ impl Drop for TerminalRunGuard { fn drop(&mut self) { if let Some(observer) = self.observer.take() { observer( - self.run.clone(), + crate::context::TerminalRunSummary::from_run(&self.run), false, Some("hosted invocation cancelled by caller".to_string()), ); @@ -164,6 +197,77 @@ impl AgentHarness { self.drive(state, ctx, input, false).await } + /// Resumes a run that stopped with [`AgentRun::deferred`] set (A2). + /// + /// `messages` is the deferred run's transcript (`run.messages`, which + /// still ends with the assistant tool-call row whose deferred calls are + /// unanswered) and `results` resolves every pending call: an + /// [`crate::tool::ToolApprovalDecision`] runs or denies an approval-gated + /// call, a [`crate::tool::DeferredCallResult`] injects the host's outcome + /// for an external one. The loop answers each call — executing approved + /// ones for real, with the model's or the approver's edited arguments — + /// and then continues with the next model call exactly as if the batch + /// had never paused. + /// + /// The only state needed to resume is the transcript plus `results`, so + /// this works across a process restart: persist `run.messages` and + /// `run.deferred` (both serializable), and call this from any process. + /// Check [`crate::tool::DeferredToolRequests::remaining`] first — + /// an incomplete `results` fails with [`TinyAgentsError::Validation`] + /// naming the unresolved ids before anything runs. + /// + /// Equivalent to `invoke_in_context(state, ctx.with_deferred_results(results), messages)`, + /// preceded by a [`AgentHarness::reconcile_tool_effects`] pass scoped to + /// exclude the calls `results` is about to answer. + /// + /// A call this run's own `execution_deferral` filed (mid-execution + /// `ApprovalRequired`/`CallDeferred`) is settled in the tool-effect + /// ledger as [`crate::tool::ToolEffectStatus::Deferred`] the moment it + /// pauses — not left `started` — so [`AgentHarness::reconcile_tool_effects`] + /// (which only reconciles rows still `started`) does not treat it as a + /// crash artifact on its own. The `excluded` set passed here is + /// defense-in-depth on top of that: even if a call's row is unexpectedly + /// still `started` (the `Deferred` settle write is best-effort and only + /// logs on failure), excluding every id `results` answers guarantees this + /// call never receives a synthesized "interrupted" answer that would + /// pre-empt `results`'s real one. + /// + /// A genuinely crashed **sibling** call in the same batch — one with no + /// entry in `results` and a ledger row still `started` because the + /// process died before it could pause or settle — is not excluded, and + /// is reconciled normally (re-executed or answered "interrupted before + /// settlement" per its [`tinytools::ToolReplay`] policy) before the loop + /// resumes. + /// + /// Reconciling a genuine crash with no live `results` at all (a host + /// resuming from durable state after a real process crash, with no + /// deferral in flight) remains a host's explicit, separate call to + /// [`AgentHarness::reconcile_tool_effects`] — this method's own + /// reconcile pass only ever excludes ids `results` names, so it is a + /// strict addition, never a replacement, for that path. + pub async fn resume_deferred( + &self, + state: &State, + ctx: RunContext, + messages: Vec, + results: crate::tool::DeferredToolResults, + ) -> Result { + let mut messages = messages; + if ctx.tool_effect_ledger.is_some() { + let run_id = ctx.run_id().as_str().to_string(); + let excluded: std::collections::HashSet = results + .approvals + .keys() + .chain(results.calls.keys()) + .cloned() + .collect(); + self.reconcile_tool_effects(&ctx, &run_id, &mut messages, &excluded) + .await?; + } + self.invoke_in_context(state, ctx.with_deferred_results(results), messages) + .await + } + /// Streaming counterpart of [`AgentHarness::invoke`]. /// /// Behaves exactly like [`AgentHarness::invoke`] except each model call is @@ -336,7 +440,8 @@ impl AgentHarness { // A paused run is resumable, not finished: reporting it // `completed` is what made "paused for a human" look identical // to "the model produced an empty final answer". - let paused = terminal.run.paused.is_some(); + // A deferred run (A2) is resumable for the same reason. + let paused = terminal.run.paused.is_some() || terminal.run.deferred.is_some(); if paused { status.mark_interrupted(); } else { diff --git a/crates/tinyagents-harness/src/agent_loop/handoff_transform.rs b/crates/tinyagents-harness/src/agent_loop/handoff_transform.rs new file mode 100644 index 00000000..7f0cd76c --- /dev/null +++ b/crates/tinyagents-harness/src/agent_loop/handoff_transform.rs @@ -0,0 +1,372 @@ +//! Cross-provider handoff transform. +//! +//! A run's transcript can span more than one provider or model across its +//! lifetime: an explicit [`ModelRequest::model`][tinyinference_llm::model::ModelRequest::model] +//! override, a fallback chain, or a host-driven routing decision can each +//! hand the *next* model call a transcript whose assistant messages were +//! produced by a *different* provider. Left alone, that transcript can carry +//! content the new target cannot replay: +//! +//! * a provider-encrypted [`ContentBlock::RedactedThinking`] block, opaque +//! outside the provider that emitted it; +//! * a signed [`ContentBlock::Thinking`] block whose signature only that +//! provider can verify; +//! * tool-call ids shaped for the origin provider (for example an OpenAI +//! Responses id) that the target provider rejects outright (Anthropic caps +//! `tool_use`/`tool_result` ids at 64 characters matching +//! `^[a-zA-Z0-9_-]{1,64}$`); +//! * an image block when the target model has no vision input. +//! +//! [`prepare_for_model`] runs as a pure, read-only-in/owned-out pass over the +//! working transcript immediately before a +//! [`ModelRequest`][tinyinference_llm::model::ModelRequest] is dispatched. It +//! never touches `Turn`/`RunQueue` bookkeeping (a different concern owned +//! elsewhere in the loop) and never mutates its input: same-origin +//! transcripts — the overwhelming common case, a run that never switches +//! provider — are returned as +//! [`Cow::Borrowed`][std::borrow::Cow::Borrowed] with zero allocation. +//! +//! # What counts as "foreign" +//! +//! An [`AssistantMessage`] with an [`origin`][AssistantMessage::origin] is +//! foreign when that origin differs from the call's `target_origin` in +//! *any* of `provider`, `api`, or `model`. An assistant message with no +//! origin (replayed from a journal written before this field existed, or +//! authored directly by the host) is foreign only if it structurally +//! carries content the target cannot accept — the same-origin optimism a +//! stamped message gets is not extended to a message the harness cannot +//! actually verify came from the target. +//! +//! [`Message::User`] and [`Message::Tool`] carry no origin at all (only a +//! model produces an [`AssistantMessage`]), so they are always checked +//! structurally: an image block is downgraded whenever the target lacks +//! vision input, regardless of which turn produced the message. +//! +//! # Tool-call id remapping +//! +//! A foreign assistant message's non-conforming tool-call ids are rewritten +//! to a target-conforming shape (ASCII alphanumeric/`_`/`-`, truncated to +//! [`ModelProfile::max_tool_call_id_len`] when set, de-duplicated against +//! every id already in the transcript) through one id map built for the +//! whole call. Every [`Message::Tool`] whose `tool_call_id` answers a +//! remapped call is rewritten with the same mapping, so a tool result never +//! ends up orphaned from the call it answers. + +use std::borrow::Cow; +use std::collections::{HashMap, HashSet}; + +use regex::Regex; +use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message, MessageOrigin}; +use tinyinference_llm::model::ModelProfile; + +/// Placeholder text substituted for an image block the target model cannot +/// accept. +const IMAGE_PLACEHOLDER: &str = "[image omitted: not supported by the target model]"; + +/// Outcome of [`prepare_for_model`]: the (possibly unchanged) messages plus +/// how many messages were rewritten, so the caller can decide whether to +/// emit [`crate::events::AgentEvent::HandoffTransformApplied`]. +pub(super) struct HandoffTransformOutcome<'a> { + /// The transcript to send to the model. Borrowed when nothing changed. + pub(super) messages: Cow<'a, [Message]>, + /// Number of messages rewritten (0 when `messages` is [`Cow::Borrowed`]). + pub(super) changes: usize, +} + +/// Derives the [`MessageOrigin`] a resolved model's [`ModelProfile`] +/// represents, for use as [`prepare_for_model`]'s `target_origin`. +/// +/// [`ModelProfile`] does not itself carry an API-surface tag (a request can +/// reach the same OpenAI-family provider through Chat Completions or the +/// Responses API, and both share one profile shape), so `api` is a +/// best-effort hint derived from `provider`: Anthropic's native Messages API +/// is named exactly, every other known provider is assumed to be an +/// OpenAI-compatible Chat Completions endpoint (true for OpenAI itself and +/// every local/compatible preset), and an unknown/absent provider gets an +/// empty `api`. This under-distinguishes an OpenAI Responses-API call from a +/// Chat Completions call against the same model id — both compare equal to +/// this function's own output, so a within-run switch between the two APIs +/// is not treated as foreign. A caller with more precise knowledge (for +/// example a host that resolved a real API surface) should build a +/// [`MessageOrigin`] directly instead of calling this helper. +pub(super) fn target_origin_for(profile: &ModelProfile) -> MessageOrigin { + let provider = profile.provider.clone().unwrap_or_default(); + let api = match provider.as_str() { + "anthropic" => "messages", + "" => "", + _ => "chat_completions", + } + .to_string(); + MessageOrigin { + provider, + api, + model: profile.model.clone().unwrap_or_default(), + } +} + +/// Prepares a transcript for a call against `target`/`target_origin`, +/// rewriting only the messages a cross-provider handoff makes unsafe to +/// replay verbatim. See the module documentation for the exact rules. +/// +/// Pure and allocation-free on the common path: a transcript with no +/// foreign content (including every same-origin run, which is most runs) +/// returns [`Cow::Borrowed`] over `messages`. +pub(super) fn prepare_for_model<'a>( + messages: &'a [Message], + target: &ModelProfile, + target_origin: &MessageOrigin, +) -> HandoffTransformOutcome<'a> { + let id_pattern = target + .tool_call_id_pattern + .as_deref() + .and_then(|pattern| Regex::new(pattern).ok()); + + // Pass 1: decide which assistant tool-call ids need remapping, and seed + // the "already used" set with every id already in the transcript so a + // freshly minted id can never collide with one that was already fine. + let mut used_ids: HashSet = HashSet::new(); + let mut id_map: HashMap = HashMap::new(); + for message in messages { + if let Message::Assistant(assistant) = message { + let foreign = is_foreign(assistant, target, target_origin, id_pattern.as_ref()); + for call in &assistant.tool_calls { + if foreign && !tool_call_id_conforms(&call.id, target, id_pattern.as_ref()) { + let new_id = mint_tool_call_id(&call.id, target, &mut used_ids); + id_map.insert(call.id.clone(), new_id); + } else { + used_ids.insert(call.id.clone()); + } + } + } + } + + if id_map.is_empty() + && !messages.iter().any(|message| match message { + Message::Assistant(assistant) => { + is_foreign(assistant, target, target_origin, id_pattern.as_ref()) + } + Message::User(user) => content_needs_image_downgrade(&user.content, target), + Message::Tool(tool) => content_needs_image_downgrade(&tool.content, target), + Message::System(_) | Message::Custom(_) => false, + }) + { + return HandoffTransformOutcome { + messages: Cow::Borrowed(messages), + changes: 0, + }; + } + + let mut changes = 0usize; + let mut out = Vec::with_capacity(messages.len()); + for message in messages { + match message { + Message::Assistant(assistant) => { + if is_foreign(assistant, target, target_origin, id_pattern.as_ref()) { + out.push(Message::Assistant(transform_assistant( + assistant, target, &id_map, + ))); + changes += 1; + } else { + out.push(message.clone()); + } + } + Message::Tool(tool) => { + let remapped_id = id_map.get(&tool.tool_call_id); + let needs_image_downgrade = content_needs_image_downgrade(&tool.content, target); + if remapped_id.is_some() || needs_image_downgrade { + let mut tool = tool.clone(); + if let Some(new_id) = remapped_id { + tool.tool_call_id = new_id.clone(); + } + if needs_image_downgrade { + tool.content = downgrade_images(tool.content, target); + } + out.push(Message::Tool(tool)); + changes += 1; + } else { + out.push(message.clone()); + } + } + Message::User(user) => { + if content_needs_image_downgrade(&user.content, target) { + let mut user = user.clone(); + user.content = downgrade_images(user.content, target); + out.push(Message::User(user)); + changes += 1; + } else { + out.push(message.clone()); + } + } + other => out.push(other.clone()), + } + } + + HandoffTransformOutcome { + messages: Cow::Owned(out), + changes, + } +} + +/// Whether `assistant` was produced by a different provider/api/model than +/// `target_origin` (or, lacking a stamped origin, structurally carries +/// content the target cannot accept). +fn is_foreign( + assistant: &AssistantMessage, + target: &ModelProfile, + target_origin: &MessageOrigin, + id_pattern: Option<&Regex>, +) -> bool { + match &assistant.origin { + Some(origin) => origin != target_origin, + None => { + assistant.content.iter().any(|block| { + matches!(block, ContentBlock::RedactedThinking { .. }) + || matches!( + block, + ContentBlock::Thinking { + signature: Some(_), + .. + } + ) + || (matches!(block, ContentBlock::Image(_)) && !target.modalities.image_in) + }) || assistant + .tool_calls + .iter() + .any(|call| !tool_call_id_conforms(&call.id, target, id_pattern)) + } + } +} + +/// Whether `content` carries an image block the target cannot accept. +fn content_needs_image_downgrade(content: &[ContentBlock], target: &ModelProfile) -> bool { + !target.modalities.image_in + && content + .iter() + .any(|block| matches!(block, ContentBlock::Image(_))) +} + +/// Replaces every [`ContentBlock::Image`] with a text placeholder. +fn downgrade_images(content: Vec, target: &ModelProfile) -> Vec { + if target.modalities.image_in { + return content; + } + content + .into_iter() + .map(|block| match block { + ContentBlock::Image(_) => ContentBlock::Text(IMAGE_PLACEHOLDER.to_string()), + other => other, + }) + .collect() +} + +/// Rewrites a foreign assistant message: drops +/// [`ContentBlock::RedactedThinking`], converts a signed +/// [`ContentBlock::Thinking`] to plain text (or drops it when empty), +/// downgrades images the target cannot accept, and remaps any tool-call id +/// found in `id_map`. +fn transform_assistant( + assistant: &AssistantMessage, + target: &ModelProfile, + id_map: &HashMap, +) -> AssistantMessage { + let content = assistant + .content + .iter() + .cloned() + .filter_map(|block| match block { + ContentBlock::RedactedThinking { .. } => None, + ContentBlock::Thinking { + text, + signature: Some(_), + } => { + if text.is_empty() { + None + } else { + Some(ContentBlock::Text(text)) + } + } + ContentBlock::Image(_) if !target.modalities.image_in => { + Some(ContentBlock::Text(IMAGE_PLACEHOLDER.to_string())) + } + other => Some(other), + }) + .collect(); + let tool_calls = assistant + .tool_calls + .iter() + .cloned() + .map(|mut call| { + if let Some(new_id) = id_map.get(&call.id) { + call.id = new_id.clone(); + } + call + }) + .collect(); + AssistantMessage { + id: assistant.id.clone(), + content, + tool_calls, + usage: assistant.usage, + // The message no longer verbatim-replays what the origin provider + // produced (thinking stripped, ids remapped): it is no longer a + // faithful record of that origin, so clear it rather than leave a + // stale claim. + origin: None, + } +} + +/// Whether `id` already satisfies the target's shape/length constraints. A +/// target with neither constraint accepts every id. +fn tool_call_id_conforms(id: &str, target: &ModelProfile, id_pattern: Option<&Regex>) -> bool { + if let Some(max_len) = target.max_tool_call_id_len + && id.chars().count() > max_len + { + return false; + } + match id_pattern { + Some(pattern) => pattern.is_match(id), + None => true, + } +} + +/// Mints a target-conforming replacement for a non-conforming tool-call id: +/// disallowed characters become `_`, the result is truncated to +/// [`ModelProfile::max_tool_call_id_len`] when set, and a numeric suffix is +/// added (shrinking the base further if needed) until the candidate is +/// unique against every id `used` records. The winning candidate is +/// inserted into `used` before returning so a later collision is not +/// possible. +fn mint_tool_call_id(id: &str, target: &ModelProfile, used: &mut HashSet) -> String { + let sanitized: String = id + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || c == '_' || c == '-' { + c + } else { + '_' + } + }) + .collect(); + let sanitized = if sanitized.is_empty() { + "tc".to_string() + } else { + sanitized + }; + let max_len = target.max_tool_call_id_len.unwrap_or(usize::MAX); + let truncated: String = sanitized.chars().take(max_len).collect(); + + let mut candidate = truncated.clone(); + let mut suffix = 1u32; + while used.contains(&candidate) { + let suffix_str = format!("-{suffix}"); + let keep = max_len.saturating_sub(suffix_str.chars().count()).max(1); + let base: String = truncated.chars().take(keep).collect(); + candidate = format!("{base}{suffix_str}"); + suffix += 1; + } + used.insert(candidate.clone()); + candidate +} + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-harness/src/agent_loop/handoff_transform/test.rs b/crates/tinyagents-harness/src/agent_loop/handoff_transform/test.rs new file mode 100644 index 00000000..7296d7bc --- /dev/null +++ b/crates/tinyagents-harness/src/agent_loop/handoff_transform/test.rs @@ -0,0 +1,446 @@ +//! Unit tests for the cross-provider handoff transform. + +use std::borrow::Cow; + +use tinyinference_llm::message::{AssistantMessage, ContentBlock, ImageRef, Message, ToolMessage}; +use tinyinference_llm::model::{Modalities, ModelProfile}; +use tinyinference_llm::tool::ToolCall; + +use super::*; + +fn anthropic_target() -> (ModelProfile, MessageOrigin) { + let profile = ModelProfile { + provider: Some("anthropic".to_string()), + model: Some("claude-opus-4-6".to_string()), + modalities: Modalities { + image_in: true, + ..Modalities::default() + }, + tool_call_id_pattern: Some("^[a-zA-Z0-9_-]{1,64}$".to_string()), + max_tool_call_id_len: Some(64), + ..ModelProfile::default() + }; + let origin = MessageOrigin { + provider: "anthropic".to_string(), + api: "messages".to_string(), + model: "claude-opus-4-6".to_string(), + }; + (profile, origin) +} + +fn openai_origin() -> MessageOrigin { + MessageOrigin { + provider: "openai".to_string(), + api: "responses".to_string(), + model: "gpt-5".to_string(), + } +} + +fn assistant(content: Vec, tool_calls: Vec) -> AssistantMessage { + AssistantMessage { + id: None, + content, + tool_calls, + usage: None, + origin: None, + } +} + +// --------------------------------------------------------------------------- +// Same-origin: no allocation, nothing rewritten. +// --------------------------------------------------------------------------- + +#[test] +fn same_origin_transcript_is_untouched_and_borrowed() { + let (profile, target_origin) = anthropic_target(); + let mut same_origin = assistant( + vec![ + ContentBlock::thinking("reasoning"), + ContentBlock::Text("hi".into()), + ], + vec![ToolCall::new("toolu_ok", "search", serde_json::json!({}))], + ); + same_origin.origin = Some(target_origin.clone()); + let messages = vec![ + Message::system("sys"), + Message::user("hello"), + Message::Assistant(same_origin), + Message::Tool(ToolMessage { + tool_call_id: "toolu_ok".into(), + content: vec![ContentBlock::Text("42".into())], + trusted_verbatim: false, + artifact: None, + }), + ]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + assert_eq!(outcome.changes, 0); + assert!(matches!(outcome.messages, Cow::Borrowed(_))); + assert_eq!(outcome.messages.as_ref(), messages.as_slice()); +} + +// --------------------------------------------------------------------------- +// Thinking: redacted dropped, signed converted to text, unsigned untouched. +// --------------------------------------------------------------------------- + +#[test] +fn foreign_redacted_thinking_is_dropped_and_signed_thinking_becomes_text() { + let (profile, target_origin) = anthropic_target(); + let mut foreign = assistant( + vec![ + ContentBlock::RedactedThinking { + data: "opaque".into(), + }, + ContentBlock::Thinking { + text: "signed reasoning".into(), + signature: Some("sig-123".into()), + }, + ContentBlock::Text("visible answer".into()), + ], + vec![], + ); + foreign.origin = Some(openai_origin()); + let messages = vec![Message::Assistant(foreign)]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + assert_eq!(outcome.changes, 1); + let Message::Assistant(rewritten) = &outcome.messages[0] else { + panic!("expected an assistant message"); + }; + // Redacted thinking is gone entirely. + assert!( + !rewritten + .content + .iter() + .any(|block| matches!(block, ContentBlock::RedactedThinking { .. })) + ); + // The signed thinking became plain visible text (no signature to replay). + assert!( + rewritten + .content + .contains(&ContentBlock::Text("signed reasoning".to_string())) + ); + assert!( + rewritten + .content + .contains(&ContentBlock::Text("visible answer".to_string())) + ); + // The rewritten message no longer claims the origin provider verbatim. + assert!(rewritten.origin.is_none()); +} + +#[test] +fn foreign_unsigned_thinking_is_left_alone() { + let (profile, target_origin) = anthropic_target(); + let mut foreign = assistant( + vec![ + ContentBlock::Thinking { + text: "unsigned reasoning".into(), + signature: None, + }, + ContentBlock::Text("answer".into()), + ], + vec![ToolCall::new( + "call-needs-fix!", + "search", + serde_json::json!({}), + )], + ); + foreign.origin = Some(openai_origin()); + let messages = vec![Message::Assistant(foreign)]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + let Message::Assistant(rewritten) = &outcome.messages[0] else { + panic!("expected an assistant message"); + }; + assert!(rewritten.content.iter().any(|block| matches!( + block, + ContentBlock::Thinking { + signature: None, + .. + } + ))); +} + +#[test] +fn empty_signed_thinking_is_dropped_rather_than_becoming_an_empty_text_block() { + let (profile, target_origin) = anthropic_target(); + let mut foreign = assistant( + vec![ + ContentBlock::Thinking { + text: String::new(), + signature: Some("sig".into()), + }, + ContentBlock::Text("answer".into()), + ], + vec![], + ); + foreign.origin = Some(openai_origin()); + let messages = vec![Message::Assistant(foreign)]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + let Message::Assistant(rewritten) = &outcome.messages[0] else { + panic!("expected an assistant message"); + }; + assert_eq!(rewritten.content, vec![ContentBlock::Text("answer".into())]); +} + +// --------------------------------------------------------------------------- +// Tool-call id normalization: assistant call id and matching tool result id. +// --------------------------------------------------------------------------- + +#[test] +fn foreign_tool_call_ids_are_normalized_and_tool_results_follow() { + let (profile, target_origin) = anthropic_target(); + let long_id = "resp_call_".to_string() + &"x".repeat(80); + let mut foreign = assistant( + vec![ContentBlock::Text("checking".into())], + vec![ToolCall::new( + long_id.clone(), + "search", + serde_json::json!({"q":"x"}), + )], + ); + foreign.origin = Some(openai_origin()); + let messages = vec![ + Message::Assistant(foreign), + Message::Tool(ToolMessage { + tool_call_id: long_id.clone(), + content: vec![ContentBlock::Text("result".into())], + trusted_verbatim: false, + artifact: None, + }), + ]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + assert_eq!(outcome.changes, 2); + let Message::Assistant(rewritten) = &outcome.messages[0] else { + panic!("expected an assistant message"); + }; + let new_id = rewritten.tool_calls[0].id.clone(); + assert_ne!(new_id, long_id, "the id must be rewritten"); + assert!(new_id.chars().count() <= 64); + assert!( + new_id + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-') + ); + + let Message::Tool(tool_result) = &outcome.messages[1] else { + panic!("expected a tool message"); + }; + assert_eq!( + tool_result.tool_call_id, new_id, + "the tool result must follow the same remapping" + ); +} + +#[test] +fn conforming_tool_call_ids_are_left_untouched_even_on_a_foreign_message() { + let (profile, target_origin) = anthropic_target(); + let mut foreign = assistant( + vec![ContentBlock::Text("checking".into())], + vec![ToolCall::new("toolu_fine", "search", serde_json::json!({}))], + ); + foreign.origin = Some(openai_origin()); + let messages = vec![Message::Assistant(foreign)]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + let Message::Assistant(rewritten) = &outcome.messages[0] else { + panic!("expected an assistant message"); + }; + assert_eq!(rewritten.tool_calls[0].id, "toolu_fine"); +} + +#[test] +fn mint_tool_call_id_deduplicates_against_used_ids() { + let profile = ModelProfile { + max_tool_call_id_len: Some(6), + ..ModelProfile::default() + }; + let mut used = std::collections::HashSet::new(); + let first = mint_tool_call_id("abc!def", &profile, &mut used); + assert_eq!(first, "abc_de"); + let second = mint_tool_call_id("abc!def", &profile, &mut used); + assert_ne!(second, first, "a colliding id must be disambiguated"); + assert!(second.chars().count() <= 6); +} + +// --------------------------------------------------------------------------- +// Image downgrade: assistant content, and user/tool content (no origin). +// --------------------------------------------------------------------------- + +fn no_vision_target() -> (ModelProfile, MessageOrigin) { + let profile = ModelProfile { + provider: Some("openai".to_string()), + model: Some("gpt-5-mini".to_string()), + ..ModelProfile::default() + }; + let origin = MessageOrigin { + provider: "openai".to_string(), + api: "chat_completions".to_string(), + model: "gpt-5-mini".to_string(), + }; + (profile, origin) +} + +#[test] +fn foreign_assistant_image_is_downgraded_when_target_lacks_vision() { + let (profile, target_origin) = no_vision_target(); + let mut foreign = assistant( + vec![ + ContentBlock::Image(ImageRef { + url: "data:image/png;base64,xyz".into(), + mime_type: Some("image/png".into()), + }), + ContentBlock::Text("see above".into()), + ], + vec![], + ); + foreign.origin = Some(MessageOrigin { + provider: "anthropic".into(), + api: "messages".into(), + model: "claude-opus-4-6".into(), + }); + let messages = vec![Message::Assistant(foreign)]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + let Message::Assistant(rewritten) = &outcome.messages[0] else { + panic!("expected an assistant message"); + }; + assert!( + !rewritten + .content + .iter() + .any(|block| matches!(block, ContentBlock::Image(_))) + ); + assert!( + rewritten + .content + .iter() + .any(|block| matches!(block, ContentBlock::Text(text) if text.contains("image"))) + ); +} + +#[test] +fn user_image_is_downgraded_regardless_of_origin_because_users_carry_none() { + let (profile, target_origin) = no_vision_target(); + let messages = vec![Message::User(tinyinference_llm::message::UserMessage { + content: vec![ + ContentBlock::Image(ImageRef { + url: "https://example.com/pic.png".into(), + mime_type: None, + }), + ContentBlock::Text("what is this?".into()), + ], + })]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + assert_eq!(outcome.changes, 1); + let Message::User(rewritten) = &outcome.messages[0] else { + panic!("expected a user message"); + }; + assert!( + !rewritten + .content + .iter() + .any(|block| matches!(block, ContentBlock::Image(_))) + ); +} + +#[test] +fn image_is_untouched_when_the_target_supports_vision() { + let (profile, target_origin) = anthropic_target(); + let messages = vec![Message::User(tinyinference_llm::message::UserMessage { + content: vec![ContentBlock::Image(ImageRef { + url: "https://example.com/pic.png".into(), + mime_type: None, + })], + })]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + assert_eq!(outcome.changes, 0); + assert!(matches!(outcome.messages, Cow::Borrowed(_))); +} + +// --------------------------------------------------------------------------- +// Legacy (no-origin) messages: foreign only when structurally unacceptable. +// --------------------------------------------------------------------------- + +#[test] +fn legacy_message_with_clean_content_is_not_touched() { + let (profile, target_origin) = anthropic_target(); + let legacy = assistant(vec![ContentBlock::Text("hello".into())], vec![]); + assert!(legacy.origin.is_none()); + let messages = vec![Message::Assistant(legacy)]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + assert_eq!(outcome.changes, 0); + assert!(matches!(outcome.messages, Cow::Borrowed(_))); +} + +#[test] +fn legacy_message_with_redacted_thinking_is_treated_as_foreign() { + let (profile, target_origin) = anthropic_target(); + let legacy = assistant( + vec![ + ContentBlock::RedactedThinking { + data: "opaque".into(), + }, + ContentBlock::Text("hello".into()), + ], + vec![], + ); + let messages = vec![Message::Assistant(legacy)]; + + let outcome = prepare_for_model(&messages, &profile, &target_origin); + + assert_eq!(outcome.changes, 1); + let Message::Assistant(rewritten) = &outcome.messages[0] else { + panic!("expected an assistant message"); + }; + assert!( + !rewritten + .content + .iter() + .any(|block| matches!(block, ContentBlock::RedactedThinking { .. })) + ); +} + +// --------------------------------------------------------------------------- +// `target_origin_for` +// --------------------------------------------------------------------------- + +#[test] +fn target_origin_for_names_anthropic_messages_api() { + let profile = ModelProfile { + provider: Some("anthropic".to_string()), + model: Some("claude-opus-4-6".to_string()), + ..ModelProfile::default() + }; + let origin = target_origin_for(&profile); + assert_eq!(origin.provider, "anthropic"); + assert_eq!(origin.api, "messages"); + assert_eq!(origin.model, "claude-opus-4-6"); +} + +#[test] +fn target_origin_for_assumes_chat_completions_for_other_known_providers() { + let profile = ModelProfile { + provider: Some("ollama".to_string()), + model: Some("llama3".to_string()), + ..ModelProfile::default() + }; + let origin = target_origin_for(&profile); + assert_eq!(origin.api, "chat_completions"); +} diff --git a/crates/tinyagents-harness/src/agent_loop/mod.rs b/crates/tinyagents-harness/src/agent_loop/mod.rs index ee098fa0..3cf1e7ab 100644 --- a/crates/tinyagents-harness/src/agent_loop/mod.rs +++ b/crates/tinyagents-harness/src/agent_loop/mod.rs @@ -98,7 +98,7 @@ use std::sync::Arc; use std::time::Duration; use crate::cache::{ResponseCache, cache_key}; -use crate::context::{MiddlewareControl, RunConfig, RunContext}; +use crate::context::{LoopTarget, MiddlewareControl, RunConfig, RunContext}; use crate::error::{Result, TinyAgentsError}; use crate::events::{AgentEvent, HarnessRunStatus, LimitKind}; use crate::ids::{CallId, ComponentId, HarnessPhase}; @@ -107,7 +107,7 @@ use crate::middleware::{ ToolBaseCall, }; use crate::model_registry::{ResolvedModelBinding, model_eligible}; -use crate::runtime::{AgentHarness, InvalidArgsPolicy, UnknownToolPolicy}; +use crate::runtime::{AgentHarness, EndStrategy, InvalidArgsPolicy, UnknownToolPolicy}; use crate::structured::{StructuredExtractor, StructuredStrategy}; use futures::StreamExt; use serde_json::Value; @@ -120,12 +120,22 @@ use tinyinference_llm::tool::{ToolCall, ToolSchema}; mod dialect; mod entry; +mod handoff_transform; mod model_call; +pub mod phases; mod run_loop; -mod stream; +pub(crate) mod stream; +mod tool_changes; mod tools; pub use stream::AgentStreamItem; +pub(crate) use stream::{StreamRunner, invoke_stream_with_runner}; +#[cfg(test)] +mod deferred_test; +#[cfg(test)] +mod rich_tool_test; +#[cfg(test)] +mod run_queue_test; #[cfg(test)] mod test; diff --git a/crates/tinyagents-harness/src/agent_loop/model_call.rs b/crates/tinyagents-harness/src/agent_loop/model_call.rs index 0070d30a..5318268d 100644 --- a/crates/tinyagents-harness/src/agent_loop/model_call.rs +++ b/crates/tinyagents-harness/src/agent_loop/model_call.rs @@ -60,21 +60,21 @@ impl AgentHarness { } let resolution = host_run.host.models.resolve(&resolve); let (budget, bound) = self.model_call_budget(ctx); - let model = match budget { - Some(remaining) => tokio::select! { - biased; - _ = ctx.cancellation.cancelled() => return Err(TinyAgentsError::Cancelled), - result = tokio::time::timeout(remaining, resolution) => result.map_err(|_| TinyAgentsError::Timeout(format!("host model resolution for run `{}` exceeded its {bound}", ctx.run_id())))?, - }, - None => tokio::select! { - biased; - _ = ctx.cancellation.cancelled() => return Err(TinyAgentsError::Cancelled), - result = resolution => result, - }, - }.map_err(|error| match error { - TinyAgentsError::Cancelled | TinyAgentsError::Timeout(_) => error, - _ => { tinyagents_tracing::warn!(agent_id = %host_run.agent_id, "[host] model resolution failed"); TinyAgentsError::Model("host model resolution failed".to_string()) } - })?; + let model = ctx + .bounded(budget, resolution, || { + format!( + "host model resolution for run `{}` exceeded its {bound}", + ctx.run_id() + ) + }) + .await + .map_err(|error| match error { + TinyAgentsError::Cancelled | TinyAgentsError::Timeout(_) => error, + _ => { + tracing::warn!(agent_id = %host_run.agent_id, "[host] model resolution failed"); + TinyAgentsError::Model("host model resolution failed".to_string()) + } + })?; let name = model .profile() .and_then(|profile| profile.model.clone()) @@ -175,14 +175,14 @@ impl AgentHarness { }); if side_effecting_provider { - tinyagents_tracing::debug!( + tracing::debug!( call_id = %call_id.as_str(), provider = "claude-code", "[cache] response cache disabled for side-effecting provider" ); } else if decision.is_none() { let reason = self.cache_skip_reason(request); - tinyagents_tracing::debug!( + tracing::debug!( call_id = %call_id.as_str(), reason = reason.as_str(), "[cache] response cache not consulted for this model call" @@ -197,7 +197,7 @@ impl AgentHarness { let looked_up = match cache.get(key).await { Ok(hit) => hit, Err(error) => { - tinyagents_tracing::warn!( + tracing::warn!( call_id = %call_id.as_str(), %error, "[cache] response-cache lookup failed; treating as a miss" @@ -257,7 +257,7 @@ impl AgentHarness { } let injected = policy.protect_prompt_prefix && apply_prompt_cache_breakpoints(&mut breakpointed); - tinyagents_tracing::debug!( + tracing::debug!( call_id = %call_id.as_str(), protect_prompt_prefix = policy.protect_prompt_prefix, prompt_cache_key_injected = injected, @@ -285,7 +285,7 @@ impl AgentHarness { .as_ref() .map(|resolved| resolved.name.as_str()); if served_by.is_some_and(|name| name != primary_name) { - tinyagents_tracing::debug!( + tracing::debug!( call_id = %call_id.as_str(), primary = %primary_name, served_by = served_by.unwrap_or_default(), @@ -298,7 +298,7 @@ impl AgentHarness { // The provider call already succeeded and was paid for. // Discarding its answer because the cache is unavailable would // be strictly worse than not caching. - tinyagents_tracing::warn!( + tracing::warn!( call_id = %call_id.as_str(), %error, "[cache] response-cache write failed; returning the response uncached" @@ -360,7 +360,7 @@ impl AgentHarness { ) -> Result { let content = cached.message.content.clone(); let tool_calls = cached.tool_calls().to_vec(); - tinyagents_tracing::debug!( + tracing::debug!( call_id = %call_id.as_str(), text_len = cached.text().len(), tool_calls = tool_calls.len(), @@ -395,6 +395,7 @@ impl AgentHarness { call_id: call.id.clone(), content: serde_json::to_string(&call.arguments).unwrap_or_default(), tool_name: Some(call.name.clone()), + ..Default::default() }), }); } @@ -442,8 +443,41 @@ impl AgentHarness { ); } if saw_streamed_content { + // Same rule as the live streaming path (see the matching comment + // in `invoke_model_streaming_once`): keep the cached response's + // own `Thinking` blocks (with their signature) verbatim unless + // the synthetic replay deltas were actually transformed by + // `on_model_delta`, since a signed thinking block must be + // replayed byte-for-byte ahead of a tool call on the next turn. + let cached_reasoning: String = cached + .message + .content + .iter() + .filter_map(|block| match block { + tinyinference_llm::message::ContentBlock::Thinking { text, .. } => { + Some(text.as_str()) + } + _ => None, + }) + .collect(); + let reasoning_untransformed = cached_reasoning == streamed_reasoning; + let mut transformed_content = Vec::new(); - if !streamed_reasoning.is_empty() { + if reasoning_untransformed { + transformed_content.extend( + cached + .message + .content + .iter() + .filter(|block| { + matches!( + block, + tinyinference_llm::message::ContentBlock::Thinking { .. } + ) + }) + .cloned(), + ); + } else if !streamed_reasoning.is_empty() { transformed_content.push(tinyinference_llm::message::ContentBlock::Thinking { text: streamed_reasoning, signature: None, @@ -549,7 +583,6 @@ impl AgentHarness { // for the call to run to completion. `cancelled()` is // cancel-safe, and the pre-call `is_cancelled()` check above // still short-circuits before the request is ever issued. - let cancellation = ctx.cancellation.clone(); let fut = async { model .invoke(state, request.clone()) @@ -563,11 +596,14 @@ impl AgentHarness { bound, fut, ); - tokio::select! { - biased; - _ = cancellation.cancelled() => Err(TinyAgentsError::Cancelled), - result = budgeted => result, - } + // `with_call_budget` already applies its own deadline, so + // this only needs to race cancellation against an + // otherwise-unbounded future — `bounded`'s `None` arm, + // which never calls `timeout_message`. + ctx.bounded(None, budgeted, || { + unreachable!("with_call_budget already applies its own timeout") + }) + .await }; match attempt_result { Ok(response) => break Ok(response), @@ -575,6 +611,17 @@ impl AgentHarness { // `RunLimits::max_retries_per_call` is a hard ceiling // that a looser `RetryPolicy::max_attempts` cannot // exceed; whichever is stricter wins. + // A registered `RetryMiddleware` (or any other + // `ModelMiddleware::overrides_retry`) already retries + // the whole wrap onion around this base call. Retrying + // again here would multiply attempts + // (`mw.max_attempts × policy.retry.max_attempts × + // |fallback|` for one logical failure) and emit + // `RetryScheduled` for attempts the middleware cannot + // see, so the base call skips its own retry loop and + // defers entirely to the middleware (I-7); the + // fallback chain below is unaffected. + let retry_overridden = self.middleware.has_retry_override(); let max_attempts = self .policy .retry @@ -585,7 +632,7 @@ impl AgentHarness { // uses), applying the harness ceiling by capping a // cloned policy first so the two sites cannot drift. let capped = self.policy.retry.clone().with_max_attempts(max_attempts); - if capped.should_retry_error(attempt, &error) { + if !retry_overridden && capped.should_retry_error(attempt, &error) { // Compute the backoff from the *pre-increment* // attempt number: `attempt == 0` is the first // retry and must sleep `initial_backoff_ms` @@ -602,7 +649,7 @@ impl AgentHarness { // for a streaming call *is* the signal that // every delta seen so far for this `call_id` // must be dropped. - tinyagents_tracing::warn!( + tracing::warn!( call_id = %call_id.as_str(), discarded_deltas = deltas_emitted, attempt, @@ -635,6 +682,7 @@ impl AgentHarness { if response.resolved_model.is_none() { response.resolved_model = Some(resolved); } + split_thinking_tags(&mut response, model.profile()); return Ok(response); } Err(error) => { @@ -803,10 +851,23 @@ impl AgentHarness { match budget { Some(budget) => match tokio::time::timeout(budget, fut).await { Ok(result) => result, - Err(_) => Err(TinyAgentsError::Timeout(format!( - "{what} for run `{run_id}` exceeded its {bound} ({} ms)", - budget.as_millis() - ))), + Err(_) => { + let message = format!( + "{what} for run `{run_id}` exceeded its {bound} ({} ms)", + budget.as_millis() + ); + // Only the per-model-call ceiling is retryable: it means + // this one call wedged, not that the run is out of time. + // Every other bound this helper is used with (the run's + // remaining wall-clock budget, for model calls, tool + // calls, host resolution, tool authorization/screening, + // and host turn preparation) is terminal. + if bound == PER_CALL_BOUND_LABEL { + Err(TinyAgentsError::CallTimeout(message)) + } else { + Err(TinyAgentsError::Timeout(message)) + } + } }, None => fut.await, } @@ -860,6 +921,17 @@ impl AgentHarness { let mut transformed_tools = StreamAccumulator::new(); let mut saw_tool_delta = false; + // Some providers pad the very first streamed text chunk with + // whitespace that is a wire-format artifact, not content (see + // `ModelProfile::ignore_streamed_leading_whitespace`). Stripped once, + // on the first delta that actually carries non-whitespace text; + // deltas consisting only of leading whitespace are dropped outright + // rather than surfaced empty. + let mut strip_leading_whitespace = model + .profile() + .map(|profile| profile.ignore_streamed_leading_whitespace) + .unwrap_or(false); + // Clone the cheap token so the cancellation future does not borrow // `ctx` for the duration of the stream loop (the body still needs // `&mut ctx` for events and middleware). @@ -953,7 +1025,20 @@ impl AgentHarness { _ => None, }; - if let Some(message_delta) = message_delta { + if let Some(mut message_delta) = message_delta { + if strip_leading_whitespace && !message_delta.text.is_empty() { + let stripped = message_delta.text.trim_start(); + if stripped.is_empty() { + message_delta.text.clear(); + } else if stripped.len() == message_delta.text.len() { + // No leading whitespace to strip in this delta; the + // next delta carrying text is no longer the first. + strip_leading_whitespace = false; + } else { + message_delta.text = stripped.to_string(); + strip_leading_whitespace = false; + } + } saw_tool_delta |= message_delta.tool_call.is_some(); // Build the middleware-facing delta first (it needs owned // copies of the fields), then move `message_delta` into the @@ -1039,11 +1124,53 @@ impl AgentHarness { { // Deltas represent only text/thinking, so preserve terminal // blocks that cannot be streamed as a `ModelDelta` (JSON, - // images, and provider extensions). Provider signatures on - // thinking are intentionally discarded: a transformed block - // can no longer be replayed as the signed raw one. + // images, and provider extensions). + // + // A signed `Thinking` block must be replayed *verbatim* on + // the next model call when thinking + tool calls are both in + // play (Anthropic requires the exact signed block ahead of a + // `tool_use`); synthesizing a fresh, unsigned block here would + // make that replay fail. So the terminal provider blocks are + // kept as-is unless a delta middleware actually rewrote the + // reasoning text: compare the concatenated `Thinking` text + // that crossed `on_model_delta` against the terminal + // response's own `Thinking` text. Equal means no middleware + // touched it — keep the terminal blocks (signature intact). + // Different means the delta stream was transformed — fall + // back to a synthetic, unsigned block built from what + // actually crossed the middleware boundary, same as before. + // `RedactedThinking` carries no reasoning text at all (it is + // opaque), so it is always kept verbatim. + let terminal_reasoning: String = response + .message + .content + .iter() + .filter_map(|block| match block { + tinyinference_llm::message::ContentBlock::Thinking { text, .. } => { + Some(text.as_str()) + } + _ => None, + }) + .collect(); + let reasoning_untransformed = terminal_reasoning == streamed_reasoning; + let mut content = Vec::new(); - if !streamed_reasoning.is_empty() { + if reasoning_untransformed { + content.extend( + response + .message + .content + .iter() + .filter(|block| { + matches!( + block, + tinyinference_llm::message::ContentBlock::Thinking { .. } + ) + }) + .cloned(), + ); + streamed_reasoning.clear(); + } else if !streamed_reasoning.is_empty() { content.push(tinyinference_llm::message::ContentBlock::Thinking { text: std::mem::take(&mut streamed_reasoning), signature: None, @@ -1093,9 +1220,79 @@ impl AgentHarness { accumulator.push(&item); } - Ok(accumulator.finish()?) + let mut response = accumulator.finish()?; + split_thinking_tags(&mut response, model.profile()); + Ok(response) } } + +/// Splits a model's inline `...`-tagged reasoning span out of a +/// `Text` content block into a dedicated [`ContentBlock::Thinking`] block, +/// for a model whose [`ModelProfile::thinking_tags`] declares the tag pair it +/// emits inline instead of on a distinct reasoning channel. +/// +/// A no-op when the profile declares no tag pair, or the response carries no +/// text block containing both tags. Only the first tagged span in each text +/// block is extracted — every provider that uses this convention emits at +/// most one reasoning span ahead of the visible answer — and any text before +/// or after the span is preserved as ordinary `Text` blocks in the same +/// position. +/// +/// [`ContentBlock::Thinking`]: tinyinference_llm::message::ContentBlock::Thinking +/// [`ModelProfile::thinking_tags`]: tinyinference_llm::model::ModelProfile::thinking_tags +pub(super) fn split_thinking_tags( + response: &mut tinyinference_llm::model::ModelResponse, + profile: Option<&tinyinference_llm::model::ModelProfile>, +) { + let Some((open, close)) = profile.and_then(|p| p.thinking_tags.as_ref()) else { + return; + }; + if open.is_empty() || close.is_empty() { + return; + } + + let mut rebuilt = Vec::with_capacity(response.message.content.len()); + for block in response.message.content.drain(..) { + match block { + tinyinference_llm::message::ContentBlock::Text(text) => { + match split_one(&text, open, close) { + Some((before, thinking, after)) => { + if !before.is_empty() { + rebuilt.push(tinyinference_llm::message::ContentBlock::Text(before)); + } + if !thinking.is_empty() { + rebuilt.push(tinyinference_llm::message::ContentBlock::Thinking { + text: thinking, + signature: None, + }); + } + if !after.is_empty() { + rebuilt.push(tinyinference_llm::message::ContentBlock::Text(after)); + } + } + None => rebuilt.push(tinyinference_llm::message::ContentBlock::Text(text)), + } + } + other => rebuilt.push(other), + } + } + response.message.content = rebuilt; +} + +/// Splits `text` on the first `open`/`close` tag pair, returning +/// `(before, inside, after)` with the tags themselves removed and each +/// segment trimmed of the whitespace/newlines the tags typically pad. `None` +/// when the text does not contain a complete `open`...`close` span. +fn split_one(text: &str, open: &str, close: &str) -> Option<(String, String, String)> { + let open_idx = text.find(open)?; + let after_open = open_idx + open.len(); + let close_rel = text[after_open..].find(close)?; + let close_idx = after_open + close_rel; + let before = text[..open_idx].trim().to_string(); + let inside = text[after_open..close_idx].trim().to_string(); + let after = text[close_idx + close.len()..].trim().to_string(); + Some((before, inside, after)) +} /// The innermost model call wrapped by the model-wrap onion. /// /// Implements [`ModelBaseCall`] over the harness's cache + retry + fallback core @@ -1178,7 +1375,7 @@ impl ModelCallBase<'_, State, Ctx> { if binding.resolved.source == ModelResolutionSource::RequestOverride && binding.resolved.name == requested => { - tinyagents_tracing::debug!( + tracing::debug!( call_id = %self.call_id.as_str(), from = %self.resolved.name, to = %binding.resolved.name, @@ -1187,7 +1384,7 @@ impl ModelCallBase<'_, State, Ctx> { Ok(binding) } _ => { - tinyagents_tracing::warn!( + tracing::warn!( call_id = %self.call_id.as_str(), requested = %requested, resolved = %self.resolved.name, @@ -1246,12 +1443,13 @@ impl ToolBaseCall for ToolCall settings.resolve(self.dispatch.tool().timeout_policy(&call.arguments)) }); let timeout_result = super::tools::timeout_result(&call, timeout); - let future = async { - self.dispatch - .execute(state, call.arguments, self.options, ctx) - .await - .map_err(super::tools::map_tool_dispatch_error) - }; + let future = super::tools::execute_tool_recovering_model_retry(self.dispatch.execute( + state, + CallId::new(call.id), + call.arguments, + self.options, + ctx, + )); match timeout.and_then(|resolved| resolved.deadline) { Some(deadline) => match tokio::time::timeout(deadline, future).await { Ok(result) => result, diff --git a/crates/tinyagents-harness/src/agent_loop/phases.rs b/crates/tinyagents-harness/src/agent_loop/phases.rs new file mode 100644 index 00000000..0a25f7df --- /dev/null +++ b/crates/tinyagents-harness/src/agent_loop/phases.rs @@ -0,0 +1,207 @@ +//! Typed phase contracts and the [`LoopDriver`] seam for driving the agent +//! loop as discrete steps instead of the monolithic [`super::run_loop`] body. +//! +//! # Why this module exists +//! +//! `tinyagents-graph` depends on `tinyagents-harness` (never the reverse), so +//! a compiled-graph rendition of the agent loop (`tinyagents_graph::agent_loop`, +//! A5 in `docs/runtime-comparison/feature-gaps.md`) must live in the graph +//! crate. This module is the harness-side half of that boundary: it defines +//! the data that crosses a node edge in the graph rendition (a turn's plan, +//! a model call's outcome, a tool batch's outcome, and how a turn settled), +//! and the [`LoopDriver`] trait that lets an [`AgentHarness`] delegate +//! `invoke`/`invoke_with_status` to an alternate engine — the graph crate's +//! `GraphLoopDriver` — instead of [`super::run_loop`]. +//! +//! # Scope +//! +//! [`super::run_loop`]'s body (`run_loop_body` in `run_loop.rs`) is a single +//! ~1,100-line function that interleaves request building, host-model +//! routing, cross-provider handoff transforms, truncated-empty-response +//! recovery, every structured-output strategy, host budget admission, and +//! the A6 end-strategy resolution — all behaviorally load-bearing and +//! covered by the harness's ~1,174-test baseline. Splitting that function +//! into four clean phase functions **in place** (rewriting `run_loop_body` +//! itself to call them) was judged too risky to attempt as part of this +//! change: it would touch the single highest-blast-radius function in the +//! crate with no incremental way to verify each extracted phase preserves +//! every one of those behaviors. +//! +//! Instead, this module defines the phase *contracts* — the types a graph +//! node reads and writes — and the [`LoopDriver`] hook that lets +//! `tinyagents-graph` supply a complete alternate implementation of those +//! phases (`tinyagents_graph::agent_loop`). That implementation intentionally +//! covers a **subset** of `run_loop_body`'s behavior (documented on +//! `tinyagents_graph::agent_loop::compile_loop`): the common tool-calling / +//! structured-output / limit / interrupt / steering paths, without host-model +//! routing, cross-provider handoff transforms, the deferred-tool discovery +//! bridge, or truncated-empty-response recovery. [`RunPolicy::execution`] +//! defaults to [`LoopExecution::Direct`], so every existing caller keeps +//! running the full-fidelity `run_loop_body` unless it explicitly opts in. +//! +//! [`AgentHarness`]: crate::runtime::AgentHarness +//! [`RunPolicy::execution`]: crate::runtime::RunPolicy::execution +//! [`LoopExecution::Direct`]: crate::runtime::LoopExecution::Direct + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; + +use crate::context::RunContext; +use crate::error::Result; +use crate::events::HarnessRunStatus; +use crate::middleware::AgentRun; +use crate::runtime::AgentHarness; +use tinyinference_llm::message::Message; +use tinyinference_llm::model::{ModelRequest, ModelResponse}; +use tinyinference_llm::tool::ToolCall; + +/// The structured-output plan resolved for one turn. +/// +/// Mirrors the `(StructuredStrategy, name, schema)` tuple `run_loop_body` +/// computes internally, using a plain string tag for the strategy instead of +/// the private [`crate::structured::StructuredStrategy`] enum so this type +/// can be `Serialize`/`Deserialize` and cross a graph node boundary (and, +/// eventually, a checkpoint) without exposing that internal type. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct StructuredPlan { + /// `"provider_schema"` (native `response_format`) or `"tool_call"` (a + /// synthetic tool the model is asked to call with the schema's shape). + pub strategy: String, + /// The schema's name, also used as the synthetic tool's name under the + /// `"tool_call"` strategy. + pub schema_name: String, + /// The JSON Schema describing the desired output shape. + pub schema: serde_json::Value, +} + +/// The strategy tag for [`StructuredPlan::strategy`] under provider-native +/// schema mode. +pub const STRATEGY_PROVIDER_SCHEMA: &str = "provider_schema"; +/// The strategy tag for [`StructuredPlan::strategy`] under the tool-call +/// fallback. +pub const STRATEGY_TOOL_CALL: &str = "tool_call"; + +/// What one turn intends to send: the built [`ModelRequest`] plus the +/// resolved structured-output plan (if any). +/// +/// Produced by a `plan_turn`-shaped step (`tinyagents_graph::agent_loop`'s +/// `plan` node) and consumed by a `call_model`-shaped step. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TurnPlan { + /// The request ready to dispatch to the resolved model. + pub request: ModelRequest, + /// The structured-output plan for this turn, when the run requested one. + pub structured: Option, +} + +/// The outcome of dispatching one model call. +/// +/// Produced by a `call_model`-shaped step (`tinyagents_graph::agent_loop`'s +/// `model` node) and consumed by the routing decision that picks `tools` or +/// `settle` next. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ModelOutcome { + /// The harness-assigned correlation id for this call. + pub call_id: String, + /// The provider's response. + pub response: ModelResponse, +} + +/// The outcome of executing one turn's batch of tool calls. +/// +/// Produced by an `execute_tool_batch`-shaped step +/// (`tinyagents_graph::agent_loop`'s `tools` node). +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ToolBatchOutcome { + /// The tool-result messages to append to the transcript, in call order. + pub results: Vec, + /// Names of calls that reached a tool executor, in execution order (see + /// [`AgentRun::executed_tools`]). + pub executed_tools: Vec, +} + +/// How a turn settled: whether the run is finished, and any structured +/// output extracted. +/// +/// Produced by a `settle_turn`-shaped step (`tinyagents_graph::agent_loop`'s +/// `settle` node). +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct Settlement { + /// Whether the run should end after this turn. + pub finished: bool, + /// The extracted structured output, when one was requested and + /// successfully extracted/validated. + pub structured: Option, + /// Which schema variant matched, under + /// [`crate::structured::StructuredStrategy::ToolCallUnion`]. `None` for + /// every other strategy. + pub structured_variant: Option, +} + +/// A seam that lets an [`AgentHarness`] delegate its loop execution to an +/// alternate engine instead of the built-in [`super::run_loop`]. +/// +/// `tinyagents-graph` is the only intended implementor +/// (`GraphLoopDriver`, see `tinyagents_graph::agent_loop`): harness cannot +/// depend on the graph crate (dependency direction is graph -> harness), so +/// this trait — plus [`AgentHarness::with_loop_driver`] — is the hook the +/// graph crate uses to plug itself in without an inverted dependency. +/// +/// A driver's [`Self::drive`] has the exact same contract as +/// [`super::run_loop`] (which it replaces at the call site in +/// `agent_loop::entry::drive_collecting`): it owns the whole run from +/// `RunStarted` through `before_agent`/`after_agent` middleware and the +/// terminal `RunCompleted`/`RunFailed`/pause event, writing every message it +/// produces onto `run.messages` before returning (so the transcript is +/// preserved on every exit path, including an error, exactly as the direct +/// loop preserves it — see `agent_loop`'s module docs on "Exits"). +#[async_trait] +pub trait LoopDriver: Send + Sync { + /// Drives one run to completion (or a deliberate pause/error), mirroring + /// [`super::run_loop`]'s contract. + #[allow(clippy::too_many_arguments)] + async fn drive( + &self, + harness: &AgentHarness, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + input: Vec, + streaming: bool, + ) -> Result<()>; +} + +/// Executes one turn's batch of tool calls and reports what changed. +/// +/// This is a thin, behavior-preserving wrapper over the same +/// [`AgentHarness::execute_tools`][super::AgentHarness::execute_tools] +/// serial-admission / serial-or-concurrent-execution / ordered-fold pipeline +/// the direct loop uses (see `agent_loop::tools`) — unlike [`TurnPlan`] and +/// friends, this phase's *implementation*, not just its data contract, is +/// reused as-is, so a graph-driven tool batch preserves the exact ordering, +/// concurrency-eligibility, budget/limit, and middleware semantics the direct +/// loop guarantees. `messages` and `run` are mutated in place, exactly as +/// `execute_tools` does; the returned [`ToolBatchOutcome`] additionally +/// reports just the slice each produced, for a caller (a graph node) that +/// wants the batch's own delta rather than diffing the whole transcript +/// itself. +pub async fn execute_tool_batch( + harness: &AgentHarness, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + messages: &mut Vec, + tool_calls: Vec, +) -> Result { + let messages_before = messages.len(); + let executed_before = run.executed_tools.len(); + harness + .execute_tools(state, ctx, run, status, messages, tool_calls) + .await?; + Ok(ToolBatchOutcome { + results: messages[messages_before..].to_vec(), + executed_tools: run.executed_tools[executed_before..].to_vec(), + }) +} diff --git a/crates/tinyagents-harness/src/agent_loop/rich_tool_test.rs b/crates/tinyagents-harness/src/agent_loop/rich_tool_test.rs new file mode 100644 index 00000000..9f28147a --- /dev/null +++ b/crates/tinyagents-harness/src/agent_loop/rich_tool_test.rs @@ -0,0 +1,509 @@ +//! Tests for B1/B2 consumption in the agent loop: what a tool sees on its +//! [`ToolExecutionContext`] during a real invocation (call id, store, typed +//! state view, `custom` events) and what the loop does with a rich +//! [`tinytools::ToolResult`] (`follow_up` becomes a user message after the +//! batch's tool rows; `metadata` reaches events and the run but never the +//! transcript). + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use serde_json::json; + +use crate::context::{RunConfig, RunContext}; +use crate::events::AgentEvent; +use crate::ids::CallId; +use crate::runtime::AgentHarness; +use crate::store::namespaced::{InMemoryNamespacedStore, Namespace, NamespacedStore}; +use crate::testkit::{EventRecorder, ScriptedModel}; +use crate::tool::ToolExecutionContext; +use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; +use tinyinference_llm::model::ModelResponse; +use tinyinference_llm::tool::ToolCall; +use tinyinference_llm::usage::Usage; +use tinytools::{FileData, ImageData, Tool, ToolCallOptions, ToolContent, ToolResult}; + +// ── Helpers ───────────────────────────────────────────────────────────────── + +#[derive(Debug, PartialEq)] +struct AppState { + tenant: &'static str, +} + +/// What a tool observed on its harness context during one invocation. +#[derive(Clone, Debug, Default)] +struct Observed { + call_id: Option, + had_store: bool, + tenant: Option<&'static str>, +} + +/// A tool that downcasts its run context to the harness type, records what it +/// saw, writes to the store, emits a custom event, and returns `result`. +struct ContextTool { + name: &'static str, + result: ToolResult, + observed: Arc>>, +} + +impl ContextTool { + fn new(name: &'static str, result: ToolResult) -> (Arc, Arc>>) { + let observed = Arc::new(Mutex::new(Vec::new())); + let tool = Arc::new(Self { + name, + result, + observed: Arc::clone(&observed), + }); + (tool, observed) + } +} + +#[async_trait] +impl Tool for ContextTool { + fn name(&self) -> &str { + self.name + } + fn description(&self) -> &str { + "context-aware tool" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object"}) + } + async fn execute(&self, _arguments: serde_json::Value) -> anyhow::Result { + unreachable!("the harness always dispatches through execute_with_context") + } + async fn execute_with_context( + &self, + _arguments: serde_json::Value, + _options: ToolCallOptions, + context: Option<&dyn tinytools::ToolRunContext>, + ) -> anyhow::Result { + let harness = context + .and_then(tinytools::ToolRunContext::host_extension) + .and_then(|any| any.downcast_ref::()); + let mut seen = Observed::default(); + if let Some(harness) = harness { + seen.call_id = Some(harness.call_id.clone()); + seen.had_store = harness.store.is_some(); + seen.tenant = harness.state::().map(|s| s.tenant); + if let Some(store) = &harness.store { + store + .put( + &Namespace::from("calls"), + harness.call_id.as_str(), + json!({"tool": self.name}), + ) + .await?; + } + harness.custom(json!({"stage": "working"})); + } + self.observed.lock().unwrap().push(seen); + Ok(self.result.clone()) + } +} + +fn response(tool_calls: Vec, text: &str) -> ModelResponse { + let content = if text.is_empty() { + Vec::new() + } else { + vec![ContentBlock::Text(text.to_string())] + }; + ModelResponse { + message: AssistantMessage { + id: None, + content, + tool_calls, + usage: Some(Usage::new(1, 1)), + origin: None, + }, + usage: Some(Usage::new(1, 1)), + finish_reason: Some("stop".to_string()), + raw: None, + resolved_model: None, + continue_turn: None, + served_from_cache: false, + correlation: None, + resolved_route: None, + } +} + +fn tool_turn(calls: &[(&str, &str)]) -> ModelResponse { + response( + calls + .iter() + .map(|(id, name)| ToolCall::new(*id, *name, json!({}))) + .collect(), + "", + ) +} + +fn final_turn(text: &str) -> ModelResponse { + response(Vec::new(), text) +} + +/// A compact role rendering of a transcript for ordering assertions. +fn shape(messages: &[Message]) -> Vec { + messages + .iter() + .map(|message| match message { + Message::System(_) => "system".to_string(), + Message::User(_) => format!("user:{}", message.text()), + Message::Assistant(a) if !a.tool_calls.is_empty() => { + format!("assistant:tools[{}]", a.tool_calls.len()) + } + Message::Assistant(_) => "assistant".to_string(), + Message::Tool(t) => format!("tool:{}", t.tool_call_id), + Message::Custom(_) => "custom".to_string(), + }) + .collect() +} + +struct Fixture { + harness: AgentHarness<()>, + model: Arc, + recorder: EventRecorder, +} + +fn fixture(responses: Vec) -> Fixture { + let model = Arc::new(ScriptedModel::new(responses)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::clone(&model) as _); + Fixture { + harness, + model, + recorder: EventRecorder::new(), + } +} + +impl Fixture { + fn ctx(&self, run_id: &str) -> RunContext<()> { + RunContext::new(RunConfig::new(run_id), ()).with_events(self.recorder.sink()) + } +} + +// ── Context parity (B1) ───────────────────────────────────────────────────── + +#[tokio::test] +async fn a_tool_sees_the_call_id_the_model_used() { + let mut fx = fixture(vec![tool_turn(&[("call-xyz", "ctx")]), final_turn("done")]); + let (tool, observed) = ContextTool::new("ctx", ToolResult::success("ok")); + fx.harness.register_tool(tool); + + let run = fx + .harness + .invoke_in_context(&(), fx.ctx("call-id"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let seen = observed.lock().unwrap(); + assert_eq!(seen.len(), 1); + assert_eq!(seen[0].call_id, Some(CallId::new("call-xyz"))); + // The transcript answers the same id. + assert!( + run.messages + .iter() + .any(|message| matches!(message, Message::Tool(t) if t.tool_call_id == "call-xyz")) + ); +} + +#[tokio::test] +async fn concurrent_calls_each_see_their_own_call_id() { + let mut fx = fixture(vec![ + tool_turn(&[("call-a", "ctx"), ("call-b", "ctx")]), + final_turn("done"), + ]); + let (tool, observed) = ContextTool::new("ctx", ToolResult::success("ok")); + fx.harness.register_tool(tool); + + fx.harness + .invoke_in_context(&(), fx.ctx("concurrent"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let mut ids: Vec<_> = observed + .lock() + .unwrap() + .iter() + .map(|seen| seen.call_id.clone().unwrap().as_str().to_string()) + .collect(); + ids.sort(); + assert_eq!(ids, ["call-a", "call-b"]); +} + +#[tokio::test] +async fn store_is_none_without_one_and_the_run_store_with_one() { + let mut fx = fixture(vec![tool_turn(&[("c1", "ctx")]), final_turn("done")]); + let (tool, observed) = ContextTool::new("ctx", ToolResult::success("ok")); + fx.harness.register_tool(tool); + fx.harness + .invoke_in_context(&(), fx.ctx("no-store"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + assert!(!observed.lock().unwrap()[0].had_store); + + let mut fx = fixture(vec![tool_turn(&[("c2", "ctx")]), final_turn("done")]); + let (tool, observed) = ContextTool::new("ctx", ToolResult::success("ok")); + fx.harness.register_tool(tool); + let store = Arc::new(InMemoryNamespacedStore::new()); + let ctx = fx + .ctx("with-store") + .with_namespaced_store(Arc::clone(&store) as Arc); + fx.harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + assert!(observed.lock().unwrap()[0].had_store); + // The tool's write landed on the host's store, keyed by the call id. + let item = store + .get(&Namespace::from("calls"), "c2") + .await + .unwrap() + .expect("tool wrote through the run store"); + assert_eq!(item.value, json!({"tool": "ctx"})); +} + +#[tokio::test] +async fn state_view_is_typed_when_attached_and_absent_otherwise() { + let mut fx = fixture(vec![tool_turn(&[("c1", "ctx")]), final_turn("done")]); + let (tool, observed) = ContextTool::new("ctx", ToolResult::success("ok")); + fx.harness.register_tool(tool); + fx.harness + .invoke_in_context(&(), fx.ctx("no-state"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + assert_eq!(observed.lock().unwrap()[0].tenant, None); + + let mut fx = fixture(vec![tool_turn(&[("c2", "ctx")]), final_turn("done")]); + let (tool, observed) = ContextTool::new("ctx", ToolResult::success("ok")); + fx.harness.register_tool(tool); + let ctx = fx + .ctx("with-state") + .with_state_view(Arc::new(AppState { tenant: "acme" })); + fx.harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + assert_eq!(observed.lock().unwrap()[0].tenant, Some("acme")); +} + +#[tokio::test] +async fn custom_events_land_on_the_run_stream_inside_the_call() { + let mut fx = fixture(vec![tool_turn(&[("c1", "ctx")]), final_turn("done")]); + let (tool, _) = ContextTool::new("ctx", ToolResult::success("ok")); + fx.harness.register_tool(tool); + fx.harness + .invoke_in_context(&(), fx.ctx("custom"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let events = fx.recorder.events(); + let position = |pred: &dyn Fn(&AgentEvent) -> bool| events.iter().position(pred).unwrap(); + let started = position(&|e| matches!(e, AgentEvent::ToolStarted { .. })); + let custom = position(&|e| { + matches!( + e, + AgentEvent::Custom { call_id: Some(id), payload } + if id.as_str() == "c1" && payload == &json!({"stage": "working"}) + ) + }); + let completed = position(&|e| matches!(e, AgentEvent::ToolCompleted { .. })); + assert!( + started < custom && custom < completed, + "custom event is ordered inside its call: {:?}", + fx.recorder.kinds() + ); +} + +// ── Rich returns (B2) ─────────────────────────────────────────────────────── + +#[tokio::test] +async fn follow_up_becomes_a_user_message_after_the_tool_result() { + let mut fx = fixture(vec![tool_turn(&[("c1", "shot")]), final_turn("done")]); + let (tool, _) = ContextTool::new( + "shot", + ToolResult::success("clicked").with_follow_up([ToolContent::Text { + text: "Here is the page after the click.".into(), + }]), + ); + fx.harness.register_tool(tool); + + let run = fx + .harness + .invoke_in_context(&(), fx.ctx("follow-up"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!( + shape(&run.messages), + [ + "user:go", + "assistant:tools[1]", + "tool:c1", + "user:Here is the page after the click.", + "assistant", + ] + ); + // The next model call saw it. + let requests = fx.model.requests(); + assert_eq!( + shape(&requests[1].messages)[3], + "user:Here is the page after the click." + ); + // The tool-result row itself does not carry the follow-up. + let Message::Tool(row) = &run.messages[2] else { + panic!("tool row"); + }; + assert_eq!(row.content, vec![ContentBlock::Text("clicked".into())]); +} + +#[tokio::test] +async fn follow_ups_in_a_batch_come_after_every_tool_row_in_call_order() { + let mut fx = fixture(vec![ + tool_turn(&[("c1", "first"), ("c2", "second")]), + final_turn("done"), + ]); + let (first, _) = ContextTool::new( + "first", + ToolResult::success("1").with_follow_up([ToolContent::Text { + text: "after first".into(), + }]), + ); + let (second, _) = ContextTool::new( + "second", + ToolResult::success("2").with_follow_up([ToolContent::Text { + text: "after second".into(), + }]), + ); + fx.harness.register_tool(first); + fx.harness.register_tool(second); + + let run = fx + .harness + .invoke_in_context(&(), fx.ctx("batch"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + // Both tool rows stay adjacent to the assistant row (a provider + // requirement); the follow-ups trail the batch in call order. + assert_eq!( + shape(&run.messages), + [ + "user:go", + "assistant:tools[2]", + "tool:c1", + "tool:c2", + "user:after first", + "user:after second", + "assistant", + ] + ); +} + +#[tokio::test] +async fn follow_up_image_becomes_an_image_block_and_a_file_a_placeholder() { + let mut fx = fixture(vec![tool_turn(&[("c1", "shot")]), final_turn("done")]); + let (tool, _) = ContextTool::new( + "shot", + ToolResult::success("ok").with_follow_up([ + ToolContent::Image { + media_type: "image/png".into(), + data: ImageData::Url("https://example.test/shot.png".into()), + }, + ToolContent::Image { + media_type: "image/jpeg".into(), + data: ImageData::Base64("AA==".into()), + }, + ToolContent::File { + name: "report.pdf".into(), + media_type: "application/pdf".into(), + data: FileData::Path("/tmp/report.pdf".into()), + }, + ToolContent::Json { + data: json!({"k": 1}), + }, + ]), + ); + fx.harness.register_tool(tool); + + let run = fx + .harness + .invoke_in_context(&(), fx.ctx("blocks"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let Message::User(follow_up) = &run.messages[3] else { + panic!( + "expected the follow-up user message, got {:?}", + run.messages[3] + ); + }; + assert_eq!( + follow_up.content, + vec![ + ContentBlock::Image(tinyinference_llm::message::ImageRef { + url: "https://example.test/shot.png".into(), + mime_type: Some("image/png".into()), + }), + ContentBlock::Image(tinyinference_llm::message::ImageRef { + url: "data:image/jpeg;base64,AA==".into(), + mime_type: Some("image/jpeg".into()), + }), + ContentBlock::Text("[file report.pdf (application/pdf)]".into()), + ContentBlock::Json(json!({"k": 1})), + ] + ); +} + +#[tokio::test] +async fn metadata_reaches_the_event_and_the_run_but_never_the_transcript() { + const MARKER: &str = "host-only-marker-7f3a"; + let mut fx = fixture(vec![tool_turn(&[("c1", "meta")]), final_turn("done")]); + let (tool, _) = ContextTool::new( + "meta", + ToolResult::success("visible").with_metadata(json!({"marker": MARKER})), + ); + fx.harness.register_tool(tool); + + let run = fx + .harness + .invoke_in_context(&(), fx.ctx("metadata"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + // On the event. + let completed = fx + .recorder + .events() + .into_iter() + .find_map(|event| match event { + AgentEvent::ToolCompleted { + call_id, metadata, .. + } => Some((call_id, metadata)), + _ => None, + }) + .expect("ToolCompleted emitted"); + assert_eq!(completed.0, CallId::new("c1")); + assert_eq!(completed.1, Some(json!({"marker": MARKER}))); + + // On the run. + assert_eq!(run.executed_tools, ["meta"]); + assert_eq!(run.tool_metadata.len(), 1); + assert_eq!(run.tool_metadata[0].call_id, CallId::new("c1")); + assert_eq!(run.tool_metadata[0].tool_name, "meta"); + assert_eq!(run.tool_metadata[0].metadata, json!({"marker": MARKER})); + + // Nowhere in what the model was sent, nor in the run transcript — not + // even in the tool row's host-side artifact. + for request in fx.model.requests() { + let wire = serde_json::to_string(&request.messages).unwrap(); + assert!( + !wire.contains(MARKER), + "metadata leaked to the model: {wire}" + ); + } + let transcript = serde_json::to_string(&run.messages).unwrap(); + assert!( + !transcript.contains(MARKER), + "metadata leaked into the transcript: {transcript}" + ); +} diff --git a/crates/tinyagents-harness/src/agent_loop/run_loop.rs b/crates/tinyagents-harness/src/agent_loop/run_loop.rs index 2976ed2d..330baa13 100644 --- a/crates/tinyagents-harness/src/agent_loop/run_loop.rs +++ b/crates/tinyagents-harness/src/agent_loop/run_loop.rs @@ -5,7 +5,9 @@ //! Split out of `agent_loop/mod.rs`; see that module's doc comment for //! the full loop lifecycle, limits, and backoff design. +use super::handoff_transform; use super::model_call::ModelCallBase; +use super::tool_changes; use super::*; impl AgentHarness { @@ -21,6 +23,12 @@ impl AgentHarness { input: Vec, streaming: bool, ) -> Result<()> { + // The tracker's wall-clock start is stamped when the context is + // constructed (`RunContext::new`), not necessarily when the run + // actually begins doing work — a context built ahead of time and + // queued would otherwise burn down its deadline before the first + // model call. Restart it here, at the true top of the run (M-8). + ctx.limits.restart(); let mut messages = input; // The body borrows the working transcript rather than owning it so the // transcript survives **every** exit path, not just the successful one. @@ -31,11 +39,18 @@ impl AgentHarness { .run_loop_body(state, ctx, run, status, &mut messages, streaming) .await; run.messages = std::mem::take(&mut messages); + // A4: the `Collect` lane is delivered on the run, never on the + // transcript, and on every exit path — a host that pushed + // observations during a run that then failed still gets them back. + if let Some(queue) = ctx.run_queue.clone() { + run.collected + .extend(queue.drain(crate::run_queue::QueueLane::Collect).await); + } let exit = match outcome { Ok(exit) => exit, Err(error) => { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::agent_loop", run_id = %ctx.run_id(), messages = run.messages.len(), @@ -51,7 +66,7 @@ impl AgentHarness { match exit { LoopExit::Finished | LoopExit::LimitStop(_) => { if let LoopExit::LimitStop(kind) = &exit { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::agent_loop", run_id = %ctx.run_id(), limit_kind = ?kind, @@ -77,7 +92,7 @@ impl AgentHarness { }), }); status.set_last_event(record.id); - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::agent_loop", run_id = %ctx.run_id(), checkpoint = pause.paused_at_checkpoint, @@ -85,6 +100,31 @@ impl AgentHarness { ); run.paused = Some(pause); } + LoopExit::Deferred(requests) => { + // Like a pause, a deferral is not a completion: the run is + // waiting on a human decision or host-side execution for the + // calls listed in `requests`. The transcript already carries + // the assistant's tool-call row and every non-deferred + // sibling's result, so persisting `run.messages` + + // `run.deferred` is all a host needs to resume later. + let record = ctx.emit(AgentEvent::ControlApplied { + control: "deferred".to_string(), + detail: format!( + "{} approval(s), {} external call(s) pending", + requests.approvals.len(), + requests.calls.len() + ), + }); + status.set_last_event(record.id); + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + approvals = requests.approvals.len(), + calls = requests.calls.len(), + "[agent_loop] run deferred on pending tool calls" + ); + run.deferred = Some(requests); + } } Ok(()) @@ -128,7 +168,7 @@ impl AgentHarness { ); let effective_tool_calls = resolve_call_cap(ctx.config.max_tool_calls, self.policy.limits.max_tool_calls); - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::agent_loop", run_id = %ctx.run_id(), config_model_calls = ?ctx.config.max_model_calls, @@ -156,13 +196,15 @@ impl AgentHarness { // `tool_call` bridge, whose two schemas are appended *after* the // name-sorted direct set so the cached prefix is unchanged by them. // The same host allow-list gates both halves: deferral only ever - // subtracts from what the host admitted. - let allowed_tools = crate::runtime::host_invocation_binding::(ctx)? - .map(|binding| binding.allowed_tools); + // subtracts from what the host admitted. `resolve_tool_allowlist` + // (not a raw read of `binding.allowed_tools`) is what applies I-9's + // fail-closed default, so an empty declared list denies every tool + // here exactly as it does for the direct set below. + let allowed_tools = self.resolve_tool_allowlist(ctx)?; let host_allows = |name: &str| { allowed_tools .as_ref() - .is_none_or(|allowed| allowed.is_empty() || allowed.contains(name)) + .is_none_or(|allowed| allowed.contains(name)) }; let mut tool_schemas = self .tools @@ -170,9 +212,66 @@ impl AgentHarness { .into_iter() .filter(|schema| host_allows(&schema.name)) .collect::>(); + // Composable toolset chain (gap B3, `AgentHarness::with_toolset`): + // additive to the registry's own `Direct` schemas above — a name the + // registry already advertises keeps the registry's declaration, so a + // registered tool always wins a collision. This run's toolset is + // consulted once here, matching the registry's own once-per-run + // schema build a few lines up (the comment above explains why: the + // resulting request tool list feeds the provider prompt cache, so + // rebuilding it every turn would defeat that cache). A caller that + // genuinely needs true per-turn variance can still call + // [`crate::tool::toolset::ToolSet::tools`] directly from a + // `before_model` middleware, which *does* run every turn. + if let Some(toolset) = &self.toolset { + let existing: std::collections::HashSet<&str> = tool_schemas + .iter() + .map(|schema| schema.name.as_str()) + .collect(); + let extra: Vec<_> = toolset + .tools(ctx) + .await? + .into_iter() + .filter(|tool| tool.exposure() == tinytools::ToolExposure::Direct) + .filter(|tool| host_allows(tool.name())) + .filter(|tool| !existing.contains(tool.name())) + .map(|tool| crate::tool::provider_schema(tool.as_ref())) + .collect(); + tool_schemas.extend(extra); + // Keep the combined set name-sorted: every consumer of + // `tool_schemas` below (and the provider request it feeds) relies + // on the sort for wire-byte/prompt-cache stability. + tool_schemas.sort_by(|left, right| left.name.cmp(&right.name)); + } + // Provider projection applies once, to the full combined set + // (registry + toolset), so a toolset-supplied schema reaches the + // wire cleaned exactly like a registered one. if let Some(preparation) = &self.policy.tool_schemas { tool_schemas = crate::tool::prepare_tool_schemas(&tool_schemas, preparation); } + // B6 (`docs/runtime-comparison/plan.md`): `declared_tool_schemas` + // tracks what the transcript has actually been told about the + // toolset chain's tools so far (folded or patched in, turn by turn, + // by the loop below), so a later turn's live toolset resolution can + // be diffed against it instead of against the wire list — the wire + // list also carries the bridge schemas captured into + // `bridge_schemas` next, which never change within a run and so are + // deliberately excluded from the diff. + // + // Starts empty rather than seeded from the merge above: nothing has + // been recorded on the transcript yet, so the loop's first-turn diff + // (below) always fires when a toolset is installed, declaring the + // full initial toolset-supplied set as one patch (turn 1 has no + // prior cached prefix to protect, so there is no cost to always + // recording it). This is what makes + // [`tinyinference_llm::message::replay_system_state`] able to + // reconstruct the *complete* effective tool set from the transcript + // alone, not just later deltas — the alternative (seeding from the + // merge above) would leave the initial toolset-only tools + // permanently undeclared on the wire-only `tool_schemas` snapshot + // computed here, which a replay can never see. + let mut declared_tool_schemas: Vec = Vec::new(); + let mut bridge_schemas: Vec = Vec::new(); let deferred_catalog = self.deferred_catalog(&host_allows); if !deferred_catalog.is_empty() { // A host-registered `tool_search`/`tool_call` keeps its slot: the @@ -202,19 +301,11 @@ impl AgentHarness { } for schema in bridge { if self.tools.dispatch(&schema.name).is_none() { - tool_schemas.push(schema); + tool_schemas.push(schema.clone()); + bridge_schemas.push(schema); } } } - // The dialect itself is resolved per turn, once the model for that - // turn is known (see the `run_dialect` binding below, right after - // `binding`): `Auto` needs the model's capability to decide between - // `Native` and the documented `Xml` fallback, and that capability is - // not known this early. `tool_schemas` — what the text protocols need - // a registry built from (the *prepared* direct set plus the discovery - // bridge, i.e. exactly what is rendered into the catalogue and can - // come back as a call) — is fixed for the whole run and captured here. - // Fail closed on a structured-output schema whose name collides with a // registered tool *or* the intrinsic discovery bridge. Under the // tool-call strategy the schema is sent as an extra `function` entry, @@ -265,6 +356,26 @@ impl AgentHarness { }); status.set_last_event(record.id); + // Resume (A2): the caller supplied decisions for the tool calls a + // previous run left pending on this transcript. Apply them — answer + // denials and host-supplied results, run approved calls — before + // spending a model call, so the model's next turn sees every call + // answered. An approved call that defers *again* is settled exactly + // like a fresh deferral below. + if let Some(results) = ctx.take_deferred_results() { + let pending = pending_tool_calls(messages)?; + status.mark_running(HarnessPhase::Tools); + let deferred = self + .apply_deferred_results(state, ctx, run, status, messages, pending, results) + .await?; + if let Some(exit) = self + .settle_deferred(state, ctx, run, status, messages, deferred) + .await? + { + return Ok(exit); + } + } + // Truncated-empty recovery state (see `RunPolicy::truncated_empty_retries`). // These persist across the retry `continue` within a single logical turn: // `boosted_max_tokens` overrides the next request's cap, `truncation_base` @@ -277,6 +388,12 @@ impl AgentHarness { let mut boosted_max_tokens: Option = None; let mut truncation_base: Option = None; + // Output-validation retry state (see `RunPolicy::output_retry`, A3). + // Scoped to the whole run rather than reset per turn: `max_attempts` + // is a run-wide ceiling on re-asks, matching `retries.output` in + // Pydantic AI rather than a per-turn allowance. + let mut output_retry_attempts: u8 = 0; + loop { // Safe cancellation checkpoint: if an orchestrator requested // cooperative cancellation, stop before doing any further work @@ -313,8 +430,10 @@ impl AgentHarness { // `after_tool`/`wrap_tool` was honored one full model call late — // an extra billable provider round trip after a guardrail, or a // human gate, had already said stop. - if let Some(exit) = self.apply_pending_control(ctx, run, status)? { - return Ok(exit); + match self.apply_pending_control(ctx, run, status, messages)? { + ControlEffect::None => {} + ControlEffect::ContinueLoop => continue, + ControlEffect::Exit(exit) => return Ok(exit), } // Fail-closed limit and deadline checks before each model call. @@ -339,7 +458,7 @@ impl AgentHarness { ctx.emit(AgentEvent::LimitReached { kind: LimitKind::ModelCalls, }); - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::agent_loop", run_id = %ctx.run_id(), "[agent_loop] model-call cap reached; stopping with the partial run" @@ -359,7 +478,7 @@ impl AgentHarness { self.policy.limits.behavior, crate::limits::LimitBehavior::StopWithPartial ) { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::agent_loop", run_id = %ctx.run_id(), "[agent_loop] model-call cap reached; policy asks to stop with the \ @@ -371,6 +490,61 @@ impl AgentHarness { } } + // B6 (`docs/runtime-comparison/plan.md`, `declare_tool_changes`): + // re-consult the toolset chain (documented as "called once per + // turn", `ToolSet::tools`) and diff its live set against what + // this transcript has declared so far. A caller whose toolset + // never varies turn to turn sees no diff and pays nothing here — + // this only fires for a genuine mid-run change. Deliberately + // runs before the request/`ModelStarted` below, so the patch (if + // any) is part of *this* turn's request. + if let Some(toolset) = &self.toolset { + let mut live_schemas: Vec = self + .tools + .schemas() + .into_iter() + .filter(|schema| host_allows(&schema.name)) + .collect(); + let existing: std::collections::HashSet<&str> = live_schemas + .iter() + .map(|schema| schema.name.as_str()) + .collect(); + let extra: Vec<_> = toolset + .tools(ctx) + .await? + .into_iter() + .filter(|tool| tool.exposure() == tinytools::ToolExposure::Direct) + .filter(|tool| host_allows(tool.name())) + .filter(|tool| !existing.contains(tool.name())) + .map(|tool| crate::tool::provider_schema(tool.as_ref())) + .collect(); + live_schemas.extend(extra); + live_schemas.sort_by(|left, right| left.name.cmp(&right.name)); + if let Some(preparation) = &self.policy.tool_schemas { + live_schemas = crate::tool::prepare_tool_schemas(&live_schemas, preparation); + } + if let Some(patch) = + tool_changes::diff_tool_set(&declared_tool_schemas, &live_schemas) + { + // A cheap, non-mutating preview resolution against the + // transcript as it stands (pre-patch) decides fold vs. + // insert. It is a pure registry lookup (no network call, + // see `ModelRegistry::resolve_request`), so this stays + // proportional to the diff it gates. An unresolved + // preview conservatively folds (`false`): folding is + // always correct, only less cache-friendly. + let mid_conversation = self + .models + .resolve_request(&ModelRequest::new(messages.clone()), None, None) + .and_then(|binding| binding.model.profile().cloned()) + .is_some_and(|profile| profile.mid_conversation_system_messages); + tool_changes::apply_tool_change_patch(messages, patch, mid_conversation); + declared_tool_schemas = live_schemas.clone(); + tool_schemas = declared_tool_schemas.clone(); + tool_schemas.extend(bridge_schemas.clone()); + } + } + // Build the request from the working transcript, tool schemas, and // policy response format. Go through `PromptBuilder` rather than // constructing `ModelRequest` directly: a provider KV cache needs @@ -427,53 +601,35 @@ impl AgentHarness { .run_before_model(ctx, state, &mut request) .await?; - // `ToolDispatcher::Native` is documented as *forcing* provider-native - // tool calls, unlike `Auto`'s "native when available, else Xml". - // `RunDialect::resolve` maps both to the same `Native` variant (it - // only decides whether *this* host renders a text protocol), so - // without a capability requirement that promise was unenforceable: - // a model profile lacking `tool_calling` could still be resolved, - // and a provider adapter is free to fall back to its own text - // encoding for such a profile. Requiring the capability makes - // resolution itself fail closed for an incapable model. An - // adapter's own *runtime* degrade after a live "tools not - // supported" provider response is a separate, adapter-internal - // reliability behavior this host-side dialect selection has no - // visibility into or control over. - // - // Checked against `request.tools` (the *effective* tool set), - // not the pre-`before_model` `tool_schemas` snapshot: a run that - // starts with no tools but whose `before_model` middleware adds - // some must still be gated — checking the earlier snapshot would - // silently let those middleware-added tools reach an - // incapable-of-native-tool-calling model. - // - // Also gated on an `Auto` structured-output format even when - // `request.tools` is still empty here: `StructuredStrategy` - // resolution (below, after `binding`) only ever appends a - // synthetic tool-call schema for a model whose profile already - // has `tool_calling` (`StructuredStrategy::for_profile`'s - // `ToolCall` arm), so requiring it up front is what makes that - // later fact true rather than merely hoped for — by the time - // structured planning knows whether a schema tool is needed the - // model is already resolved, too late to gate resolution on. - // Requiring the capability here is conservatively broader than - // strictly necessary for a model that would have used - // `ProviderSchema` instead, but never wrong: a fail-closed - // requirement narrowing the candidate pool is the point of this - // gate. - let structured_output_may_need_tool_calling = matches!( - self.policy.default_response_format, - Some(ResponseFormat::Auto { .. }) - ); + // A forced native dialect cannot silently select a model that + // lacks provider-native tool calling. This has to happen after + // `before_model`, because middleware may add tools, and before + // model resolution, because the resolver is the capability gate. + // An automatic structured response also needs this gate: its + // fallback may become a native schema tool after selection. if matches!( self.policy.tool_dialect, crate::config::ToolDispatcher::Native - ) && (!request.tools.is_empty() || structured_output_may_need_tool_calling) + ) && (!request.tools.is_empty() + || matches!(request.response_format, Some(ResponseFormat::Auto { .. }))) { - let mut required = request.required_capabilities.clone().unwrap_or_default(); - required.tool_calling = true; - request.required_capabilities = Some(required); + request + .required_capabilities + .get_or_insert_default() + .tool_calling = true; + } + + // Safe checkpoint: a control requested from `before_model_control` + // (for example `BudgetMiddleware` finding the budget already + // exhausted) is honored **before** the model is actually + // dispatched, not one billable call late. Without this checkpoint + // the queued control would only be drained at the next one (after + // this response comes back), spending exactly the call the + // control was raised to prevent. + match self.apply_pending_control(ctx, run, status, messages)? { + ControlEffect::None => {} + ControlEffect::ContinueLoop => continue, + ControlEffect::Exit(exit) => return Ok(exit), } // Resolve the model for the event/log name before invoking. @@ -497,42 +653,6 @@ impl AgentHarness { }; let model_name = binding.resolved.name.clone(); - // Resolved per turn (not once for the whole run) because `Auto` - // needs the *resolved* model's capability, known only now: - // `ToolDispatcher::Auto` is documented as "provider-native tool - // calls when the provider supports them, otherwise Xml", but - // mapping it to the same host-side-no-op behavior as `Native` - // (as an earlier version of this dialect resolution did) left - // that fallback unenforced — a model with `tool_calling: false` - // selected under `Auto` would receive a request that still - // depended on provider-native tools, with no host-rendered text - // protocol and no adapter guaranteed to supply one. `Native` - // stays forced regardless of capability (it fails closed at - // resolution instead, via the capability requirement above); - // `Xml`/`Pformat` stay forced as explicit opt-ins. - let effective_dispatcher = match self.policy.tool_dialect { - crate::config::ToolDispatcher::Auto => { - // A model with *no declared profile at all* is unknown, - // not incapable — treated as capable (the historical - // behavior, and correct for hosts/tests that never - // bother declaring a profile). Only an explicit - // `tool_calling: false` triggers the documented Xml - // fallback. - if binding - .model - .profile() - .is_none_or(|profile| profile.tool_calling) - { - crate::config::ToolDispatcher::Native - } else { - crate::config::ToolDispatcher::Xml - } - } - other => other, - }; - let run_dialect = - super::dialect::RunDialect::resolve(effective_dispatcher, &tool_schemas); - // An explicit request override that resolution skipped (unknown // name, missing capability, or provider-retired) falls through to // a lower-priority candidate by documented fail-closed semantics; @@ -548,6 +668,62 @@ impl AgentHarness { }); } + // Cross-provider handoff: rewrite any part of the outgoing + // transcript that a mid-session provider/model switch left + // unsafe to replay verbatim (foreign signed/redacted thinking, + // non-conforming tool-call ids, unsupported images) right before + // this request is sent. A no-op (same-origin run, the common + // case) allocates nothing — see `handoff_transform`. Runs before + // the schema/reasoning adjustments below so a rewritten + // transcript (rather than the pre-handoff one) is what those + // adjustments and the eventual request see. + if let Some(profile) = binding.model.profile() { + let target_origin = handoff_transform::target_origin_for(profile); + let outcome = handoff_transform::prepare_for_model( + &request.messages, + profile, + &target_origin, + ); + let changes = outcome.changes; + if changes > 0 { + request.messages = outcome.messages.into_owned(); + ctx.emit(AgentEvent::HandoffTransformApplied { changes }); + } + } + + // Apply the resolved model's schema transform (for example + // stripping `$defs` a provider rejects) to every tool schema + // already attached to the request. This is the same wire-shape + // adjustment `SchemaPreparation::schema_transform` performs, run + // here because it depends on the resolved binding's profile, + // which is only known once resolution above has run. + if let Some(transform) = binding + .model + .profile() + .and_then(|profile| profile.schema_transform.as_ref()) + { + for tool in request.tools.iter_mut() { + tool.parameters = transform.apply(&tool.parameters); + } + } + + // A caller asking for a *named* reasoning effort (for example + // `ReasoningEffort::High`) gets whatever generic token that name + // implies unless the resolved model's profile maps that name to + // something more specific for this exact model (a provider-tuned + // `budget_tokens`, typically). Only fill in a name the profile + // actually maps and only when the caller has not already pinned + // an explicit `budget_tokens` — an explicit budget is the + // caller's own override and must win over the profile default. + if let Some(profile) = binding.model.profile() + && let Some(reasoning) = request.reasoning.as_ref() + && reasoning.budget_tokens.is_none() + && let Some(effort) = reasoning.effort + && let Some(mapped) = profile.thinking_level_map.get(effort.as_str()) + { + request.reasoning = Some(mapped.clone()); + } + // Resolve the structured-output plan against the resolved model. // `Auto` consults the model profile to choose provider-native schema // mode versus a tool-call fallback; an explicit `JsonSchema` always @@ -555,7 +731,74 @@ impl AgentHarness { // the final response below. let structured_plan: Option<(StructuredStrategy, String, Value)> = match request.response_format.clone() { + Some(ResponseFormat::Auto { name, schema }) + if matches!( + self.policy.structured_strategy_override, + Some(crate::runtime::StructuredStrategyOverride::Prompted { .. }) + ) => + { + let template = match &self.policy.structured_strategy_override { + Some(crate::runtime::StructuredStrategyOverride::Prompted { + template, + }) => template.clone(), + _ => unreachable!("guarded by the match arm above"), + }; + let schema = crate::tool::apply_profile_schema_transform( + &schema, + binding.model.profile(), + ); + request.response_format = Some(ResponseFormat::Text); + let instructions = template.clone().unwrap_or_else(|| { + crate::structured::default_prompted_template().to_string() + }); + let schema_text = serde_json::to_string_pretty(&schema).unwrap_or_default(); + request.messages.insert( + 0, + Message::system(format!( + "{instructions}\n\nJSON Schema for `{name}`:\n{schema_text}" + )), + ); + Some((StructuredStrategy::Prompted { template }, name, schema)) + } + Some(ResponseFormat::Auto { name, schema }) + if matches!( + self.policy.structured_strategy_override, + Some(crate::runtime::StructuredStrategyOverride::ToolCallUnion { .. }) + ) => + { + let variants = match &self.policy.structured_strategy_override { + Some(crate::runtime::StructuredStrategyOverride::ToolCallUnion { + variants, + }) => variants.clone(), + _ => unreachable!("guarded by the match arm above"), + }; + request.response_format = Some(ResponseFormat::Text); + for (variant_name, variant_schema) in &variants { + let variant_schema = crate::tool::apply_profile_schema_transform( + variant_schema, + binding.model.profile(), + ); + let schema_tool = ToolSchema { + name: variant_name.clone(), + description: format!("Return the result as `{variant_name}`."), + parameters: variant_schema, + format: tinyinference_llm::tool::ToolFormat::Json, + }; + request.tools.push(match &self.policy.tool_schemas { + Some(preparation) => { + crate::tool::prepare_tool_schema(&schema_tool, preparation) + } + None => schema_tool, + }); + } + let _ = schema; + Some((StructuredStrategy::ToolCallUnion, name, Value::Null)) + } Some(ResponseFormat::Auto { name, schema }) => { + let schema = crate::tool::apply_profile_schema_transform( + &schema, + binding.model.profile(), + ); let strategy = StructuredStrategy::for_profile(binding.model.profile()); match strategy { StructuredStrategy::ProviderSchema => { @@ -596,7 +839,7 @@ impl AgentHarness { if tool_schemas.is_empty() { request.tool_choice = ToolChoice::Tool(name.clone()); } else { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::agent_loop", run_id = %ctx.run_id(), schema_name = %name, @@ -606,10 +849,46 @@ impl AgentHarness { ); } } + // A profile whose `default_structured_mode` is + // `Prompted` reaches this arm too (not only + // through the dedicated + // `structured_strategy_override` arm above): the + // schema goes into the system segment instead of + // a provider API field, mirroring the override + // arm's construction. + StructuredStrategy::Prompted { ref template } => { + request.response_format = Some(ResponseFormat::Text); + let instructions = template.clone().unwrap_or_else(|| { + crate::structured::default_prompted_template().to_string() + }); + let schema_text = + serde_json::to_string_pretty(&schema).unwrap_or_default(); + request.messages.insert( + 0, + Message::system(format!( + "{instructions}\n\nJSON Schema for `{name}`:\n{schema_text}" + )), + ); + } + // `for_profile` never returns `ToolCallUnion`; + // that strategy is reached exclusively through + // the dedicated `structured_strategy_override` + // arm above. + StructuredStrategy::ToolCallUnion => unreachable!( + "StructuredStrategy::for_profile never returns ToolCallUnion" + ), } Some((strategy, name, schema)) } Some(ResponseFormat::JsonSchema { name, schema }) => { + let schema = crate::tool::apply_profile_schema_transform( + &schema, + binding.model.profile(), + ); + request.response_format = Some(ResponseFormat::JsonSchema { + name: name.clone(), + schema: schema.clone(), + }); Some((StructuredStrategy::ProviderSchema, name, schema)) } _ => None, @@ -624,32 +903,31 @@ impl AgentHarness { // a recovered call against; the catalogue already advertises it // because it is rendered fresh from `tools` on every call. let offered_tool_count = request.tools.len(); - // An empty recovery when the effective choice is `None`: a - // `before_model` middleware asking for no tool calls this turn - // must actually get none. `apply_to_request` below already skips - // its rewrite for `None`, but that alone left recovery/the - // stream scrubber still treating every offered name as - // recognizable — so a model that narrated `` markup - // as plain text anyway would still have it parsed and dispatched - // as a real, side-effecting call despite the explicit - // prohibition. An empty `offered` list makes every grammar in - // `tinytools-agent` decline to recognize anything as a call. - let recovery = if request.tool_choice == ToolChoice::None { - super::dialect::TextRecovery::default() - } else { - super::dialect::TextRecovery { - offered: Arc::new(request.tools.clone()), - registry: run_dialect.registry_for(&request.tools), - } - }; // Whether this turn could possibly have accepted a tool call at - // all — reuses `recovery.offered`, which is already empty - // exactly when no tools were offered or the effective choice was + // all: tools were offered and the effective choice is not // `None`. Read below by the dropped-tool-call nudge: nudging a // model to "issue the call" when no call could ever have been // accepted wastes up to `dropped_tool_call_nudges` model calls - // asking for something impossible before falling through. - let tools_available_this_turn = !recovery.offered.is_empty(); + // asking for something impossible before falling through. Also + // gates whether `recovery` below is populated at all: an empty + // recovery makes every grammar in `tinytools-agent` decline to + // recognize anything as a call, so a model that narrated + // ``-shaped markup as plain text while explicitly + // told not to call anything is never misread as a real, + // side-effecting call. + let tools_available_this_turn = + offered_tool_count > 0 && request.tool_choice != ToolChoice::None; + let dialect = + super::dialect::RunDialect::resolve(self.policy.tool_dialect, &request.tools); + let forced_text_dialect = dialect.is_text(); + let recovery = if tools_available_this_turn { + super::dialect::TextRecovery { + offered: Arc::new(request.tools.clone()), + registry: dialect.registry_for(&request.tools), + } + } else { + super::dialect::TextRecovery::default() + }; // Applied before budget preflight below: for a text dialect this // rewrite folds the protocol block and full tool catalogue into // `request.messages` and clears `request.tools`, and that is the @@ -658,15 +936,13 @@ impl AgentHarness { // `max_input_tokens` pass admission on the small structured // request and then send a materially larger rendered-text one, // defeating the pre-call budget limit. - run_dialect.apply_to_request(&mut request); + dialect.apply_to_request(&mut request); // A host budget is acquired only for an explicit host-driven run. - // Do it after structured-output planning and the dialect - // rewrite: a synthetic schema tool and, for a text dialect, the - // rendered protocol/catalogue text are both part of the actual - // provider request and must be included in its estimate. The - // permit remains alive through response accounting, so - // cancellation or a provider error still releases it through + // Do it after structured-output planning: a synthetic schema tool + // is part of the provider request and must be included in its + // estimate. The permit remains alive through response accounting, + // so cancellation or a provider error still releases it through // Drop. let host_budget = if let Some(host_run) = crate::runtime::host_invocation_binding::(ctx)? @@ -685,10 +961,7 @@ impl AgentHarness { }; let hint = budget.compression_hint(&context_state); if hint.is_advised() { - tinyagents_tracing::debug!( - ?hint, - "[host] budget gate advised context compression" - ); + tracing::debug!(?hint, "[host] budget gate advised context compression"); apply_host_budget_compression(ctx, &mut request.messages, hint)?; } let estimate = crate::host::CallEstimate::new( @@ -696,34 +969,21 @@ impl AgentHarness { crate::token_estimation::estimate_slice_tokens(&request.messages), request.max_tokens.unwrap_or_default() as u64, ) - .with_agent(host_run.agent_id) + .with_agent(host_run.agent_id.clone()) .with_thread( ctx.thread_id() .cloned() .unwrap_or_else(|| ctx.run_id().as_str().into()), ) .with_tool_count(offered_tool_count); - let permit = match self.call_budget(ctx) { - Some(remaining) => tokio::select! { - biased; - _ = ctx.cancellation.cancelled() => { - return Err(TinyAgentsError::Cancelled); - } - acquired = tokio::time::timeout(remaining, budget.acquire(&estimate)) => { - acquired.map_err(|_| TinyAgentsError::Timeout(format!( - "budget admission for run `{}` exceeded its remaining wall-clock deadline", - ctx.run_id() - )))?? - } - }, - None => tokio::select! { - biased; - _ = ctx.cancellation.cancelled() => { - return Err(TinyAgentsError::Cancelled); - } - acquired = budget.acquire(&estimate) => acquired?, - }, - }; + let permit = ctx + .bounded(self.call_budget(ctx), budget.acquire(&estimate), || { + format!( + "budget admission for run `{}` exceeded its remaining wall-clock deadline", + ctx.run_id() + ) + }) + .await?; Some((budget.clone(), permit)) } else { None @@ -731,10 +991,14 @@ impl AgentHarness { } else { None }; - let call_id = CallId::new(format!("{}-model-{}", ctx.run_id(), run.model_calls + 1)); status.mark_running(HarnessPhase::Model); status.active_model_call = Some(call_id.clone()); + // Mirrored onto the context so `ModelMiddleware` (e.g. + // `RetryMiddleware`) can correlate its own events with the exact + // call id the loop uses instead of deriving an uncorrelated one + // (I-7). Cleared right after the wrap onion returns, below. + ctx.active_model_call = Some(call_id.clone()); // Captured here (where the call actually starts) so the completed // event carries a real start time for duration-aware exporters. let model_started_at_ms = crate::ids::now_ms(); @@ -744,13 +1008,30 @@ impl AgentHarness { }); status.set_last_event(record.id); + // Captured before `binding.model` moves into `base` below: decides + // whether text-dialect recovery should even be attempted for this + // call's response (see the call site after the model returns). A + // forced text dialect always recovers regardless of this flag — + // the model can only answer in text, so parsing it is the + // protocol, not a fallback. Under `Native`, `Auto` skips a model + // whose resolved profile reports native tool calling, since such + // a model that still answered in prose was explaining or quoting + // the format, not making a call (I-2). + let text_dialect_recovery_enabled = match self.policy.text_dialect_recovery { + crate::runtime::TextDialectRecovery::Off => false, + crate::runtime::TextDialectRecovery::On => true, + crate::runtime::TextDialectRecovery::Auto => !binding + .model + .profile() + .map(|profile| profile.tool_calling) + .unwrap_or(false), + }; + // The real model call (cache + retry + fallback core) is the // innermost base of the model-wrap onion. Lifecycle `before_model` // already ran above; the wrap onion runs here; lifecycle // `after_model` runs below — so ordering is: // before_model -> wrap onion (outer..inner..base) -> after_model. - // `recovery` and the dialect rewrite were computed above, before - // budget preflight (see the comment there for why). let base = ModelCallBase { harness: self, call_id: call_id.clone(), @@ -774,23 +1055,37 @@ impl AgentHarness { // model-wrap onion, so truncated-empty recovery can compute the next // (doubled) budget from what was actually sent. let attempt_max_tokens = request.max_tokens; - let mut response = self + let (mut response, wrap_control) = self .middleware .run_wrapped_model(ctx, state, request, &base) .await? - .into_response(); + .into_response_with_control(); + // A `ModelMiddleware::wrap_model` that short-circuited with + // `MiddlewareModelOutcome::Command` carries no real response (see + // that variant's docs); queue its control the same way a + // lifecycle hook's control-outcome return would, so the next safe + // checkpoint (right below, after this turn's bookkeeping) applies + // it instead of the placeholder response being mistaken for a + // real completion. + if let Some(control) = wrap_control { + ctx.request_control(control); + } // Providers occasionally put a text-dialect call in visible // content even when a native tool channel was offered, and a // forced text dialect always does. Read the response through - // every grammar the protocol crate knows, but only when the - // provider did not already supply structured calls. - super::dialect::recover_text_calls( - &mut response, - &call_id, - &recovery.offered, - recovery.registry.as_deref(), - ); + // every grammar the protocol crate knows (`recover_text_calls`), + // but only when the provider did not already supply structured + // calls. `recovery` is already empty when this turn offered no + // tools or the effective tool choice was `None` (computed + // above); the `forced_text_dialect || text_dialect_recovery_enabled` + // gate additionally skips a resolved model whose profile reports + // native tool calling under `RunPolicy::text_dialect_recovery`'s + // `Auto` default. The fenced-code guard and the audit event live + // in the wrapper. + if forced_text_dialect || text_dialect_recovery_enabled { + recover_text_dialect_calls(ctx, &mut response, &call_id, &recovery); + } // Account for the completed provider response before fallible // response middleware. A middleware rejection must not erase @@ -800,13 +1095,14 @@ impl AgentHarness { run.steps += 1; status.model_calls = run.model_calls; status.active_model_call = None; + ctx.active_model_call = None; // A cache replay consumed no provider tokens, so folding its usage // into the run's totals reports spend that never happened. The // saving is surfaced through the cache-hit event instead of being // buried in the spend total. if let Some(usage) = response.usage { if response.served_from_cache { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::agent_loop", run_id = %ctx.run_id(), call_id = %call_id, @@ -863,8 +1159,10 @@ impl AgentHarness { // Safe checkpoint: honor any control outcome a middleware requested // during this turn (for example an early-exit tool or a budget stop // hook), before executing further tools. - if let Some(exit) = self.apply_pending_control(ctx, run, status)? { - return Ok(exit); + match self.apply_pending_control(ctx, run, status, messages)? { + ControlEffect::None => {} + ControlEffect::ContinueLoop => continue, + ControlEffect::Exit(exit) => return Ok(exit), } let tool_calls = response.tool_calls().to_vec(); @@ -876,54 +1174,154 @@ impl AgentHarness { // sibling call in the same turn — a turn returning // `[search(...), my_schema(...)]` broke out with `search` never // executed and no event to say so. - let structured_call_name = match &structured_plan { - Some((StructuredStrategy::ToolCall, name, _)) => Some(name.clone()), - _ => None, + let structured_call_names: Vec = match &structured_plan { + Some((StructuredStrategy::ToolCall, name, _)) => vec![name.clone()], + Some((StructuredStrategy::ToolCallUnion, _, _)) => { + match &self.policy.structured_strategy_override { + Some(crate::runtime::StructuredStrategyOverride::ToolCallUnion { + variants, + }) => variants.iter().map(|(n, _)| n.clone()).collect(), + _ => Vec::new(), + } + } + _ => Vec::new(), }; let (structured_hits, real_tool_calls): (Vec, Vec) = - match &structured_call_name { - Some(name) => tool_calls + if structured_call_names.is_empty() { + (Vec::new(), tool_calls.clone()) + } else { + tool_calls .iter() .cloned() - .partition(|call| &call.name == name), - None => (Vec::new(), tool_calls.clone()), + .partition(|call| structured_call_names.contains(&call.name)) }; let structured_tool_hit = !structured_hits.is_empty(); if structured_tool_hit && !real_tool_calls.is_empty() { - // Record the structured payload the model already produced, - // then run the real tools it asked for in the same turn and let - // the loop continue; the model finishes on a later turn. - if let Some((strategy, name, schema)) = &structured_plan { - let extractor = - StructuredExtractor::new(*strategy, name.clone(), schema.clone()); - match extractor.extract(&response) { - Ok(output) => run.structured = Some(output.value), - Err(error) => tinyagents_tracing::debug!( - target: "tinyagents::agent_loop", - run_id = %ctx.run_id(), - %error, - "[agent_loop] structured extraction failed on a mixed turn; \ - continuing with the real tool calls" - ), - } - } + // A6: one turn asked to both answer (the structured-output + // schema call) and run further tools. `RunPolicy::end_strategy` + // decides what happens to the two, replacing the old + // ad-hoc "record and keep going" behavior with three named, + // documented outcomes (`EndStrategy`). let record = ctx.emit(AgentEvent::ControlApplied { control: "structured_with_tool_calls".to_string(), detail: format!( - "structured output recorded alongside {} real tool call(s); \ - the run continues", + "{:?} end_strategy handling {} real tool call(s) alongside a \ + structured-output call", + self.policy.end_strategy, real_tool_calls.len() ), }); status.set_last_event(record.id); - // Every requested `tool_call_id` must be answered or the - // transcript is malformed for the next provider call. + if matches!(self.policy.end_strategy, EndStrategy::Early) { + // Finish immediately: the structured answer wins outright, + // and the accompanying tool calls never run. Every + // requested `tool_call_id` — structured hits and the + // skipped real calls alike — still needs an answer or the + // transcript is malformed for a future replay. + if let Some((strategy, name, schema)) = &structured_plan { + let extractor = self.build_structured_extractor(strategy, name, schema); + match extractor.extract(&response) { + Ok(output) => { + run.structured = Some(output.value); + run.structured_variant = output.variant; + } + Err(error) => tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + %error, + "[agent_loop] structured extraction failed on a mixed turn \ + under EndStrategy::Early" + ), + } + } + for call in &structured_hits { + messages.push(Message::tool( + call.id.clone(), + "Structured output recorded.", + )); + } + for call in &real_tool_calls { + messages.push(Message::tool( + call.id.clone(), + "run stopped before this tool call was executed \ + (EndStrategy::Early: the structured answer ends the run first)", + )); + } + run.final_response = Some(response); + if self + .continue_from_queue_at_finish(ctx, status, messages) + .await + { + continue; + } + return Ok(LoopExit::Finished); + } + + if matches!(self.policy.end_strategy, EndStrategy::Graceful) { + // Record the answer now (it will not be asked for again), + // but let the requested tools actually run before ending + // the run — their side effects and results are not + // silently dropped, unlike `Early`. + if let Some((strategy, name, schema)) = &structured_plan { + let extractor = self.build_structured_extractor(strategy, name, schema); + match extractor.extract(&response) { + Ok(output) => { + run.structured = Some(output.value); + run.structured_variant = output.variant; + } + Err(error) => tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + %error, + "[agent_loop] structured extraction failed on a mixed turn \ + under EndStrategy::Graceful" + ), + } + } + for call in &structured_hits { + messages.push(Message::tool( + call.id.clone(), + "Structured output recorded.", + )); + } + status.mark_running(HarnessPhase::Tools); + let deferred = self + .execute_tools(state, ctx, run, status, messages, real_tool_calls) + .await?; + if let Some(exit) = self + .settle_deferred(state, ctx, run, status, messages, deferred) + .await? + { + return Ok(exit); + } + if let ControlEffect::Exit(exit) = + self.apply_pending_control(ctx, run, status, messages)? + { + return Ok(exit); + } + run.final_response = Some(response); + if self + .continue_from_queue_at_finish(ctx, status, messages) + .await + { + continue; + } + return Ok(LoopExit::Finished); + } + + // `EndStrategy::Exhaustive`: the output tool this turn is + // ignored outright (never recorded) — the run keeps going + // exactly as if only the real tool calls had been requested. + // It only finishes once a later turn's output-tool call has + // no accompanying function-tool calls. + debug_assert!(matches!(self.policy.end_strategy, EndStrategy::Exhaustive)); for call in &structured_hits { messages.push(Message::tool( call.id.clone(), - "Structured output recorded. Continue with the remaining tool calls.", + "Structured output noted but not final yet; finish the remaining tool \ + calls first (EndStrategy::Exhaustive).", )); } @@ -942,13 +1340,26 @@ impl AgentHarness { ); status.mark_running(HarnessPhase::Tools); - self.execute_tools(state, ctx, run, status, messages, real_tool_calls) + let deferred = self + .execute_tools(state, ctx, run, status, messages, real_tool_calls) .await?; + if let Some(exit) = self + .settle_deferred(state, ctx, run, status, messages, deferred) + .await? + { + return Ok(exit); + } + + // Turn boundary (A4): same steer drain as the plain tool path. + self.apply_queued_lane(ctx, status, messages, crate::run_queue::QueueLane::Steer) + .await; // Safe checkpoint: a control requested from `after_tool` / // `wrap_tool` is honored here, at the edge it was raised on. - if let Some(exit) = self.apply_pending_control(ctx, run, status)? { - return Ok(exit); + match self.apply_pending_control(ctx, run, status, messages)? { + ControlEffect::None => {} + ControlEffect::ContinueLoop => continue, + ControlEffect::Exit(exit) => return Ok(exit), } continue; } @@ -1041,11 +1452,58 @@ impl AgentHarness { // Final response: optionally extract structured output using the // resolved plan (provider-native schema or tool-call arguments). + // + // A3: extraction failure (schema-invalid/unparseable) or a + // registered `OutputValidator` rejecting an otherwise + // schema-valid value with `TinyAgentsError::ModelRetry` no + // longer immediately fails the run. Both feed the same + // output-validation retry loop — re-ask the model with the + // error as a repair prompt, bounded by + // `RunPolicy::output_retry.max_attempts` — because a + // schema-valid-but-wrong answer and a malformed one are the + // same failure from the caller's perspective: the model needs + // another turn to fix it. if let Some((strategy, name, schema)) = &structured_plan { - let extractor = - StructuredExtractor::new(*strategy, name.clone(), schema.clone()); - let output = extractor.extract(&response)?; - run.structured = Some(output.value); + let extractor = self.build_structured_extractor(strategy, name, schema); + let outcome = extractor.extract_outcome(&response); + let variant = outcome.variant.clone(); + let error = match outcome.value { + Some(value) => match &self.output_validator { + Some(validator) => match validator.validate(ctx, state, &value).await { + Ok(()) => { + run.structured = Some(value); + run.structured_variant = variant; + None + } + Err(TinyAgentsError::ModelRetry(message)) => Some(message), + Err(other) => return Err(other), + }, + None => { + run.structured = Some(value); + run.structured_variant = variant; + None + } + }, + None => outcome.error, + }; + if let Some(error) = error { + if output_retry_attempts < self.policy.output_retry.max_attempts { + output_retry_attempts += 1; + let record = ctx.emit(AgentEvent::OutputRetry { + attempt: output_retry_attempts, + error: error.clone(), + }); + status.set_last_event(record.id); + let prompt = self + .policy + .output_retry + .message_template + .replace("{error}", &error); + messages.push(Message::user(prompt)); + continue; + } + return Err(TinyAgentsError::StructuredOutput(error)); + } } // An empty provider completion — no text, no tool calls, and no // structured output — must not silently become the terminal @@ -1063,6 +1521,14 @@ impl AgentHarness { return Err(TinyAgentsError::EmptyResponse); } run.final_response = Some(response); + // Natural finish (A4): queued steering or a follow-up turns + // "done" into "one more turn" instead of returning. + if self + .continue_from_queue_at_finish(ctx, status, messages) + .await + { + continue; + } return Ok(LoopExit::Finished); } @@ -1082,47 +1548,332 @@ impl AgentHarness { // `agent_loop/tools.rs` for the dispatch rules and the semantics // preserved in each mode. status.mark_running(HarnessPhase::Tools); - self.execute_tools(state, ctx, run, status, messages, real_tool_calls) + let deferred = self + .execute_tools(state, ctx, run, status, messages, real_tool_calls) .await?; + // A2: a batch that deferred calls either resolves them inline + // (handler registered) or ends the run here with the pending + // requests; the non-deferred siblings' results are already on + // the transcript. + if let Some(exit) = self + .settle_deferred(state, ctx, run, status, messages, deferred) + .await? + { + return Ok(exit); + } + + // Turn boundary (A4): every tool result of this batch is on the + // transcript, so queued steering can be applied now — never + // mid-batch — before the next model call sees it. + self.apply_queued_lane(ctx, status, messages, crate::run_queue::QueueLane::Steer) + .await; + + // Turn boundary: give every middleware a chance to end the run + // based on the whole turn's tool results rather than any single + // call (see `Middleware::should_stop_after_turn`). A `Middleware` + // hook could already have requested `JumpTo(End)` from + // `after_tool_control`; this is the aggregate counterpart for a + // decision that only makes sense once the whole turn has settled. + if self.middleware.any_should_stop_after_turn(ctx, run) { + ctx.request_control(MiddlewareControl::JumpTo(LoopTarget::End)); + } // Safe checkpoint: honor a control requested from `after_tool` / // `wrap_tool` at the edge it was raised on, rather than a model // call later. - if let Some(exit) = self.apply_pending_control(ctx, run, status)? { - return Ok(exit); + match self.apply_pending_control(ctx, run, status, messages)? { + ControlEffect::None => {} + ControlEffect::ContinueLoop => continue, + ControlEffect::Exit(exit) => return Ok(exit), + } + } + } + + /// Takes the pending items of `lane` from the run's queue (per + /// [`RunPolicy::queue_mode`][crate::runtime::RunPolicy::queue_mode]), + /// appends them to the working transcript, and emits + /// [`AgentEvent::QueuedMessageApplied`] (A4). Returns whether anything + /// was applied. A run without a queue never applies anything. + async fn apply_queued_lane( + &self, + ctx: &mut RunContext, + status: &mut HarnessRunStatus, + messages: &mut Vec, + lane: crate::run_queue::QueueLane, + ) -> bool { + let Some(queue) = ctx.run_queue.clone() else { + return false; + }; + let items = queue.take(lane, self.policy.queue_mode).await; + if items.is_empty() { + return false; + } + let count = items.len(); + messages.extend(items); + let record = ctx.emit(AgentEvent::QueuedMessageApplied { lane, count }); + status.set_last_event(record.id); + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + lane = lane.as_str(), + count, + "[agent_loop] applied queued messages to the transcript" + ); + true + } + + /// The natural-finish queue boundary (A4): the model produced a final + /// answer, so pending `Steer` items (first) or, when there are none, + /// `Followup` items are appended and the loop runs another turn instead + /// of returning. Returns whether the loop should continue. Only reached + /// from the paths where the *model* finished — a middleware stop, a + /// limit stop, a pause, or a deferral is terminal and leaves the queue + /// untouched for the host. + async fn continue_from_queue_at_finish( + &self, + ctx: &mut RunContext, + status: &mut HarnessRunStatus, + messages: &mut Vec, + ) -> bool { + if ctx.run_queue.is_none() { + return false; + } + self.apply_queued_lane(ctx, status, messages, crate::run_queue::QueueLane::Steer) + .await + || self + .apply_queued_lane(ctx, status, messages, crate::run_queue::QueueLane::Followup) + .await + } + + /// Settles the calls a batch deferred (A2). + /// + /// Returns `Ok(None)` when nothing was deferred, or when a registered + /// [`crate::tool::DeferredToolHandler`] resolved every pending call and + /// the loop can continue. Returns `Ok(Some(LoopExit::Deferred))` when + /// the caller must resolve the requests — no handler, or an approved + /// call deferred a second time (surfaced rather than re-asked, so a + /// handler and a tool that never agree cannot spin). + async fn settle_deferred( + &self, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + messages: &mut Vec, + deferred: crate::tool::DeferredToolRequests, + ) -> Result> { + if deferred.is_empty() { + return Ok(None); + } + let Some(handler) = &self.deferred_tool_handler else { + return Ok(Some(LoopExit::Deferred(deferred))); + }; + let results = handler.handle(&deferred).await?; + let pending: Vec = deferred + .approvals + .iter() + .chain(deferred.calls.iter()) + .cloned() + .collect(); + let again = self + .apply_deferred_results(state, ctx, run, status, messages, pending, results) + .await?; + if again.is_empty() { + return Ok(None); + } + Ok(Some(LoopExit::Deferred(again))) + } + + /// Applies host decisions to `pending` deferred calls (A2): every call + /// must be resolved (`Validation` error naming the missing ids + /// otherwise). Host-supplied results and denials are answered without + /// running a tool; approvals run the tool now through the ordinary + /// serial pipeline, with the model's or the approver's edited + /// arguments. Returns whatever the approved calls deferred *again*. + #[allow(clippy::too_many_arguments)] + async fn apply_deferred_results( + &self, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + messages: &mut Vec, + pending: Vec, + mut results: crate::tool::DeferredToolResults, + ) -> Result { + let missing: Vec<&str> = pending + .iter() + .filter(|call| !results.resolves(&CallId::new(call.id.clone()))) + .map(|call| call.id.as_str()) + .collect(); + if !missing.is_empty() { + return Err(TinyAgentsError::Validation(format!( + "cannot resume: deferred tool calls still unresolved: [{}]", + missing.join(", ") + ))); + } + let mut deferred = crate::tool::DeferredToolRequests::default(); + // Follow-up user messages (B2) trail the whole resumed batch, for + // the same provider-ordering reason as in `execute_tools`. + let mut follow_ups = Vec::new(); + for mut call in pending { + let call_id = CallId::new(call.id.clone()); + if let Some(outcome) = results.calls.remove(&call_id) { + follow_ups.extend( + self.recover_tool_call( + state, + ctx, + run, + status, + messages, + &call, + outcome.into_tool_result(), + ) + .await?, + ); + continue; + } + let decision = results + .approvals + .remove(&call_id) + .expect("every pending call was validated as resolved above"); + match decision { + crate::tool::ToolApprovalDecision::Deny { message } => { + let record = ctx.emit(AgentEvent::ToolDenied { + call_id, + message: message.clone(), + }); + status.set_last_event(record.id); + follow_ups.extend( + self.recover_tool_call( + state, + ctx, + run, + status, + messages, + &call, + tinytools::ToolResult::error(message), + ) + .await?, + ); + } + decision => { + if let crate::tool::ToolApprovalDecision::ApproveWithArgs(arguments) = decision + { + call.arguments = arguments; + } + let record = ctx.emit(AgentEvent::ToolApproved { call_id }); + status.set_last_event(record.id); + ctx.mark_call_approved(call.id.clone()); + follow_ups.extend( + self.execute_tool_serially( + state, + ctx, + run, + status, + messages, + call, + &mut deferred, + ) + .await?, + ); + } } } + super::tools::append_follow_ups(messages, follow_ups); + Ok(deferred) } /// Drains any pending [`MiddlewareControl`] and turns it into a loop /// decision. /// - /// Returns `Ok(None)` when nothing was requested, `Ok(Some(exit))` when the - /// loop must stop, and `Err` for - /// [`MiddlewareControl::Interrupt`]. Called at every safe checkpoint — the - /// top of an iteration, after the model call, and after tool execution — so - /// a control raised anywhere in a turn takes effect on that turn. + /// Returns [`ControlEffect::None`] when nothing was requested (or the + /// pending request needed no loop-level action, e.g. + /// [`MiddlewareControl::UpdateState`]), [`ControlEffect::ContinueLoop`] + /// when the current turn must be abandoned in favor of a fresh iteration + /// (`JumpTo(Model)`), and [`ControlEffect::Exit`] when the run is done. + /// `Err` surfaces [`MiddlewareControl::Interrupt`]. Called at every safe + /// checkpoint — the top of an iteration, after the model call, and after + /// tool execution — so a control raised anywhere in a turn takes effect on + /// that turn. fn apply_pending_control( &self, ctx: &mut RunContext, run: &mut AgentRun, status: &mut HarnessRunStatus, - ) -> Result> { + messages: &mut Vec, + ) -> Result { let Some(control) = ctx.take_control() else { - return Ok(None); + return Ok(ControlEffect::None); }; + // `UpdateState` is applied (queued, really — see `RunContext:: + // push_state_update`) silently: it carries no loop-level decision, so + // audit-logging it as a `ControlApplied` event alongside jumps and + // stops would be noise. It still shows up wherever the host inspects + // `RunContext::take_state_updates`. + if let MiddlewareControl::UpdateState(update) = control { + ctx.push_state_update(update); + return Ok(ControlEffect::None); + } let record = ctx.emit(AgentEvent::ControlApplied { control: control.kind().to_string(), detail: match &control { + MiddlewareControl::Continue => String::new(), + MiddlewareControl::JumpTo(target) => format!("{target:?}"), + MiddlewareControl::UpdateState(_) => unreachable!("handled above"), MiddlewareControl::StopWithFinal(text) => text.clone(), MiddlewareControl::Interrupt { node, message } => format!("{node}: {message}"), }, }); status.set_last_event(record.id); match control { + MiddlewareControl::Continue => Ok(ControlEffect::None), + MiddlewareControl::UpdateState(_) => unreachable!("handled above"), + MiddlewareControl::JumpTo(LoopTarget::Tools) => { + // Tool execution already runs whenever the turn produced real + // tool calls; there is nothing else to route to when it did + // not. Either way, this is a no-op at the loop level. + Ok(ControlEffect::None) + } + MiddlewareControl::JumpTo(LoopTarget::Model) => { + // Abandon whatever the rest of this turn would have done + // (typically: running tools the model just requested) and go + // straight to a fresh model call. Close out any tool calls on + // the last assistant row first so the transcript stays + // replayable (see the `StopWithFinal` arm below for why). + Self::close_unanswered_tool_calls( + messages, + "run jumped back to the model before this tool call was executed", + ); + Ok(ControlEffect::ContinueLoop) + } + MiddlewareControl::JumpTo(LoopTarget::End) => { + Self::close_unanswered_tool_calls( + messages, + "run stopped before this tool call was executed", + ); + if run.final_response.is_none() { + let text = Self::last_assistant_text(messages); + run.final_response = Some(ModelResponse::assistant(text)); + } + Ok(ControlEffect::Exit(LoopExit::Finished)) + } MiddlewareControl::StopWithFinal(text) => { + // The most recently appended assistant row may carry + // `tool_calls` that were never answered — e.g. a middleware + // requesting `StopWithFinal` right after the model turn that + // requested them, before `execute_tools` ever ran. Left as + // is, `run.messages`/`messages` end with an assistant row + // whose tool calls have no matching tool message, which a + // provider rejects (400) if the transcript is ever replayed + // (M-1). Append a synthetic tool result for each unanswered + // call so the transcript stays replayable. + Self::close_unanswered_tool_calls( + messages, + "run stopped before this tool call was executed", + ); run.final_response = Some(ModelResponse::assistant(text)); - Ok(Some(LoopExit::Finished)) + Ok(ControlEffect::Exit(LoopExit::Finished)) } MiddlewareControl::Interrupt { node, message } => { Err(TinyAgentsError::Interrupted { node, message }) @@ -1130,6 +1881,65 @@ impl AgentHarness { } } + /// The text of the most recent assistant message, or empty when there is + /// none. Used to synthesize a final response for + /// [`MiddlewareControl::JumpTo`]`(`[`LoopTarget::End`]`)`, which (unlike + /// [`MiddlewareControl::StopWithFinal`]) carries no text of its own. + fn last_assistant_text(messages: &[Message]) -> String { + messages + .iter() + .rev() + .find(|message| matches!(message, Message::Assistant(_))) + .map(Message::text) + .unwrap_or_default() + } + + /// Appends a synthetic [`Message::tool`] result for every tool call on + /// the last message that is still unanswered, so the transcript stays + /// replayable through a provider that requires every `tool_calls` entry + /// on an assistant message to have a matching tool result before the next + /// turn (M-1). A no-op when the last message is not an unanswered + /// assistant tool-call row. + fn close_unanswered_tool_calls(messages: &mut Vec, reason: &str) { + let Some(Message::Assistant(last)) = messages.last() else { + return; + }; + if last.tool_calls.is_empty() { + return; + } + let synthetic: Vec = last + .tool_calls + .iter() + .map(|call| Message::tool(call.id.clone(), reason)) + .collect(); + messages.extend(synthetic); + } + + /// Builds the [`StructuredExtractor`] for a resolved `structured_plan` + /// entry (A6). + /// + /// [`StructuredStrategy::ToolCallUnion`] needs its variant list, which + /// `structured_plan`'s `(strategy, name, schema)` tuple has nowhere to + /// carry — the variants live on + /// [`crate::runtime::RunPolicy::structured_strategy_override`] instead, + /// which this reaches back into rather than widening the tuple. Every + /// other strategy builds the extractor directly from the tuple as + /// before. + fn build_structured_extractor( + &self, + strategy: &StructuredStrategy, + name: &str, + schema: &Value, + ) -> StructuredExtractor { + if matches!(strategy, StructuredStrategy::ToolCallUnion) + && let Some(crate::runtime::StructuredStrategyOverride::ToolCallUnion { variants }) = + &self.policy.structured_strategy_override + { + return StructuredExtractor::new_union(name, variants.clone()); + } + StructuredExtractor::new(strategy.clone(), name.to_string(), schema.clone()) + } + /// Resolves the effective response-cache decision for `request`. /// /// Returns `Some((cache, key))` when a [`ResponseCache`] is attached to the @@ -1177,23 +1987,14 @@ impl AgentHarness { budget: &Arc, usage: &tinyinference_llm::usage::Usage, ) -> Result<()> { - let cancellation = ctx.cancellation.clone(); let recording = budget.record(usage); - match self.call_budget(ctx) { - Some(remaining) => tokio::select! { - biased; - _ = cancellation.cancelled() => Err(TinyAgentsError::Cancelled), - result = tokio::time::timeout(remaining, recording) => result.map_err(|_| TinyAgentsError::Timeout(format!( - "budget usage recording for run `{}` exceeded its remaining wall-clock deadline", - ctx.run_id() - )))?, - }, - None => tokio::select! { - biased; - _ = cancellation.cancelled() => Err(TinyAgentsError::Cancelled), - result = recording => result, - }, - } + ctx.bounded(self.call_budget(ctx), recording, || { + format!( + "budget usage recording for run `{}` exceeded its remaining wall-clock deadline", + ctx.run_id() + ) + }) + .await } } @@ -1329,6 +2130,78 @@ fn apply_host_budget_compression( Ok(()) } +/// Whether `recover_text_dialect_calls` should even attempt to parse `text`. +/// +/// Fenced code blocks are always skipped regardless of +/// [`crate::runtime::TextDialectRecovery`]: a model demonstrating +/// `` syntax inside a ``` fence — explaining the format, echoing a +/// worked example — is manifestly not making a call, and recovering it would +/// silently execute quoted documentation as a real action. +fn text_dialect_markup_only_in_fenced_code(text: &str) -> bool { + let mut in_fence = false; + let mut saw_marker_outside_fence = false; + let mut saw_marker_anywhere = false; + for line in text.lines() { + if line.trim_start().starts_with("```") { + in_fence = !in_fence; + continue; + } + if line.contains("( + ctx: &RunContext, + response: &mut tinyinference_llm::model::ModelResponse, + model_call_id: &CallId, + recovery: &super::dialect::TextRecovery, +) { + if recovery.offered.is_empty() { + return; + } + + if text_dialect_markup_only_in_fenced_code(&response.text()) { + return; + } + + let before = response.message.tool_calls.len(); + super::dialect::recover_text_calls( + response, + model_call_id, + &recovery.offered, + recovery.registry.as_deref(), + ); + let recovered = response.message.tool_calls.len().saturating_sub(before); + if recovered == 0 { + return; + } + + ctx.emit(AgentEvent::ControlApplied { + control: "text_dialect_recovered".to_string(), + detail: format!( + "recovered {recovered} text-dialect tool call(s) from model call `{model_call_id}`" + ), + }); +} + /// The re-prompt sent when a model signalled a tool call it did not make. /// Deliberately terse and instruction-free beyond the one thing needed: the /// task and the tools are already in the transcript. @@ -1336,6 +2209,42 @@ const DROPPED_TOOL_CALL_NUDGE: &str = "Your previous turn indicated a tool call included. If you meant to call a tool, issue the actual tool call now; otherwise answer \ directly."; +/// The tool calls on the transcript's last assistant row that have no +/// matching tool-result row after it — the calls a previous run deferred +/// (A2). Errors when the transcript has nothing to resume. +fn pending_tool_calls(messages: &[Message]) -> Result> { + let Some(assistant_at) = messages + .iter() + .rposition(|message| matches!(message, Message::Assistant(_))) + else { + return Err(TinyAgentsError::Validation( + "cannot resume: the transcript has no assistant tool-call row".to_string(), + )); + }; + let Message::Assistant(assistant) = &messages[assistant_at] else { + unreachable!("rposition matched an assistant row"); + }; + let answered: std::collections::HashSet<&str> = messages[assistant_at + 1..] + .iter() + .filter_map(|message| match message { + Message::Tool(tool) => Some(tool.tool_call_id.as_str()), + _ => None, + }) + .collect(); + let pending: Vec = assistant + .tool_calls + .iter() + .filter(|call| !answered.contains(call.id.as_str())) + .cloned() + .collect(); + if pending.is_empty() { + return Err(TinyAgentsError::Validation( + "cannot resume: the transcript has no unanswered tool calls".to_string(), + )); + } + Ok(pending) +} + /// Resolves one run-scoped call cap from the per-run [`RunConfig`] value and /// the harness-wide [`crate::runtime::RunPolicy`] value. /// @@ -1364,3 +2273,112 @@ fn reset_truncated_empty_recovery( *boosted_max_tokens = None; *truncation_base = None; } + +#[cfg(test)] +mod recovery_tests { + use std::sync::Arc; + + use super::recover_text_dialect_calls; + use crate::agent_loop::dialect::TextRecovery; + use crate::context::{RunConfig, RunContext}; + use crate::ids::CallId; + use tinyinference_llm::model::ModelResponse; + use tinyinference_llm::tool::ToolSchema; + + fn offered(names: &[&str]) -> TextRecovery { + TextRecovery { + offered: Arc::new( + names + .iter() + .map(|name| ToolSchema::new(*name, "", serde_json::json!({"type": "object"}))) + .collect(), + ), + registry: None, + } + } + + #[test] + fn text_dialect_markup_is_not_recovered_when_the_request_offered_no_tools() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("recovery-test"), ()); + let mut response = ModelResponse::assistant( + "shell{\"command\":\"id\"}", + ); + + recover_text_dialect_calls( + &ctx, + &mut response, + &CallId::new("model-1"), + &TextRecovery::default(), + ); + + assert!(response.message.tool_calls.is_empty()); + assert!(response.text().contains("")); + } + + /// I-2 regression: `RunPolicy::text_dialect_recovery` resolving to off + /// (what `Auto` yields for a model whose profile reports native tool + /// calling) is represented as an empty `TextRecovery`, exactly like a + /// turn that offered no tools — so `` markup the model merely + /// quoted must not be executed. + #[test] + fn text_dialect_markup_is_not_recovered_when_the_policy_disables_it() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("recovery-test"), ()); + let mut response = ModelResponse::assistant( + r#"{"name": "shell", "arguments": {"command": "id"}}"#, + ); + + recover_text_dialect_calls( + &ctx, + &mut response, + &CallId::new("model-1"), + &TextRecovery::default(), + ); + + assert!(response.message.tool_calls.is_empty()); + assert!(response.text().contains("")); + } + + /// I-2 regression: a final answer that quotes `` markup inside + /// a fenced code block must never be executed, even when recovery is + /// otherwise enabled and tools were offered. + #[test] + fn text_dialect_markup_inside_a_fenced_code_block_is_never_recovered() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("recovery-test"), ()); + let mut response = ModelResponse::assistant( + "Here is the format:\n```\n{\"name\": \"shell\", \"arguments\": {}}\n```\n", + ); + + recover_text_dialect_calls( + &ctx, + &mut response, + &CallId::new("model-1"), + &offered(&["shell"]), + ); + + assert!( + response.message.tool_calls.is_empty(), + "markup quoted inside a fenced code block must not become a real call" + ); + assert!(response.text().contains("")); + } + + /// Sanity check for the fenced-code-block guard: markup outside any fence + /// is still recovered when the policy and tool offer both allow it. + #[test] + fn text_dialect_markup_outside_a_fenced_code_block_is_recovered() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("recovery-test"), ()); + let mut response = ModelResponse::assistant( + r#"{"name": "shell", "arguments": {"command": "id"}}"#, + ); + + recover_text_dialect_calls( + &ctx, + &mut response, + &CallId::new("model-1"), + &offered(&["shell"]), + ); + + assert_eq!(response.message.tool_calls.len(), 1); + assert_eq!(response.message.tool_calls[0].name, "shell"); + } +} diff --git a/crates/tinyagents-harness/src/agent_loop/run_queue_test.rs b/crates/tinyagents-harness/src/agent_loop/run_queue_test.rs new file mode 100644 index 00000000..e3cb36d0 --- /dev/null +++ b/crates/tinyagents-harness/src/agent_loop/run_queue_test.rs @@ -0,0 +1,502 @@ +//! Tests for the `RunQueue` wiring (A4): the loop drains the `Steer` lane at +//! the turn boundary after a tool batch, the `Followup` lane when it would +//! otherwise finish, honors `RunPolicy::queue_mode`, delivers the `Collect` +//! lane on `AgentRun::collected`, and emits `QueuedMessageApplied`. + +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::json; + +use crate::context::{RunConfig, RunContext}; +use crate::events::AgentEvent; +use crate::run_queue::{QueueLane, QueueMode, RunQueue, RunQueueHandle}; +use crate::runtime::{AgentHarness, RunPolicy}; +use crate::testkit::{EventRecorder, ScriptedModel}; +use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; +use tinyinference_llm::model::{ModelRequest, ModelResponse}; +use tinyinference_llm::tool::ToolCall; +use tinyinference_llm::usage::Usage; +use tinytools::{Tool, ToolResult}; + +// ── Helpers ───────────────────────────────────────────────────────────────── + +/// A tool that returns a fixed reply and, when given a queue, pushes a steer +/// message onto it *while executing* — i.e. mid-batch. +struct QueueingTool { + name: &'static str, + reply: &'static str, + push_on_execute: Option<(RunQueueHandle, &'static str)>, +} + +impl QueueingTool { + fn plain(name: &'static str, reply: &'static str) -> Arc { + Arc::new(Self { + name, + reply, + push_on_execute: None, + }) + } + + fn steering( + name: &'static str, + reply: &'static str, + queue: RunQueueHandle, + steer: &'static str, + ) -> Arc { + Arc::new(Self { + name, + reply, + push_on_execute: Some((queue, steer)), + }) + } +} + +#[async_trait] +impl Tool for QueueingTool { + fn name(&self) -> &str { + self.name + } + fn description(&self) -> &str { + "queueing tool" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object"}) + } + async fn execute(&self, _arguments: serde_json::Value) -> anyhow::Result { + if let Some((queue, steer)) = &self.push_on_execute { + queue.push(QueueLane::Steer, Message::user(*steer)).await; + } + Ok(ToolResult::success(self.reply)) + } +} + +fn response(tool_calls: Vec, text: &str) -> ModelResponse { + let content = if text.is_empty() { + Vec::new() + } else { + vec![ContentBlock::Text(text.to_string())] + }; + ModelResponse { + message: AssistantMessage { + id: None, + content, + tool_calls, + usage: Some(Usage::new(1, 1)), + origin: None, + }, + usage: Some(Usage::new(1, 1)), + finish_reason: Some("stop".to_string()), + raw: None, + resolved_model: None, + continue_turn: None, + served_from_cache: false, + correlation: None, + resolved_route: None, + } +} + +fn tool_turn(calls: &[(&str, &str)]) -> ModelResponse { + response( + calls + .iter() + .map(|(id, name)| ToolCall::new(*id, *name, json!({}))) + .collect(), + "", + ) +} + +fn final_turn(text: &str) -> ModelResponse { + response(Vec::new(), text) +} + +/// Texts of the user messages in `request`, in order. +fn user_texts(request: &ModelRequest) -> Vec { + request + .messages + .iter() + .filter(|message| matches!(message, Message::User(_))) + .map(Message::text) + .collect() +} + +/// A compact role/text rendering of a transcript for ordering assertions. +fn shape(messages: &[Message]) -> Vec { + messages + .iter() + .map(|message| match message { + Message::System(_) => format!("system:{}", message.text()), + Message::User(_) => format!("user:{}", message.text()), + Message::Assistant(a) if !a.tool_calls.is_empty() => { + format!("assistant:tools[{}]", a.tool_calls.len()) + } + Message::Assistant(_) => format!("assistant:{}", message.text()), + Message::Tool(t) => format!("tool:{}", t.tool_call_id), + Message::Custom(_) => "custom".to_string(), + }) + .collect() +} + +fn queued_applied(events: &[AgentEvent]) -> Vec<(QueueLane, usize)> { + events + .iter() + .filter_map(|event| match event { + AgentEvent::QueuedMessageApplied { lane, count } => Some((*lane, *count)), + _ => None, + }) + .collect() +} + +struct Fixture { + harness: AgentHarness<()>, + model: Arc, + queue: RunQueueHandle, + recorder: EventRecorder, +} + +fn fixture(responses: Vec, mode: QueueMode) -> Fixture { + let model = Arc::new(ScriptedModel::new(responses)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::clone(&model) as _); + harness.with_policy(RunPolicy { + queue_mode: mode, + ..RunPolicy::default() + }); + Fixture { + harness, + model, + queue: Arc::new(RunQueue::new()), + recorder: EventRecorder::new(), + } +} + +impl Fixture { + fn ctx(&self, run_id: &str) -> RunContext<()> { + RunContext::new(RunConfig::new(run_id), ()) + .with_events(self.recorder.sink()) + .with_run_queue(Arc::clone(&self.queue)) + } +} + +// ── Steer ─────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn steer_queued_mid_tool_batch_waits_for_the_batch_to_finish() { + let mut fx = fixture( + vec![ + tool_turn(&[("call-a", "a"), ("call-b", "b")]), + final_turn("done"), + ], + QueueMode::All, + ); + // Tool `a` pushes the steer while the two-call batch is executing. + fx.harness.register_tool(QueueingTool::steering( + "a", + "a-result", + Arc::clone(&fx.queue), + "steer: be brief", + )); + fx.harness + .register_tool(QueueingTool::plain("b", "b-result")); + + let ctx = fx.ctx("steer-mid-batch"); + let run = fx + .harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let requests = fx.model.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + user_texts(&requests[0]), + vec!["go"], + "the tool-call turn never sees the steer" + ); + assert_eq!( + user_texts(&requests[1]), + vec!["go", "steer: be brief"], + "the next model call sees the steer as a user message" + ); + // Both tool results land before the steer: it never interrupted the batch. + assert_eq!( + shape(&run.messages), + vec![ + "user:go", + "assistant:tools[2]", + "tool:call-a", + "tool:call-b", + "user:steer: be brief", + "assistant:done", + ] + ); + assert_eq!(run.text().as_deref(), Some("done")); + assert_eq!( + queued_applied(&fx.recorder.events()), + vec![(QueueLane::Steer, 1)] + ); + assert_eq!(fx.queue.status().await.total, 0); +} + +#[tokio::test] +async fn one_at_a_time_applies_one_steer_per_boundary_while_all_applies_every_steer_at_once() { + for (mode, expected_second, expected_third, expected_events) in [ + ( + QueueMode::OneAtATime, + vec!["go", "s1"], + vec!["go", "s1", "s2"], + vec![(QueueLane::Steer, 1), (QueueLane::Steer, 1)], + ), + ( + QueueMode::All, + vec!["go", "s1", "s2"], + vec!["go", "s1", "s2"], + vec![(QueueLane::Steer, 2)], + ), + ] { + let mut fx = fixture( + vec![ + tool_turn(&[("call-1", "a")]), + tool_turn(&[("call-2", "a")]), + final_turn("done"), + ], + mode, + ); + fx.harness + .register_tool(QueueingTool::plain("a", "a-result")); + fx.queue.push(QueueLane::Steer, Message::user("s1")).await; + fx.queue.push(QueueLane::Steer, Message::user("s2")).await; + + let ctx = fx.ctx("queue-mode"); + let run = fx + .harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let requests = fx.model.requests(); + assert_eq!(requests.len(), 3, "{mode:?}"); + assert_eq!(user_texts(&requests[0]), vec!["go"], "{mode:?}"); + assert_eq!(user_texts(&requests[1]), expected_second, "{mode:?}"); + assert_eq!(user_texts(&requests[2]), expected_third, "{mode:?}"); + assert_eq!(run.model_calls, 3, "{mode:?}"); + assert_eq!( + queued_applied(&fx.recorder.events()), + expected_events, + "{mode:?}" + ); + assert_eq!(fx.queue.status().await.total, 0, "{mode:?}"); + } +} + +// ── Followup ──────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn followup_runs_one_more_turn_after_the_model_would_have_finished() { + // Without a follow-up this script finishes after the first response. With + // one queued, the loop appends it as a user turn and keeps going until + // the follow-up turn (which itself uses a tool) reaches its own final. + let mut fx = fixture( + vec![ + final_turn("first answer"), + tool_turn(&[("call-1", "a")]), + final_turn("second answer"), + ], + QueueMode::All, + ); + fx.harness + .register_tool(QueueingTool::plain("a", "a-result")); + fx.queue + .push(QueueLane::Followup, Message::user("and then?")) + .await; + + let ctx = fx.ctx("followup"); + let run = fx + .harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!( + run.model_calls, 3, + "one extra turn (plus its tool round trip)" + ); + assert_eq!(run.text().as_deref(), Some("second answer")); + assert_eq!( + shape(&run.messages), + vec![ + "user:go", + "assistant:first answer", + "user:and then?", + "assistant:tools[1]", + "tool:call-1", + "assistant:second answer", + ] + ); + let requests = fx.model.requests(); + assert_eq!(user_texts(&requests[0]), vec!["go"]); + assert_eq!(user_texts(&requests[1]), vec!["go", "and then?"]); + assert_eq!( + queued_applied(&fx.recorder.events()), + vec![(QueueLane::Followup, 1)] + ); + assert!( + fx.recorder + .events() + .iter() + .filter(|event| matches!(event, AgentEvent::RunCompleted { .. })) + .count() + == 1, + "the run completes exactly once, after the follow-up turn" + ); +} + +#[tokio::test] +async fn no_queue_attached_finishes_exactly_as_before() { + let fx = fixture(vec![final_turn("done")], QueueMode::All); + let ctx = RunContext::new(RunConfig::new("no-queue"), ()).with_events(fx.recorder.sink()); + let run = fx + .harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + assert_eq!(run.model_calls, 1); + assert!(run.collected.is_empty()); + assert!(queued_applied(&fx.recorder.events()).is_empty()); +} + +// ── Collect ───────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn collect_lane_lands_on_the_run_and_never_reaches_the_model() { + let fx = fixture(vec![final_turn("done")], QueueMode::All); + fx.queue + .push(QueueLane::Collect, Message::user("observation 1")) + .await; + fx.queue + .push(QueueLane::Collect, Message::user("observation 2")) + .await; + + let ctx = fx.ctx("collect"); + let run = fx + .harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!( + run.collected.iter().map(Message::text).collect::>(), + vec!["observation 1", "observation 2"] + ); + assert_eq!(shape(&run.messages), vec!["user:go", "assistant:done"]); + for request in fx.model.requests() { + assert_eq!(user_texts(&request), vec!["go"]); + } + assert!( + queued_applied(&fx.recorder.events()).is_empty(), + "collected items are not applied to the transcript" + ); + assert_eq!(fx.queue.status().await.collects, 0); +} + +// ── Boundary semantics ────────────────────────────────────────────────────── + +#[tokio::test] +async fn steer_arriving_after_the_final_answer_is_applied_before_any_followup() { + // A steer that lands once the model has already answered is not lost: + // the natural-finish boundary applies it first (pi polls steering after + // every completed turn), and only a boundary with no pending steer + // takes a follow-up. Each gets its own turn, in that order. + let fx = fixture( + vec![ + final_turn("first answer"), + final_turn("steered answer"), + final_turn("followed-up answer"), + ], + QueueMode::All, + ); + fx.queue + .push(QueueLane::Followup, Message::user("follow-up")) + .await; + fx.queue + .push(QueueLane::Steer, Message::user("actually, shorter")) + .await; + + let ctx = fx.ctx("steer-at-finish"); + let run = fx + .harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!(run.model_calls, 3); + assert_eq!(run.text().as_deref(), Some("followed-up answer")); + assert_eq!( + shape(&run.messages), + vec![ + "user:go", + "assistant:first answer", + "user:actually, shorter", + "assistant:steered answer", + "user:follow-up", + "assistant:followed-up answer", + ] + ); + assert_eq!( + queued_applied(&fx.recorder.events()), + vec![(QueueLane::Steer, 1), (QueueLane::Followup, 1)] + ); + assert_eq!(fx.queue.status().await.total, 0); +} + +/// Requests `StopWithFinal` after any tool result. +struct StopAfterTool; + +#[async_trait] +impl crate::middleware::Middleware<(), ()> for StopAfterTool { + fn name(&self) -> &str { + "stop-after-tool" + } + async fn after_tool( + &self, + ctx: &mut RunContext<()>, + _state: &(), + _invocation: &crate::middleware::ToolInvocationIdentity, + _result: &mut ToolResult, + ) -> crate::error::Result<()> { + ctx.request_control(crate::context::MiddlewareControl::StopWithFinal( + "stopped by middleware".to_string(), + )); + Ok(()) + } +} + +#[tokio::test] +async fn middleware_stop_is_terminal_and_leaves_followups_queued() { + let mut fx = fixture( + vec![tool_turn(&[("call-1", "a")]), final_turn("never reached")], + QueueMode::All, + ); + fx.harness + .register_tool(QueueingTool::plain("a", "a-result")); + fx.harness.push_middleware(Arc::new(StopAfterTool)); + fx.queue + .push(QueueLane::Followup, Message::user("later")) + .await; + + let ctx = fx.ctx("middleware-stop"); + let run = fx + .harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!(run.model_calls, 1); + assert_eq!(run.text().as_deref(), Some("stopped by middleware")); + assert_eq!( + fx.queue.status().await.followups, + 1, + "a forced stop does not consume the follow-up; the host decides" + ); + assert!(queued_applied(&fx.recorder.events()).is_empty()); +} diff --git a/crates/tinyagents-harness/src/agent_loop/stream.rs b/crates/tinyagents-harness/src/agent_loop/stream.rs index 9535110a..cb31bbaf 100644 --- a/crates/tinyagents-harness/src/agent_loop/stream.rs +++ b/crates/tinyagents-harness/src/agent_loop/stream.rs @@ -31,11 +31,37 @@ use std::sync::Arc; use crate::context::{RunConfig, RunContext}; use crate::events::{EventListener, EventRecord, EventSink}; use crate::middleware::AgentRun; -use crate::runtime::AgentHarness; +use crate::runtime::{AgentHarness, InvocationRuntime}; use tinyinference_llm::message::Message; use super::PartialRunOutcome; +/// The concrete driver behind a caller-consumable stream: either the +/// durable harness borrowed for the caller's lifetime (the ordinary SDK +/// path), or an invocation-local runtime owned outright (the hosted path, +/// where the runtime is only alive as a local variable at the call site). +/// +/// Moving the `Owned` variant into the driving future (see +/// [`invoke_stream_with_runner`]) is what lets the hosted stream avoid both +/// an unsound lifetime extension and depending on field drop order: the +/// runtime's lifetime becomes exactly the future's, which the stream already +/// owns. +pub(crate) enum StreamRunner<'a, State: Send + Sync, Ctx: Send + Sync> { + Borrowed(&'a AgentHarness), + Owned(Arc>), +} + +impl std::ops::Deref for StreamRunner<'_, State, Ctx> { + type Target = AgentHarness; + + fn deref(&self) -> &AgentHarness { + match self { + StreamRunner::Borrowed(harness) => harness, + StreamRunner::Owned(runtime) => runtime.harness(), + } + } +} + /// One item yielded by [`AgentHarness::invoke_stream`]. /// /// The stream yields zero or more [`AgentStreamItem::Event`]s in emission @@ -43,6 +69,10 @@ use super::PartialRunOutcome; /// [`AgentStreamItem::Completed`] carrying the final [`AgentRun`], or /// [`AgentStreamItem::Failed`] carrying the error string. No items are produced /// after the terminal item. +// `Event` is by far the most common item — a run yields many events and +// exactly one terminal — so boxing `EventRecord` to shrink the rare `Failed` +// variant would add an allocation per streamed event for nothing. +#[allow(clippy::large_enum_variant)] #[derive(Clone, Debug)] pub enum AgentStreamItem { /// A live event emitted during the run. Carries the full @@ -156,110 +186,142 @@ impl AgentHarness { ctx: RunContext, input: Vec, ) -> impl futures::Stream + Send + 'a { - let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); - // Subscribe before driving so no event (starting with `RunStarted`) is - // missed. The listener rides the run's `EventSink`, which sub-agents - // clone, so their lifecycle events reach this stream too. - let listener: Arc = Arc::new(ChannelListener { tx }); - ctx.events.subscribe(listener.clone()); - let listener_guard = ChannelListenerGuard { - events: ctx.events.clone(), - listener, - }; + invoke_stream_with_runner(StreamRunner::Borrowed(self), state, ctx, input) + } +} + +/// Builds the caller-consumable event stream for either an ordinary +/// (borrowed-harness) or a hosted (owned-runtime) invocation. +/// +/// `runner` is moved into the driving future itself rather than dereferenced +/// up front, so an owned [`InvocationRuntime`] carried by `runner` lives +/// exactly as long as the future that needs it — no separate field, drop +/// order, or lifetime extension required on the caller's stream wrapper. +pub(crate) fn invoke_stream_with_runner<'a, State, Ctx>( + runner: StreamRunner<'a, State, Ctx>, + state: &'a State, + ctx: RunContext, + input: Vec, +) -> impl futures::Stream + Send + 'a +where + State: Send + Sync, + Ctx: Send + Sync + 'static, +{ + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + // Subscribe before driving so no event (starting with `RunStarted`) is + // missed. The listener rides the run's `EventSink`, which sub-agents + // clone, so their lifecycle events reach this stream too. + let listener: Arc = Arc::new(ChannelListener { tx }); + ctx.events.subscribe(listener.clone()); + let listener_guard = ChannelListenerGuard { + events: ctx.events.clone(), + listener, + }; - // Preserve partial work for a failed streamed run. The event stream - // remains unchanged, but terminal host capabilities need honest usage - // and executed-tool summaries for error and cancellation paths too. - let run_fut: Pin + Send + 'a>> = - Box::pin(self.invoke_streaming_in_context_collecting_partial(state, ctx, input)); + // Preserve partial work for a failed streamed run. The event stream + // remains unchanged, but terminal host capabilities need honest usage + // and executed-tool summaries for error and cancellation paths too. + // + // `runner` is moved into this async block rather than dereferenced + // beforehand: the generated state machine owns it (and, for the hosted + // `Owned` variant, the `Arc` inside it) for exactly as + // long as the future borrows from it across the `.await` below, which is + // what async/await's normal self-referential generator lowering makes + // sound without any unsafe code. + let run_fut: Pin + Send + 'a>> = + Box::pin(async move { + let runner = runner; + runner + .invoke_streaming_in_context_collecting_partial(state, ctx, input) + .await + }); - futures::stream::unfold( - ( + futures::stream::unfold( + ( + Phase::Running { + run_fut, + listener_guard, + }, + rx, + ), + |(phase, mut rx)| async move { + match phase { Phase::Running { - run_fut, + mut run_fut, listener_guard, - }, - rx, - ), - |(phase, mut rx)| async move { - match phase { - Phase::Running { - mut run_fut, - listener_guard, - } => { - tokio::select! { - biased; - // Prefer draining ready events so the consumer sees - // fine-grained progress rather than a late burst. - maybe = rx.recv() => match maybe { - Some(record) => { - Some(( - AgentStreamItem::Event(record), - ( - Phase::Running { - run_fut, - listener_guard, - }, - rx, - ), - )) - } - None => { - // All senders dropped (the run's context — - // and every sub-agent clone of the sink — - // is gone): the run is finishing. Await it - // for the terminal item. - let terminal = terminal_item(run_fut.await); + } => { + tokio::select! { + biased; + // Prefer draining ready events so the consumer sees + // fine-grained progress rather than a late burst. + maybe = rx.recv() => match maybe { + Some(record) => { + Some(( + AgentStreamItem::Event(record), + ( + Phase::Running { + run_fut, + listener_guard, + }, + rx, + ), + )) + } + None => { + // All senders dropped (the run's context — + // and every sub-agent clone of the sink — + // is gone): the run is finishing. Await it + // for the terminal item. + let terminal = terminal_item(run_fut.await); + drop(listener_guard); + Some((terminal, (Phase::Done, rx))) + } + }, + result = &mut run_fut => { + // The run finished. Events emitted during this + // final poll may still be buffered; drain them + // ahead of the terminal item. + let terminal = terminal_item(result); + match rx.try_recv() { + Ok(record) => Some(( + AgentStreamItem::Event(record), + ( + Phase::Draining { + terminal: Box::new(terminal), + listener_guard, + }, + rx, + ), + )), + Err(_) => { drop(listener_guard); Some((terminal, (Phase::Done, rx))) } - }, - result = &mut run_fut => { - // The run finished. Events emitted during this - // final poll may still be buffered; drain them - // ahead of the terminal item. - let terminal = terminal_item(result); - match rx.try_recv() { - Ok(record) => Some(( - AgentStreamItem::Event(record), - ( - Phase::Draining { - terminal: Box::new(terminal), - listener_guard, - }, - rx, - ), - )), - Err(_) => { - drop(listener_guard); - Some((terminal, (Phase::Done, rx))) - } - } } } } - Phase::Draining { - terminal, - listener_guard, - } => match rx.try_recv() { - Ok(record) => Some(( - AgentStreamItem::Event(record), - ( - Phase::Draining { - terminal, - listener_guard, - }, - rx, - ), - )), - Err(_) => { - drop(listener_guard); - Some((*terminal, (Phase::Done, rx))) - } - }, - Phase::Done => None, } - }, - ) - } + Phase::Draining { + terminal, + listener_guard, + } => match rx.try_recv() { + Ok(record) => Some(( + AgentStreamItem::Event(record), + ( + Phase::Draining { + terminal, + listener_guard, + }, + rx, + ), + )), + Err(_) => { + drop(listener_guard); + Some((*terminal, (Phase::Done, rx))) + } + }, + Phase::Done => None, + } + }, + ) } diff --git a/crates/tinyagents-harness/src/agent_loop/test.rs b/crates/tinyagents-harness/src/agent_loop/test.rs index 32dbc104..4bc591b0 100644 --- a/crates/tinyagents-harness/src/agent_loop/test.rs +++ b/crates/tinyagents-harness/src/agent_loop/test.rs @@ -28,7 +28,7 @@ use crate::tool::ToolTimeoutSettings; use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message, MessageDelta}; use tinyinference_llm::model::{ CapabilitySet, ChatModel, ModelProfile, ModelRequest, ModelResponse, ModelStreamItem, - ResponseFormat, ToolChoice, + ReasoningConfig, ReasoningEffort, ResponseFormat, SchemaTransform, StructuredMode, ToolChoice, }; use tinyinference_llm::providers::MockModel; use tinyinference_llm::tool::{ToolCall, ToolSchema}; @@ -298,6 +298,69 @@ impl Tool for StrictLookupTool { } } +/// A [`crate::tool::toolset::ToolSet`] whose live tool set changes on its +/// second call — used to prove `agent_loop::tool_changes`'s wiring actually +/// fires on a genuine mid-run toolset change (B6). +struct DynamicToolSet { + calls: std::sync::atomic::AtomicUsize, + search: Arc, + browse: Arc, +} + +#[async_trait] +impl crate::tool::toolset::ToolSet<(), ()> for DynamicToolSet { + async fn tools(&self, _ctx: &RunContext<()>) -> Result>> { + let call = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if call == 0 { + Ok(vec![self.search.clone()]) + } else { + Ok(vec![self.search.clone(), self.browse.clone()]) + } + } + + async fn call( + &self, + name: &str, + args: serde_json::Value, + _ctx: &RunContext<()>, + ) -> Result { + let tool = if name == self.search.name() { + &self.search + } else if name == self.browse.name() { + &self.browse + } else { + return Err(TinyAgentsError::ToolNotFound(name.to_string())); + }; + tool.execute(args) + .await + .map_err(|err| TinyAgentsError::Tool(err.to_string())) + } +} + +/// Wraps [`MockModel`] to advertise a caller-supplied [`ModelProfile`] +/// instead of the fixed permissive one `MockModel::profile` returns — used to +/// exercise the `mid_conversation_system_messages = true` insert path (B6) +/// end to end, which no built-in test provider otherwise advertises. +struct ProfiledModel { + inner: MockModel, + profile: ModelProfile, +} + +#[async_trait] +impl ChatModel<()> for ProfiledModel { + fn profile(&self) -> Option<&ModelProfile> { + Some(&self.profile) + } + + async fn invoke( + &self, + state: &(), + request: ModelRequest, + ) -> tinyinference_llm::Result { + >::invoke(&self.inner, state, request).await + } +} + /// Builds a tool-call assistant response (no text, one tool call). fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { ModelResponse { @@ -306,6 +369,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod content: Vec::new(), tool_calls: vec![ToolCall::new(id, name, arguments)], usage: Some(Usage::new(7, 3)), + origin: None, }, usage: Some(Usage::new(7, 3)), finish_reason: Some("tool_calls".to_string()), @@ -333,6 +397,7 @@ fn invalid_tool_call_response(id: &str, name: &str, raw: &str) -> ModelResponse content: Vec::new(), tool_calls: vec![ToolCall::invalid(id, name, raw, reason)], usage: Some(Usage::new(7, 3)), + origin: None, }, usage: Some(Usage::new(7, 3)), finish_reason: Some("tool_calls".to_string()), @@ -353,6 +418,7 @@ fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { content: vec![ContentBlock::Text(text.to_string())], tool_calls: Vec::new(), usage: Some(Usage::new(input, output)), + origin: None, }, usage: Some(Usage::new(input, output)), finish_reason: Some("stop".to_string()), @@ -377,6 +443,7 @@ fn truncated_empty_response(reasoning_tokens: u64) -> ModelResponse { content: Vec::new(), tool_calls: Vec::new(), usage: Some(Usage::new(4, reasoning_tokens)), + origin: None, }, usage: Some(Usage::new(4, reasoning_tokens)), finish_reason: Some("length".to_string()), @@ -1108,6 +1175,293 @@ async fn model_requests_tool_then_finishes() { assert_eq!(run.messages[2].text(), "tool-output"); } +/// B6: a toolset chain whose live set changes mid-run, against a model whose +/// profile does *not* advertise `mid_conversation_system_messages` (the +/// default `MockModel` profile), gets the delta **folded** into the leading +/// system message rather than appended as a new one — no new +/// [`Message::System`] appears anywhere in the transcript, but the delta is +/// still fully recorded on the leading message. +#[tokio::test] +async fn dynamic_toolset_change_folds_into_the_leading_system_message_by_default() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("call-1", "search", json!({"q": "x"})), + text_response("done", 4, 2), + ])), + ); + let search: Arc = Arc::new(FakeTool::new("search", "search-output")); + let browse: Arc = Arc::new(FakeTool::new("browse", "browse-output")); + let toolset = Arc::new(DynamicToolSet { + calls: std::sync::atomic::AtomicUsize::new(0), + search: search.clone(), + browse: browse.clone(), + }); + harness.with_toolset(toolset.clone()); + // Only `search` needs to be dispatchable (the script's only call); a + // bridge for `browse` too would register it in `self.tools` statically + // from turn one, defeating the point of this test (the toolset chain + // alone is what makes `browse` come and go). + harness.register_tool_dispatch(Arc::new(crate::tool::toolset::ToolSetDispatchBridge::new( + toolset, search, + ))); + let _ = browse; + + let run = harness + .invoke_default( + &(), + vec![Message::system("baseline persona"), Message::user("go")], + ) + .await + .expect("run succeeds"); + + let system_messages: Vec<&Message> = run + .messages + .iter() + .filter(|message| matches!(message, Message::System(_))) + .collect(); + assert_eq!( + system_messages.len(), + 1, + "no new system message was appended" + ); + let Message::System(leading) = system_messages[0] else { + unreachable!("filtered above"); + }; + // Turn 1's diff runs against an as-yet-undeclared transcript, so it + // records the whole live set (both `search` and `browse`) in one fold — + // not just the later delta — which is what lets a replay reconstruct the + // complete effective tool set from the transcript alone. + assert_eq!( + leading + .tools_added + .iter() + .map(|schema| schema.name.as_str()) + .collect::>(), + vec!["browse", "search"] + ); +} + +/// B6 (insert path): the same toolset-change scenario against a model whose +/// profile *does* advertise `mid_conversation_system_messages` gets the delta +/// appended as exactly one new [`Message::System`] patch instead. +#[tokio::test] +async fn dynamic_toolset_change_appends_exactly_one_patch_when_the_profile_allows_it() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + let profile = ModelProfile { + tool_calling: true, + mid_conversation_system_messages: true, + ..ModelProfile::default() + }; + harness.register_model( + "mock", + Arc::new(ProfiledModel { + inner: MockModel::with_responses(vec![ + tool_call_response("call-1", "search", json!({"q": "x"})), + text_response("done", 4, 2), + ]), + profile, + }), + ); + let search: Arc = Arc::new(FakeTool::new("search", "search-output")); + let browse: Arc = Arc::new(FakeTool::new("browse", "browse-output")); + let toolset = Arc::new(DynamicToolSet { + calls: std::sync::atomic::AtomicUsize::new(0), + search: search.clone(), + browse: browse.clone(), + }); + harness.with_toolset(toolset.clone()); + harness.register_tool_dispatch(Arc::new(crate::tool::toolset::ToolSetDispatchBridge::new( + toolset, search, + ))); + let _ = browse; + + let run = harness + .invoke_default( + &(), + vec![Message::system("baseline persona"), Message::user("go")], + ) + .await + .expect("run succeeds"); + + let system_messages: Vec<&Message> = run + .messages + .iter() + .filter(|message| matches!(message, Message::System(_))) + .collect(); + // The original leading system message plus exactly one appended patch. + assert_eq!(system_messages.len(), 2, "exactly one patch was appended"); + let Message::System(patch) = system_messages[1] else { + unreachable!("filtered above"); + }; + // As in the fold-path test above, turn 1's patch declares the whole live + // set, not just a later delta. + assert_eq!( + patch + .tools_added + .iter() + .map(|schema| schema.name.as_str()) + .collect::>(), + vec!["browse", "search"] + ); + + // The reconstructed effective tool set matches what was actually offered. + let (_, effective_tools) = tinyinference_llm::message::replay_system_state(&run.messages); + let mut names: Vec<&str> = effective_tools + .iter() + .map(|schema| schema.name.as_str()) + .collect(); + names.sort(); + assert_eq!(names, vec!["browse", "search"]); +} + +/// Gap G3: a `defer_loading` capability's tool is not advertised until the +/// model calls `load_capability`, and once it does, the very next turn's +/// existing tool-change diff (B6, `agent_loop::tool_changes`) picks up the +/// change automatically and appends a patch system message — no bespoke +/// capability-specific patch wiring is needed. +#[tokio::test] +async fn defer_loading_capability_is_exposed_only_after_load_capability_and_patches_the_transcript() +{ + let profile = ModelProfile { + tool_calling: true, + mid_conversation_system_messages: true, + ..ModelProfile::default() + }; + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(ProfiledModel { + inner: MockModel::with_responses(vec![ + tool_call_response( + "call-1", + crate::capability::LOAD_CAPABILITY_TOOL_NAME, + json!({"capability": "advanced"}), + ), + // Turn 2 makes no tool call — the point of this test is the + // *advertisement* change the loop's existing tool-change diff + // (B6) picks up before this turn's request goes out, not + // `advanced-tool`'s own dispatch (a deferred capability's + // tools are advertised automatically but, like any + // `with_toolset` toolset, need an explicit + // `ToolSetDispatchBridge` to also be *callable* — see that + // type's doc comment; orthogonal to what this test covers). + text_response("done", 4, 2), + ]), + profile, + }), + ); + + // A minimal single-tool toolset behind the capability, independent of + // `defer_loading` gating (that gating is `CapabilityToolSet`'s job, one + // layer up). + struct SingleToolSet { + tool: Arc, + } + #[async_trait] + impl crate::tool::toolset::ToolSet<(), ()> for SingleToolSet { + async fn tools(&self, _ctx: &RunContext<()>) -> Result>> { + Ok(vec![self.tool.clone()]) + } + async fn call( + &self, + name: &str, + args: serde_json::Value, + _ctx: &RunContext<()>, + ) -> Result { + if name == self.tool.name() { + self.tool + .execute(args) + .await + .map_err(|err| TinyAgentsError::Tool(err.to_string())) + } else { + Err(TinyAgentsError::ToolNotFound(name.to_string())) + } + } + } + + let capability = crate::capability::Capability::new("advanced") + .with_instructions("Advanced instructions.") + .with_toolset(Arc::new(SingleToolSet { + tool: Arc::new(FakeTool::new("advanced-tool", "advanced-output")), + })) + .with_defer_loading(true); + harness.with_capability(capability); + + let run = harness + .invoke_default( + &(), + vec![Message::system("baseline persona"), Message::user("go")], + ) + .await + .expect("run succeeds"); + + assert_eq!(run.text(), Some("done".to_string())); + + let system_messages: Vec<&Message> = run + .messages + .iter() + .filter(|message| matches!(message, Message::System(_))) + .collect(); + // The original leading system message, plus one patch for turn 1 (just + // `load_capability` itself — `advanced-tool` is still gated), plus one + // patch for turn 2 (once `load_capability` ran, `advanced-tool` joins the + // live set). + assert_eq!( + system_messages.len(), + 3, + "expected the leading message plus two tool-change patches" + ); + + let Message::System(turn1_patch) = system_messages[1] else { + unreachable!("filtered above"); + }; + let turn1_added: Vec<&str> = turn1_patch + .tools_added + .iter() + .map(|schema| schema.name.as_str()) + .collect(); + assert_eq!( + turn1_added, + vec![crate::capability::LOAD_CAPABILITY_TOOL_NAME] + ); + assert!( + !turn1_added.contains(&"advanced-tool"), + "the deferred capability's tool must not be advertised before load_capability runs" + ); + + let Message::System(turn2_patch) = system_messages[2] else { + unreachable!("filtered above"); + }; + assert_eq!( + turn2_patch + .tools_added + .iter() + .map(|schema| schema.name.as_str()) + .collect::>(), + vec!["advanced-tool"], + "advanced-tool becomes advertised only on the turn after load_capability ran" + ); + + // The final effective tool set (replayed from the transcript alone) + // includes both the always-registered `load_capability` and the now + // loaded `advanced-tool`. + let (_, effective_tools) = tinyinference_llm::message::replay_system_state(&run.messages); + let mut names: Vec<&str> = effective_tools + .iter() + .map(|schema| schema.name.as_str()) + .collect(); + names.sort(); + assert_eq!( + names, + vec![ + "advanced-tool", + crate::capability::LOAD_CAPABILITY_TOOL_NAME + ] + ); +} + #[tokio::test] async fn after_tool_receives_distinct_identity_for_same_named_calls() { let mut parallel_calls = ModelResponse::assistant(""); @@ -1995,6 +2349,52 @@ async fn malformed_tool_arguments_recover_as_error_tool_result() { injected, "an error tool result should be injected into the transcript" ); +} + +/// I-13 regression: provider-invalid arguments that `relaxed_json` can +/// actually repair (unquoted object keys, here) must be recovered and the +/// call executed — not turned into a "fix your JSON" round trip the model +/// often cannot act on. Before the fix, admission short-circuited straight +/// to the tool-error path without ever trying `recover_relaxed_object`, +/// even though that module exists specifically for this input shape. +#[tokio::test] +async fn provider_invalid_arguments_recoverable_by_relaxed_json_are_repaired_and_executed() { + use crate::testkit::EventRecorder; + + let tool = Arc::new(crate::testkit::FakeTool::returning("lookup", "found it")); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + // Unquoted object key: `relaxed_json::recover_relaxed_object` + // repairs this to `{"query":"weather"}`. + invalid_tool_call_response("call-x", "lookup", "{query:\"weather\"}"), + text_response("found it", 1, 1), + ])), + ); + harness.register_tool(tool.clone()); + + let recorder = EventRecorder::new(); + let ctx = + RunContext::new(RunConfig::new("relaxed-json-repair"), ()).with_events(recorder.sink()); + let run = harness + .invoke_in_context(&(), ctx, vec![Message::user("lookup the weather")]) + .await + .expect("repaired arguments let the call execute"); + + assert_eq!(run.text().as_deref(), Some("found it")); + assert_eq!( + tool.calls(), + vec![json!({"query": "weather"})], + "the tool must receive the repaired, strict-JSON arguments" + ); + assert!( + recorder.events().iter().any(|event| matches!( + event, + AgentEvent::InvalidToolArgs { recovery, .. } if recovery == "repaired" + )), + "the repair must be observable as InvalidToolArgs{{ recovery: \"repaired\" }}" + ); // The recovery is surfaced as an `InvalidToolArgs` event. assert!( recorder @@ -2400,6 +2800,40 @@ async fn run_limits_max_retries_per_call_caps_a_looser_retry_policy() { assert_eq!(*failing.attempts.lock().unwrap(), 2); } +#[tokio::test] +async fn retry_middleware_and_run_policy_retry_do_not_multiply_attempts() { + // Regression test (I-7): `RetryMiddleware::wrap_model` retries the whole + // wrap onion, and `invoke_model_resolving` (the loop's own base call) had + // its own independent retry loop; with both configured the worst case was + // `mw.max_attempts x policy.retry.max_attempts` provider calls for one + // logical failure. A registered `RetryMiddleware` must make the base call + // skip its own retry loop, so the total attempt count is bounded by the + // middleware's `max_attempts` alone. + let mut harness: AgentHarness<()> = AgentHarness::new(); + let failing = Arc::new(FailingModel { + attempts: Mutex::new(0), + }); + harness.register_model("primary", failing.clone()); + // The middleware allows 3 attempts; the loop's own retry (if it fired + // too) would allow another 5 — 15 total if the two layers multiplied. + harness.push_model_middleware(Arc::new(crate::middleware::library::RetryMiddleware::new( + RetryPolicy::default().with_max_attempts(3), + ))); + harness.with_policy(RunPolicy { + retry: RetryPolicy::default().with_max_attempts(5), + ..RunPolicy::default() + }); + + let err = harness + .invoke_default(&(), vec![Message::user("hi")]) + .await + .expect_err("FailingModel never succeeds"); + assert!(matches!(err, TinyAgentsError::Model(_)), "got {err:?}"); + + // Bounded by the middleware's max_attempts (3), not 3 x 5. + assert_eq!(*failing.attempts.lock().unwrap(), 3); +} + #[tokio::test] async fn provider_error_401_is_not_retried() { // Regression test: before `ProviderError` was preserved structurally, a @@ -2662,6 +3096,50 @@ async fn runtime_fallback_skips_capability_ineligible_candidate() { ); } +/// I-2 end-to-end regression: under the default `Auto` text-dialect recovery +/// policy, a model whose resolved profile reports native tool calling must +/// never have `` markup it merely quotes — here, inside a fenced +/// code block explaining the format — executed as a real tool call. Before +/// the fix, `recover_text_dialect_calls` ran unconditionally whenever the +/// request offered tools and the provider returned no native calls, +/// regardless of the model's own advertised capabilities. +#[tokio::test] +async fn native_tool_calling_model_does_not_execute_quoted_text_dialect_markup() { + let tool = Arc::new(FakeTool::new("shell", "must not run")); + let model = Arc::new(ProfiledTextModel { + profile: ModelProfile { + tool_calling: true, + ..ModelProfile::default() + }, + text: "Here is the tool-call format for reference:\n\ + ```\n\ + {\"name\": \"shell\", \"arguments\": {\"command\": \"id\"}}\n\ + ```\n", + attempts: Mutex::new(0), + }); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("native", model.clone()); + harness.register_tool(tool.clone()); + + let run = harness + .invoke_default(&(), vec![Message::user("how do tool calls work?")]) + .await + .expect("run succeeds with a plain text final answer"); + + assert_eq!( + *tool.calls.lock().unwrap(), + 0, + "the quoted call must not run" + ); + assert!(run.text().unwrap_or_default().contains("")); + assert_eq!( + *model.attempts.lock().unwrap(), + 1, + "no retry/fallback needed" + ); +} + #[tokio::test] async fn invoke_with_status_reports_completed() { use crate::ids::{ExecutionStatus, HarnessPhase}; @@ -2900,6 +3378,86 @@ async fn streaming_delta_transform_controls_final_run_and_cached_response() { ); } +/// C-2 regression: a streaming turn whose terminal response carries a signed +/// `Thinking` block ahead of a tool call must keep that exact signature in +/// `run.messages`. Anthropic requires the signed thinking block to precede a +/// `tool_use` block verbatim on replay; synthesizing a fresh, unsigned block +/// from the streamed reasoning text (the old behavior) breaks that replay on +/// the very next model call. No delta middleware is registered here, so the +/// streamed reasoning text is identical to the terminal block's text and the +/// fix's "keep it verbatim" branch is exercised. +#[tokio::test] +async fn streaming_turn_keeps_a_signed_thinking_signature_ahead_of_a_tool_call() { + use crate::testkit::StreamingMock; + + let tool = Arc::new(FakeTool::new("lookup", "ok")); + let mut terminal = ModelResponse::assistant(""); + terminal.message.content = vec![tinyinference_llm::message::ContentBlock::Thinking { + text: "let me think".to_string(), + signature: Some("sig-123".to_string()), + }]; + terminal + .message + .tool_calls + .push(ToolCall::new("call-1", "lookup", json!({}))); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "stream", + Arc::new(StreamingMock::new(vec![ + ModelStreamItem::Started, + ModelStreamItem::MessageDelta(MessageDelta::reasoning("let me think")), + ModelStreamItem::ToolCallDelta(tinyinference_llm::tool::ToolDelta { + call_id: "call-1".to_string(), + content: "{}".to_string(), + tool_name: Some("lookup".to_string()), + ..Default::default() + }), + ModelStreamItem::Completed(terminal), + ])), + ); + harness.register_tool(tool.clone()); + + // Cap the run at one model call: the mock always replays the same + // scripted tool call, so a second turn would just repeat it forever. + // Only the first turn's assistant message (the one under test) is + // needed. + let ctx = RunContext::new( + RunConfig::new("thinking-signature").with_max_model_calls(1), + (), + ); + let outcome = harness + .invoke_streaming_in_context_collecting_partial(&(), ctx, vec![Message::user("go")]) + .await; + + let thinking_blocks: Vec<_> = outcome + .run + .messages + .iter() + .filter_map(|message| match message { + tinyinference_llm::message::Message::Assistant(assistant) => { + Some(assistant.content.iter()) + } + _ => None, + }) + .flatten() + .filter(|block| { + matches!( + block, + tinyinference_llm::message::ContentBlock::Thinking { .. } + ) + }) + .collect(); + assert_eq!( + thinking_blocks, + vec![&tinyinference_llm::message::ContentBlock::Thinking { + text: "let me think".to_string(), + signature: Some("sig-123".to_string()), + }], + "the terminal Thinking block's signature must survive into run.messages verbatim" + ); +} + #[tokio::test] async fn streaming_middleware_can_suppress_a_standalone_tool_delta() { use crate::testkit::StreamingMock; @@ -2919,6 +3477,7 @@ async fn streaming_middleware_can_suppress_a_standalone_tool_delta() { call_id: "blocked-call".to_string(), content: "{}".to_string(), tool_name: Some("blocked".to_string()), + ..Default::default() }), ModelStreamItem::Completed(terminal), ])), @@ -2964,6 +3523,7 @@ async fn streaming_tool_delta_transform_controls_terminal_dispatch() { call_id: "raw-call".to_string(), content: r#"{"raw":true}"#.to_string(), tool_name: Some("blocked".to_string()), + ..Default::default() }), ModelStreamItem::Completed(terminal), ])), @@ -3343,6 +3903,9 @@ async fn per_model_call_ceiling_times_out_a_slow_call_with_run_time_left() { // ceiling (20ms) is tighter than the model's 200ms sleep, so the ceiling // interrupts the call — and the error must name the ceiling, not the run's // remaining budget, so triage can tell a wedged call from an exhausted run. + // A per-call ceiling is a `CallTimeout`, not a `Timeout`: it is retryable + // and must not skip the fallback chain the way a run-deadline timeout + // does (I-1). One retry attempt is enough to prove that here. let mut harness: AgentHarness<()> = AgentHarness::new(); harness.register_model( "slow", @@ -3350,6 +3913,9 @@ async fn per_model_call_ceiling_times_out_a_slow_call_with_run_time_left() { ); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_call_ms(Some(20)), + retry: RetryPolicy::default() + .with_max_attempts(1) + .with_backoff_sleep(false), ..RunPolicy::default() }); @@ -3360,41 +3926,88 @@ async fn per_model_call_ceiling_times_out_a_slow_call_with_run_time_left() { .expect_err("a call slower than the per-call ceiling must time out"); match &err { - TinyAgentsError::Timeout(msg) => { + TinyAgentsError::CallTimeout(msg) => { assert!(msg.contains("per-model-call ceiling"), "{msg}"); } - other => panic!("expected Timeout, got {other:?}"), + other => panic!("expected CallTimeout, got {other:?}"), } } #[tokio::test] -async fn per_model_call_ceiling_bounds_calls_without_any_run_deadline() { +async fn per_model_call_ceiling_consults_the_fallback_chain_instead_of_aborting() { use std::time::Duration; - use crate::testkit::SlowModel; - - // With no run timeout and no policy wall clock, a model call used to be - // awaited unbounded. The per-call ceiling alone must bound it. + use crate::testkit::{ScriptedModel, SlowModel}; + + // Same setup as `per_model_call_ceiling_times_out_a_slow_call_with_run_time_left`, + // but with a fallback model registered. Before the fix, the per-call + // ceiling produced a plain `Timeout`, which the fallback gate in + // `invoke_model_resolving` treats as terminal ("the run itself is out of + // wall-clock budget") and returns immediately — the fallback model is + // never even consulted, let alone called. With the fix, a `CallTimeout` + // falls through to the fallback walk, so the run succeeds on the + // fallback model instead of failing. + let slow = Arc::new(SlowModel::new(Duration::from_millis(200), "too late")); + let fallback = Arc::new(ScriptedModel::replies(vec!["fallback answer"])); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model( - "slow", - Arc::new(SlowModel::new(Duration::from_millis(200), "too late")), - ); + harness.register_model("slow", slow.clone()); + harness.register_model("fallback", fallback.clone()); harness.with_policy(RunPolicy { limits: RunLimits::default().with_max_model_call_ms(Some(20)), + retry: RetryPolicy::default() + .with_max_attempts(1) + .with_backoff_sleep(false), + fallback: Some(FallbackPolicy { + models: vec!["slow".to_string(), "fallback".to_string()], + }), ..RunPolicy::default() }); - let err = harness + let config = RunConfig::new("per-call-cap-fallback").with_timeout_ms(60_000); + let run = harness + .invoke(&(), (), config, vec![Message::user("hi")]) + .await + .expect("a retryable CallTimeout must fall back instead of aborting the run"); + + assert_eq!(run.text().as_deref(), Some("fallback answer")); + assert_eq!( + fallback.requests().len(), + 1, + "the fallback chain must actually have been consulted and called" + ); +} + +#[tokio::test] +async fn per_model_call_ceiling_bounds_calls_without_any_run_deadline() { + use std::time::Duration; + + use crate::testkit::SlowModel; + + // With no run timeout and no policy wall clock, a model call used to be + // awaited unbounded. The per-call ceiling alone must bound it. + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "slow", + Arc::new(SlowModel::new(Duration::from_millis(200), "too late")), + ); + harness.with_policy(RunPolicy { + limits: RunLimits::default().with_max_model_call_ms(Some(20)), + retry: RetryPolicy::default() + .with_max_attempts(1) + .with_backoff_sleep(false), + ..RunPolicy::default() + }); + + let err = harness .invoke_default(&(), vec![Message::user("hi")]) .await .expect_err("the ceiling alone must bound an otherwise-unbounded call"); match &err { - TinyAgentsError::Timeout(msg) => { + TinyAgentsError::CallTimeout(msg) => { assert!(msg.contains("per-model-call ceiling"), "{msg}"); } - other => panic!("expected Timeout, got {other:?}"), + other => panic!("expected CallTimeout, got {other:?}"), } } @@ -3855,6 +4468,27 @@ async fn middleware_control_stops_loop_with_final_response() { assert_eq!(run.final_response.unwrap().text(), "stopped early"); // The tool was never executed because the loop stopped first. assert_eq!(run.tool_calls, 0); + + // M-1 regression: the assistant row still carries the `tool_calls` the + // model requested, but the loop must synthesize a tool result for each + // one so `run.messages` stays replayable (a provider rejects a transcript + // whose assistant `tool_calls` have no matching tool message). + // user, assistant(1 tool call), tool(synthetic). + assert_eq!(run.messages.len(), 3); + let Message::Assistant(assistant) = &run.messages[1] else { + panic!( + "expected assistant message at index 1, got {:?}", + run.messages[1] + ); + }; + assert_eq!(assistant.tool_calls.len(), 1); + let Message::Tool(tool_message) = &run.messages[2] else { + panic!( + "expected synthetic tool message at index 2, got {:?}", + run.messages[2] + ); + }; + assert_eq!(tool_message.tool_call_id, assistant.tool_calls[0].id); } /// Middleware that requests an interrupt after the first model response. @@ -4076,6 +4710,111 @@ impl Tool for ConcurrencyProbeTool { } } +/// A concurrency-safe tool that fails fast (a real dispatch error, not a +/// recoverable `ToolResult::error`), used to exercise the concurrent path's +/// first-fatal-error handling. +struct FailingConcurrentTool { + name: &'static str, +} + +#[async_trait] +impl Tool for FailingConcurrentTool { + fn name(&self) -> &str { + self.name + } + fn description(&self) -> &str { + "fails fast" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object"}) + } + fn is_concurrency_safe(&self, _arguments: &serde_json::Value) -> bool { + true + } + async fn execute(&self, _arguments: serde_json::Value) -> anyhow::Result { + Err(anyhow::anyhow!("boom")) + } +} + +/// C-3 regression: on the first fatal error in the concurrent tool path, +/// every already-started sibling call must still get exactly one terminal +/// event (`ToolFailed`), and `active_tool_calls` must end up empty — not just +/// the call that actually failed. Before the fix, siblings whose futures had +/// already resolved (via `join_all`) but were never reached by the fold after +/// the first `Err` kept their `ToolStarted` unanswered and stayed listed in +/// `active_tool_calls` even though the run had already failed. +#[tokio::test] +async fn concurrent_tool_failure_fails_every_started_sibling_before_returning() { + use crate::testkit::EventRecorder; + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![multi_tool_call_response( + // "boom" fails fast and comes first in call order, so the fold + // reaches its fatal error while "alpha" (slower, but already + // resolved by the time `join_all` returns) is still an + // unprocessed sibling — exactly the scenario the fix covers. + vec![("call-a", "boom"), ("call-b", "alpha")], + )])), + ); + harness.register_tool(Arc::new(ConcurrencyProbeTool { + name: "alpha", + reply: "alpha-out", + delay: std::time::Duration::from_millis(80), + active: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + max_seen: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + })); + harness.register_tool(Arc::new(FailingConcurrentTool { name: "boom" })); + + let recorder = EventRecorder::new(); + let ctx = RunContext::new(RunConfig::new("concurrent-fatal"), ()).with_events(recorder.sink()); + let outcome = harness + .invoke_in_context_collecting_partial(&(), ctx, vec![Message::user("go")]) + .await; + + assert!( + outcome.error.is_some(), + "a fatal sibling error must fail the turn" + ); + assert!( + outcome.status.active_tool_calls.is_empty(), + "every started call must have a terminal event before the run reports failure, \ + got active_tool_calls = {:?}", + outcome.status.active_tool_calls + ); + + let started: Vec<_> = recorder + .events() + .iter() + .filter_map(|record| match record { + AgentEvent::ToolStarted { call_id, .. } => Some(call_id.as_str().to_string()), + _ => None, + }) + .collect(); + let terminal: Vec<_> = recorder + .events() + .iter() + .filter_map(|record| match record { + AgentEvent::ToolFailed { call_id, .. } => Some(call_id.as_str().to_string()), + AgentEvent::ToolCompleted { call_id, .. } => Some(call_id.as_str().to_string()), + _ => None, + }) + .collect(); + assert_eq!(started.len(), 2, "both siblings must have started"); + assert_eq!( + terminal.len(), + 2, + "every started call must be answered by exactly one terminal event, got {terminal:?}" + ); + for call_id in &started { + assert!( + terminal.contains(call_id), + "call `{call_id}` started but has no terminal event" + ); + } +} + /// Builds an assistant response carrying several tool calls in one turn. fn multi_tool_call_response(calls: Vec<(&str, &str)>) -> ModelResponse { let tool_calls = calls @@ -4088,6 +4827,7 @@ fn multi_tool_call_response(calls: Vec<(&str, &str)>) -> ModelResponse { content: Vec::new(), tool_calls, usage: Some(Usage::new(7, 3)), + origin: None, }, usage: Some(Usage::new(7, 3)), finish_reason: Some("tool_calls".to_string()), @@ -4149,6 +4889,55 @@ async fn independent_tool_calls_in_one_turn_run_concurrently() { ); } +#[tokio::test] +async fn max_tool_concurrency_bounds_how_many_tools_run_at_once() { + // I-8 regression test: with 4 concurrency-safe tools requested in one + // turn and `RunLimits::max_tool_concurrency` set to 2, at most 2 may be + // in flight at once, even though all 4 are eligible for the concurrent + // path. `max_seen` is an atomic high-water mark, so any window where 3+ + // ran together would be caught regardless of scheduling order. + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + multi_tool_call_response(vec![ + ("call-a", "alpha"), + ("call-b", "beta"), + ("call-c", "gamma"), + ("call-d", "delta"), + ]), + text_response("done", 4, 2), + ])), + ); + let active = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let max_seen = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + for name in ["alpha", "beta", "gamma", "delta"] { + harness.register_tool(Arc::new(ConcurrencyProbeTool { + name, + reply: "out", + delay: std::time::Duration::from_millis(60), + active: active.clone(), + max_seen: max_seen.clone(), + })); + } + harness.with_policy(RunPolicy { + limits: RunLimits::default().with_max_tool_concurrency(Some(2)), + ..RunPolicy::default() + }); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!(run.tool_calls, 4); + assert_eq!( + max_seen.load(std::sync::atomic::Ordering::SeqCst), + 2, + "no more than max_tool_concurrency (2) tools should ever be in flight at once" + ); +} + #[tokio::test] async fn parallel_tool_results_keep_original_call_order_and_ids() { let mut harness: AgentHarness<()> = AgentHarness::new(); @@ -4902,3 +5691,1044 @@ async fn echo_unwrap_is_skipped_when_the_inner_value_is_still_invalid() { "the tool must not run with arguments that never validated" ); } + +// --------------------------------------------------------------------------- +// Follow-up F1: harness-side `ModelProfile` wiring +// --------------------------------------------------------------------------- + +/// A tool whose declared schema carries `$defs`, so a `SchemaTransform` that +/// strips them (or resolves refs) is observable on the wire. +struct DefsTool; + +#[async_trait] +impl Tool for DefsTool { + fn name(&self) -> &str { + "lookup" + } + fn description(&self) -> &str { + "look something up" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "$defs": {"Id": {"type": "string"}}, + "properties": {"id": {"$ref": "#/$defs/Id"}}, + }) + } + async fn execute(&self, _arguments: serde_json::Value) -> anyhow::Result { + Ok(ToolResult::success("ok")) + } +} + +#[tokio::test] +async fn resolved_profile_schema_transform_is_applied_to_tool_schemas() { + use crate::testkit::ScriptedModel; + + let model = Arc::new( + ScriptedModel::replies(vec!["done"]).with_profile(ModelProfile { + schema_transform: Some(SchemaTransform::StripDefs), + ..ModelProfile::default() + }), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", model.clone()); + harness.register_tool(Arc::new(DefsTool)); + + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let requests = model.requests(); + let request = requests.first().expect("one model call"); + let tool = request + .tools + .iter() + .find(|t| t.name == "lookup") + .expect("lookup tool advertised"); + assert!( + tool.parameters.get("$defs").is_none(), + "the resolved profile's StripDefs transform must strip `$defs` before \ + the request is sent: {:?}", + tool.parameters + ); +} + +#[tokio::test] +async fn resolved_profile_schema_transform_is_applied_to_the_structured_output_schema() { + use crate::testkit::ScriptedModel; + + let model = Arc::new( + ScriptedModel::replies(vec![r#"{"value":"hi"}"#]).with_profile(ModelProfile { + native_structured_output: true, + json_schema: true, + schema_transform: Some(SchemaTransform::StripDefs), + ..ModelProfile::default() + }), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", model.clone()); + harness.with_policy(RunPolicy { + default_response_format: Some(ResponseFormat::auto( + "answer", + json!({ + "type": "object", + "$defs": {"Id": {"type": "string"}}, + "properties": {"value": {"$ref": "#/$defs/Id"}}, + }), + )), + ..RunPolicy::default() + }); + + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let requests = model.requests(); + let request = requests.first().expect("one model call"); + let ResponseFormat::JsonSchema { schema, .. } = request + .response_format + .as_ref() + .expect("provider-native structured output was requested") + else { + panic!("expected JsonSchema, got {:?}", request.response_format); + }; + assert!( + schema.get("$defs").is_none(), + "the structured-output schema must be transformed too: {schema:?}" + ); +} + +#[tokio::test] +async fn default_structured_mode_prompted_injects_schema_into_the_system_segment() { + use crate::testkit::ScriptedModel; + + let model = Arc::new( + ScriptedModel::replies(vec![r#"{"value":"hi"}"#]).with_profile(ModelProfile { + default_structured_mode: Some(StructuredMode::Prompted), + prompted_output_template: Some("Reply with JSON matching:".to_string()), + ..ModelProfile::default() + }), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", model.clone()); + harness.with_policy(RunPolicy { + default_response_format: Some(ResponseFormat::auto( + "answer", + json!({"type": "object", "properties": {"value": {"type": "string"}}}), + )), + ..RunPolicy::default() + }); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let requests = model.requests(); + let request = requests.first().expect("one model call"); + assert_eq!(request.response_format, Some(ResponseFormat::Text)); + let system_text: String = request + .messages + .iter() + .filter(|m| matches!(m, Message::System(_))) + .map(|m| m.text()) + .collect::>() + .join("\n"); + assert!( + system_text.contains("Reply with JSON matching:"), + "the profile's prompted template must be injected: {system_text}" + ); + assert!( + system_text.contains("JSON Schema for `answer`"), + "the schema must be described in the system segment: {system_text}" + ); + + let structured = run.structured.expect("structured output present"); + assert_eq!(structured["value"], "hi"); +} + +/// A middleware that pins a named reasoning effort onto every model request, +/// standing in for a caller that only knows the generic level name (not this +/// specific model's tuned budget). +struct RequestReasoningEffort(ReasoningEffort); + +#[async_trait] +impl Middleware<()> for RequestReasoningEffort { + fn name(&self) -> &str { + "request-reasoning-effort" + } + + async fn before_model( + &self, + _ctx: &mut RunContext<()>, + _state: &(), + request: &mut ModelRequest, + ) -> Result<()> { + request.reasoning = Some(ReasoningConfig::effort(self.0)); + Ok(()) + } +} + +#[tokio::test] +async fn thinking_level_map_resolves_a_named_reasoning_effort() { + use crate::testkit::ScriptedModel; + + let mut thinking_level_map = std::collections::BTreeMap::new(); + thinking_level_map.insert( + "high".to_string(), + ReasoningConfig { + effort: Some(ReasoningEffort::High), + budget_tokens: Some(32_000), + summary: None, + }, + ); + let model = Arc::new( + ScriptedModel::replies(vec!["done"]).with_profile(ModelProfile { + thinking_level_map, + ..ModelProfile::default() + }), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", model.clone()); + harness.push_middleware(Arc::new(RequestReasoningEffort(ReasoningEffort::High))); + + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let requests = model.requests(); + let request = requests.first().expect("one model call"); + assert_eq!( + request.reasoning.as_ref().and_then(|r| r.budget_tokens), + Some(32_000), + "the profile's tuned budget for the `high` level must replace the bare \ + effort the caller asked for: {:?}", + request.reasoning + ); +} + +#[tokio::test] +async fn thinking_level_map_does_not_override_an_explicit_budget() { + use crate::testkit::ScriptedModel; + + let mut thinking_level_map = std::collections::BTreeMap::new(); + thinking_level_map.insert( + "high".to_string(), + ReasoningConfig { + effort: Some(ReasoningEffort::High), + budget_tokens: Some(32_000), + summary: None, + }, + ); + let model = Arc::new( + ScriptedModel::replies(vec!["done"]).with_profile(ModelProfile { + thinking_level_map, + ..ModelProfile::default() + }), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", model.clone()); + + struct ExplicitBudget; + #[async_trait] + impl Middleware<()> for ExplicitBudget { + fn name(&self) -> &str { + "explicit-budget" + } + async fn before_model( + &self, + _ctx: &mut RunContext<()>, + _state: &(), + request: &mut ModelRequest, + ) -> Result<()> { + request.reasoning = Some(ReasoningConfig { + effort: Some(ReasoningEffort::High), + budget_tokens: Some(1_234), + summary: None, + }); + Ok(()) + } + } + harness.push_middleware(Arc::new(ExplicitBudget)); + + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let requests = model.requests(); + let request = requests.first().expect("one model call"); + assert_eq!( + request.reasoning.as_ref().and_then(|r| r.budget_tokens), + Some(1_234), + "an explicit caller budget must win over the profile's mapped default" + ); +} + +#[tokio::test] +async fn thinking_tags_are_split_out_of_a_unary_response_into_a_thinking_block() { + use crate::testkit::ScriptedModel; + + let model = Arc::new( + ScriptedModel::new(vec![ModelResponse::assistant( + "reasoning about itthe final answer", + )]) + .with_profile(ModelProfile { + thinking_tags: Some(("".to_string(), "".to_string())), + ..ModelProfile::default() + }), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", model.clone()); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let response = run.final_response.expect("final response"); + let thinking: Vec<&str> = response + .message + .content + .iter() + .filter_map(|block| match block { + ContentBlock::Thinking { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect(); + assert_eq!(thinking, vec!["reasoning about it"]); + assert_eq!(response.text(), "the final answer"); +} + +#[tokio::test] +async fn a_model_without_thinking_tags_configured_leaves_text_untouched() { + use crate::testkit::ScriptedModel; + + let model = Arc::new(ScriptedModel::new(vec![ModelResponse::assistant( + "not a reasoning tag hereplain text", + )])); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", model.clone()); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let response = run.final_response.expect("final response"); + assert!( + response + .message + .content + .iter() + .all(|block| !matches!(block, ContentBlock::Thinking { .. })), + "no profile means no tag pair to split on" + ); +} + +#[tokio::test] +async fn streaming_ignores_leading_whitespace_on_the_first_text_delta_only() { + use crate::testkit::StreamingMock; + + let model = Arc::new( + StreamingMock::from_text_chunks([" Hello", ", world"]).with_profile(ModelProfile { + ignore_streamed_leading_whitespace: true, + ..ModelProfile::default() + }), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("stream", model); + + let run = harness + .invoke_streaming( + &(), + (), + RunConfig::new("strip-leading-ws"), + vec![Message::user("hi")], + ) + .await + .expect("streaming run succeeds"); + + assert_eq!(run.text(), Some("Hello, world".to_string())); +} + +#[tokio::test] +async fn streaming_without_the_profile_flag_keeps_leading_whitespace() { + use crate::testkit::StreamingMock; + + let model = Arc::new(StreamingMock::from_text_chunks([" Hello", ", world"])); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("stream", model); + + let run = harness + .invoke_streaming( + &(), + (), + RunConfig::new("keep-leading-ws"), + vec![Message::user("hi")], + ) + .await + .expect("streaming run succeeds"); + + assert_eq!(run.text(), Some(" Hello, world".to_string())); +} + +// --------------------------------------------------------------------------- +// Tool-effect ledger (B5) +// --------------------------------------------------------------------------- + +mod tool_effects_test { + use super::*; + use crate::ids::{CallId, RunId}; + use crate::tool::{ + LedgerFailure, ToolEffect, ToolEffectLedger, ToolEffectSettle, ToolEffectStart, + ToolEffectStatus, + }; + use std::collections::HashMap; + use tokio::sync::Notify; + + /// In-memory [`ToolEffectLedger`] test double. Optionally fails every + /// `started` write (`fail_started`) and/or signals a [`Notify`] the + /// instant a `started` write lands (`on_started`), so a test can await + /// "the ledger has recorded this call as in flight" before acting. + #[derive(Clone, Default)] + struct InMemoryToolEffectLedger { + inner: Arc>>, + fail_started: bool, + on_started: Option>, + } + + impl InMemoryToolEffectLedger { + fn get(&self, run_id: &str, call_id: &str) -> Option { + self.inner + .lock() + .unwrap() + .get(&(run_id.to_string(), call_id.to_string())) + .cloned() + } + } + + #[async_trait] + impl ToolEffectLedger for InMemoryToolEffectLedger { + async fn started(&self, start: ToolEffectStart) -> Result<()> { + if self.fail_started { + return Err(TinyAgentsError::Tool("ledger unavailable".to_string())); + } + let effect = ToolEffect { + run_id: start.run_id.as_str().to_string(), + call_id: start.call_id.as_str().to_string(), + tool: start.tool, + status: ToolEffectStatus::Started, + idempotency_key: Some(start.idempotency_key), + effect_summary: start.effect_summary, + started_at: chrono::Utc::now(), + settled_at: None, + }; + self.inner + .lock() + .unwrap() + .insert((effect.run_id.clone(), effect.call_id.clone()), effect); + if let Some(notify) = &self.on_started { + notify.notify_one(); + } + Ok(()) + } + + async fn settled(&self, settle: ToolEffectSettle) -> Result<()> { + let run_id = settle.run_id.as_str().to_string(); + let call_id = settle.call_id.as_str().to_string(); + let mut guard = self.inner.lock().unwrap(); + let entry = guard + .entry((run_id.clone(), call_id.clone())) + .or_insert_with(|| ToolEffect { + run_id, + call_id, + tool: String::new(), + status: ToolEffectStatus::Started, + idempotency_key: None, + effect_summary: None, + started_at: chrono::Utc::now(), + settled_at: None, + }); + entry.status = settle.status; + entry.settled_at = Some(chrono::Utc::now()); + if let Some(summary) = settle.effect_summary { + entry.effect_summary = Some(summary); + } + Ok(()) + } + + async fn unresolved(&self, run_id: &str) -> Result> { + Ok(self + .inner + .lock() + .unwrap() + .values() + .filter(|effect| { + effect.run_id == run_id && effect.status == ToolEffectStatus::Started + }) + .cloned() + .collect()) + } + } + + /// A tool that never returns, used to prove an interrupted call (dropped + /// mid-flight, after its `started` ledger row landed) leaves that row + /// unresolved. + struct NeverFinishesTool { + notify: Arc, + } + + #[async_trait] + impl Tool for NeverFinishesTool { + fn name(&self) -> &str { + "hang" + } + fn description(&self) -> &str { + "never returns" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object"}) + } + async fn execute(&self, _arguments: serde_json::Value) -> anyhow::Result { + // Never notified by the test, so this hangs until the caller + // drops/aborts the run future. + self.notify.notified().await; + Ok(ToolResult::success("unreachable")) + } + } + + /// A tool that declares an explicit [`tinytools::ToolReplay`] policy, used + /// to drive [`AgentHarness::reconcile_tool_effects`] down each branch. + struct ReplayTool { + name: &'static str, + replay: tinytools::ToolReplay, + } + + #[async_trait] + impl Tool for ReplayTool { + fn name(&self) -> &str { + self.name + } + fn description(&self) -> &str { + "replay-classified tool" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object"}) + } + fn policy(&self) -> tinytools::ToolPolicy { + tinytools::ToolPolicy::default().with_runtime(tinytools::ToolRuntime { + replay: self.replay, + ..tinytools::ToolRuntime::default() + }) + } + async fn execute(&self, _arguments: serde_json::Value) -> anyhow::Result { + Ok(ToolResult::success("re-executed")) + } + } + + #[tokio::test] + async fn agent_loop_writes_started_then_completed_around_a_tool_call() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("call-1", "lookup", json!({"q": "x"})), + text_response("done", 4, 2), + ])), + ); + harness.register_tool(Arc::new(FakeTool::new("lookup", "tool-output"))); + + let ledger = Arc::new(InMemoryToolEffectLedger::default()); + let ctx: RunContext<()> = + RunContext::new(RunConfig::new("run-1"), ()).with_tool_effect_ledger(ledger.clone()); + + harness + .invoke_in_context(&(), ctx, vec![Message::user("please look up")]) + .await + .expect("run succeeds"); + + let effect = ledger + .get("run-1", "call-1") + .expect("ledger recorded the call"); + assert_eq!(effect.tool, "lookup"); + assert_eq!(effect.status, ToolEffectStatus::Completed); + assert!(effect.settled_at.is_some()); + } + + #[tokio::test] + async fn crash_between_started_and_settled_leaves_an_unresolved_row() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![tool_call_response( + "call-1", + "hang", + json!({}), + )])), + ); + let hang_notify = Arc::new(Notify::new()); + harness.register_tool(Arc::new(NeverFinishesTool { + notify: hang_notify.clone(), + })); + + let started_signal = Arc::new(Notify::new()); + let ledger = Arc::new(InMemoryToolEffectLedger { + on_started: Some(started_signal.clone()), + ..Default::default() + }); + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-crash"), ()) + .with_tool_effect_ledger(ledger.clone()); + + let harness = Arc::new(harness); + let run_harness = harness.clone(); + let handle = tokio::spawn(async move { + let _ = run_harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await; + }); + + // Wait for the ledger to observe the `started` write, then drop the + // run future (abort) before the tool — which never resolves on its + // own — could possibly settle it. This simulates a process crash + // between admission and settlement (TOOL-effect equivalent of a + // mid-flight kill). + started_signal.notified().await; + handle.abort(); + let _ = handle.await; + + let unresolved = ledger + .unresolved("run-crash") + .await + .expect("ledger read succeeds"); + assert_eq!(unresolved.len(), 1); + assert_eq!(unresolved[0].call_id, "call-1"); + assert_eq!(unresolved[0].status, ToolEffectStatus::Started); + } + + #[tokio::test] + async fn reconcile_leaves_a_safe_replay_call_pending_for_re_execution() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_tool(Arc::new(ReplayTool { + name: "safe_tool", + replay: tinytools::ToolReplay::Safe, + })); + + let ledger = Arc::new(InMemoryToolEffectLedger::default()); + ledger + .started(ToolEffectStart { + run_id: RunId::new("run-1"), + call_id: CallId::new("call-1"), + tool: "safe_tool".to_string(), + idempotency_key: "key".to_string(), + effect_summary: None, + }) + .await + .unwrap(); + + let recorder = crate::testkit::EventRecorder::new(); + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-1"), ()) + .with_events(recorder.sink()) + .with_tool_effect_ledger(ledger.clone()); + + let mut messages = vec![ + Message::user("go"), + Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: vec![ToolCall::new("call-1", "safe_tool", json!({}))], + usage: None, + origin: None, + }), + ]; + + let synthesized = harness + .reconcile_tool_effects(&ctx, "run-1", &mut messages, &Default::default()) + .await + .unwrap(); + + assert!(synthesized.is_empty()); + assert_eq!( + messages.len(), + 2, + "no tool answer appended for a Safe-replay call — the loop must \ + re-execute it" + ); + assert!( + recorder + .kinds() + .contains(&"tool.effect_reconciled".to_string()) + ); + // The ledger row is untouched (still `started`): re-execution will + // settle it normally through the ordinary started/settled path. + assert_eq!( + ledger.get("run-1", "call-1").unwrap().status, + ToolEffectStatus::Started + ); + } + + #[tokio::test] + async fn reconcile_synthesizes_an_interrupted_result_for_a_never_replay_call() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_tool(Arc::new(ReplayTool { + name: "risky_tool", + replay: tinytools::ToolReplay::Never, + })); + + let ledger = Arc::new(InMemoryToolEffectLedger::default()); + ledger + .started(ToolEffectStart { + run_id: RunId::new("run-1"), + call_id: CallId::new("call-1"), + tool: "risky_tool".to_string(), + idempotency_key: "key".to_string(), + effect_summary: None, + }) + .await + .unwrap(); + + let recorder = crate::testkit::EventRecorder::new(); + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-1"), ()) + .with_events(recorder.sink()) + .with_tool_effect_ledger(ledger.clone()); + + let mut messages = vec![ + Message::user("go"), + Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: vec![ToolCall::new("call-1", "risky_tool", json!({}))], + usage: None, + origin: None, + }), + ]; + + let synthesized = harness + .reconcile_tool_effects(&ctx, "run-1", &mut messages, &Default::default()) + .await + .unwrap(); + + assert_eq!(synthesized.len(), 1); + assert!(matches!(synthesized[0], Message::Tool(_))); + assert_eq!(synthesized[0].text(), "interrupted before settlement"); + assert_eq!(messages.len(), 3, "an interrupted answer was appended"); + assert_eq!( + ledger.get("run-1", "call-1").unwrap().status, + ToolEffectStatus::Interrupted + ); + assert!( + recorder + .kinds() + .contains(&"tool.effect_reconciled".to_string()) + ); + } + + #[tokio::test] + async fn ledger_started_failure_aborts_the_run_under_the_default_policy() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![tool_call_response( + "call-1", + "lookup", + json!({}), + )])), + ); + harness.register_tool(Arc::new(FakeTool::new("lookup", "tool-output"))); + + let ledger = Arc::new(InMemoryToolEffectLedger { + fail_started: true, + ..Default::default() + }); + let ctx: RunContext<()> = + RunContext::new(RunConfig::new("run-1"), ()).with_tool_effect_ledger(ledger); + + let err = harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect_err("LedgerFailure::Abort (the default) must fail the call"); + assert!(matches!(err, TinyAgentsError::Tool(_))); + } + + #[tokio::test] + async fn ledger_started_failure_is_ignored_under_the_continue_policy() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("call-1", "lookup", json!({})), + text_response("done", 4, 2), + ])), + ); + let tool = Arc::new(FakeTool::new("lookup", "tool-output")); + harness.register_tool(tool.clone()); + + let ledger = Arc::new(InMemoryToolEffectLedger { + fail_started: true, + ..Default::default() + }); + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-1"), ()) + .with_tool_effect_ledger(ledger) + .with_tool_effect_ledger_failure(LedgerFailure::Continue); + + let run = harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("LedgerFailure::Continue must not fail the run"); + assert_eq!(run.tool_calls, 1); + } + + // ── `resume_deferred` + tool-effect ledger (reconcile hazard fix) ─────── + + use crate::tool::DeferredToolResults; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Raises `ApprovalRequired` on its first invocation (mid-execution + /// deferral); on every later invocation, either succeeds or fails per + /// `fail_on_retry`, so a test can drive both the `Completed` and `Failed` + /// post-resume ledger transitions from the same tool shape. + struct ApprovalOnceTool { + attempts: AtomicUsize, + fail_on_retry: bool, + } + + #[async_trait] + impl Tool for ApprovalOnceTool { + fn name(&self) -> &str { + "wire" + } + fn description(&self) -> &str { + "defers once, then runs for real" + } + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object"}) + } + async fn execute(&self, _arguments: serde_json::Value) -> anyhow::Result { + if self.attempts.fetch_add(1, Ordering::SeqCst) == 0 { + return Err(TinyAgentsError::ApprovalRequired { + metadata: serde_json::Value::Null, + } + .into()); + } + if self.fail_on_retry { + anyhow::bail!("boom"); + } + Ok(ToolResult::success("approved-result")) + } + } + + #[tokio::test] + async fn resume_deferred_settles_the_deferred_row_completed_without_a_synthesized_crash_answer() + { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("call-1", "wire", json!({})), + text_response("all done", 4, 2), + ])), + ); + harness.register_tool(Arc::new(ApprovalOnceTool { + attempts: AtomicUsize::new(0), + fail_on_retry: false, + })); + + let ledger = Arc::new(InMemoryToolEffectLedger::default()); + let ctx: RunContext<()> = + RunContext::new(RunConfig::new("run-1"), ()).with_tool_effect_ledger(ledger.clone()); + let first = harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("a mid-execution deferral is not an error"); + let pending = first.deferred.clone().expect("approval pending"); + assert_eq!(pending.approvals[0].id, "call-1"); + + // The call left `started` only briefly: `defer_started_tool_call` + // settles it `Deferred` the instant the deferral is filed, not left + // `started` for `reconcile_tool_effects` to mistake for a crash. + let effect = ledger + .get("run-1", "call-1") + .expect("ledger recorded the call"); + assert_eq!(effect.status, ToolEffectStatus::Deferred); + assert!(effect.settled_at.is_some()); + + let results = DeferredToolResults::new().approve("call-1"); + let ctx2: RunContext<()> = + RunContext::new(RunConfig::new("run-1"), ()).with_tool_effect_ledger(ledger.clone()); + let run = harness + .resume_deferred(&(), ctx2, first.messages.clone(), results) + .await + .expect("resume completes the run"); + + // The call must receive its real answer, never the synthesized + // "interrupted before settlement" a crash-reconcile would produce. + let answer = run.messages.iter().find_map(|message| match message { + Message::Tool(tool) if tool.tool_call_id == "call-1" => Some(message.text()), + _ => None, + }); + assert_eq!(answer.as_deref(), Some("approved-result")); + assert_eq!(run.text().as_deref(), Some("all done")); + + assert_eq!( + ledger.get("run-1", "call-1").unwrap().status, + ToolEffectStatus::Completed + ); + } + + #[tokio::test] + async fn resume_deferred_settles_the_deferred_row_failed_when_the_retry_errors() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![tool_call_response( + "call-1", + "wire", + json!({}), + )])), + ); + harness.register_tool(Arc::new(ApprovalOnceTool { + attempts: AtomicUsize::new(0), + fail_on_retry: true, + })); + + let ledger = Arc::new(InMemoryToolEffectLedger::default()); + let ctx: RunContext<()> = + RunContext::new(RunConfig::new("run-1"), ()).with_tool_effect_ledger(ledger.clone()); + let first = harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("a mid-execution deferral is not an error"); + assert_eq!( + ledger.get("run-1", "call-1").unwrap().status, + ToolEffectStatus::Deferred + ); + + let results = DeferredToolResults::new().approve("call-1"); + let ctx2: RunContext<()> = + RunContext::new(RunConfig::new("run-1"), ()).with_tool_effect_ledger(ledger.clone()); + harness + .resume_deferred(&(), ctx2, first.messages.clone(), results) + .await + .expect_err("the retried execution error propagates"); + + assert_eq!( + ledger.get("run-1", "call-1").unwrap().status, + ToolEffectStatus::Failed + ); + } + + #[tokio::test] + async fn resume_deferred_reconciles_a_crashed_sibling_but_excludes_the_call_it_resolves() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![text_response( + "all done", 4, 2, + )])), + ); + harness.register_tool(Arc::new(FakeTool::new("approve_tool", "approved-result"))); + harness.register_tool(Arc::new(ReplayTool { + name: "risky_tool", + replay: tinytools::ToolReplay::Never, + })); + + let ledger = Arc::new(InMemoryToolEffectLedger::default()); + // `call-a` was deliberately deferred mid-execution (already settled + // `Deferred`, matching what `defer_started_tool_call` does). + ledger + .started(ToolEffectStart { + run_id: RunId::new("run-x"), + call_id: CallId::new("call-a"), + tool: "approve_tool".to_string(), + idempotency_key: "key-a".to_string(), + effect_summary: None, + }) + .await + .unwrap(); + ledger + .settled(ToolEffectSettle { + run_id: RunId::new("run-x"), + call_id: CallId::new("call-a"), + status: ToolEffectStatus::Deferred, + effect_summary: None, + }) + .await + .unwrap(); + // `call-b`'s process crashed mid-flight: admitted and started, but + // never settled or deferred — the genuine crash artifact. + ledger + .started(ToolEffectStart { + run_id: RunId::new("run-x"), + call_id: CallId::new("call-b"), + tool: "risky_tool".to_string(), + idempotency_key: "key-b".to_string(), + effect_summary: None, + }) + .await + .unwrap(); + + let recorder = crate::testkit::EventRecorder::new(); + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-x"), ()) + .with_events(recorder.sink()) + .with_tool_effect_ledger(ledger.clone()); + + let messages = vec![ + Message::user("go"), + Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: vec![ + ToolCall::new("call-a", "approve_tool", json!({})), + ToolCall::new("call-b", "risky_tool", json!({})), + ], + usage: None, + origin: None, + }), + ]; + let results = DeferredToolResults::new().approve("call-a"); + + let run = harness + .resume_deferred(&(), ctx, messages, results) + .await + .expect("resume completes despite the crashed sibling"); + + // `call-b` had no live `results` entry: `reconcile_tool_effects` + // (run before `results` is applied) answered it as interrupted, + // per its default `ToolReplay::Never`. + let call_b_answer = run.messages.iter().find_map(|message| match message { + Message::Tool(tool) if tool.tool_call_id == "call-b" => Some(message.text()), + _ => None, + }); + assert_eq!( + call_b_answer.as_deref(), + Some("interrupted before settlement") + ); + assert_eq!( + ledger.get("run-x", "call-b").unwrap().status, + ToolEffectStatus::Interrupted + ); + + // `call-a` is exactly what `results` resolves: it must not receive a + // synthesized crash answer, and it ran for real through the approval + // path instead. + let call_a_answer = run.messages.iter().find_map(|message| match message { + Message::Tool(tool) if tool.tool_call_id == "call-a" => Some(message.text()), + _ => None, + }); + assert_eq!(call_a_answer.as_deref(), Some("approved-result")); + assert_eq!( + ledger.get("run-x", "call-a").unwrap().status, + ToolEffectStatus::Completed + ); + + assert!( + recorder + .kinds() + .contains(&"tool.effect_reconciled".to_string()), + "the crashed sibling was reconciled" + ); + } +} diff --git a/crates/tinyagents-harness/src/agent_loop/tool_changes.rs b/crates/tinyagents-harness/src/agent_loop/tool_changes.rs new file mode 100644 index 00000000..5b92669e --- /dev/null +++ b/crates/tinyagents-harness/src/agent_loop/tool_changes.rs @@ -0,0 +1,170 @@ +//! Transcript-carried tool-set change patches (gap B6, +//! `docs/runtime-comparison/plan.md`). +//! +//! A [`crate::tool::toolset::ToolSet`] chain's live tool set can legitimately +//! vary turn to turn (`ToolSet::tools` is documented as "called once per +//! turn"). Naively rebuilding the whole system prompt to describe a changed +//! loadout would bust a provider's cached prefix on every such change. These +//! two pure functions are the write side of that mechanism: [`diff_tool_set`] +//! computes the minimal delta between what a transcript has declared so far +//! and what is live now, and [`apply_tool_change_patch`] lays that delta onto +//! the transcript — either as a small mid-conversation +//! [`tinyinference_llm::message::Message::System`] patch (when the resolved +//! model's [`tinyinference_llm::model::ModelProfile:: +//! mid_conversation_system_messages`] allows it) or folded into the leading +//! system message (when it does not). The read side, +//! [`tinyinference_llm::message::replay_system_state`], reconstructs the +//! effective tool set from the same transcript. + +use std::collections::{BTreeMap, HashSet}; + +use tinyinference_llm::message::{Message, SystemMessage}; +use tinyinference_llm::tool::ToolSchema; + +/// Name of the [`SystemMessage::sections`] entry a tool-change patch writes +/// its human-readable summary under. +pub(super) const TOOL_CHANGES_SECTION: &str = "tool_changes"; + +/// Computes the [`SystemMessage`] patch needed to bring `previous` (what the +/// transcript has declared so far, via [`tinyinference_llm::message:: +/// replay_system_state`] or an equivalent running tally) in line with +/// `current` (the toolset chain's live set this turn), or `None` when they +/// already agree. +/// +/// Comparison is by full schema equality (name, description, parameters, +/// format): a tool whose declaration changed under the same name is reported +/// only via `tools_added` (the newer schema), never additionally via +/// `tools_removed` — [`replay_system_state`](tinyinference_llm::message::replay_system_state)'s +/// fold semantics already let a later `tools_added` entry for an existing +/// name supersede the earlier one. +pub(super) fn diff_tool_set( + previous: &[ToolSchema], + current: &[ToolSchema], +) -> Option { + let previous_by_name: BTreeMap<&str, &ToolSchema> = previous + .iter() + .map(|schema| (schema.name.as_str(), schema)) + .collect(); + let current_names: HashSet<&str> = current.iter().map(|schema| schema.name.as_str()).collect(); + + let mut tools_added: Vec = current + .iter() + .filter(|schema| previous_by_name.get(schema.name.as_str()) != Some(schema)) + .cloned() + .collect(); + tools_added.sort_by(|left, right| left.name.cmp(&right.name)); + + let mut tools_removed: Vec = previous_by_name + .keys() + .filter(|name| !current_names.contains(*name)) + .map(|name| (*name).to_string()) + .collect(); + tools_removed.sort(); + + if tools_added.is_empty() && tools_removed.is_empty() { + return None; + } + + let mut sections = BTreeMap::new(); + sections.insert( + TOOL_CHANGES_SECTION.to_string(), + Some(describe_delta(&tools_added, &tools_removed)), + ); + + Some(SystemMessage { + content: Vec::new(), + sections, + tools_added, + tools_removed, + }) +} + +/// Renders a short human-readable summary of a tool-set delta, used as the +/// patch's `tool_changes` section text. +fn describe_delta(added: &[ToolSchema], removed: &[String]) -> String { + let mut lines = Vec::new(); + if !added.is_empty() { + let names: Vec<&str> = added.iter().map(|schema| schema.name.as_str()).collect(); + lines.push(format!("Tools now available: {}.", names.join(", "))); + } + if !removed.is_empty() { + lines.push(format!( + "Tools no longer available: {}.", + removed.join(", ") + )); + } + lines.join(" ") +} + +/// Applies a tool-change `patch` to the working transcript. +/// +/// When `mid_conversation` is `true` (the resolved model's +/// `ModelProfile::mid_conversation_system_messages` allows a system message +/// anywhere in the transcript, e.g. the OpenAI-compatible chat path), `patch` +/// is appended as a new tail [`Message::System`]. Because +/// `agent_loop/run_loop.rs`'s `system_end` treats only the transcript's +/// **leading** run of `Message::System` entries as the cacheable prefix, a +/// patch appended after any non-system message (the ordinary case: a run +/// always opens with at least one user turn before any tool-set change can +/// occur) lands in the non-cacheable tail and the prefix's +/// [`crate::prompt::PromptBuilder`] fingerprint is unchanged. +/// +/// When `mid_conversation` is `false` (for example Anthropic, whose Messages +/// API hoists every `Message::System` in the transcript into one leading +/// `system` array regardless of position — a "mid-conversation" system +/// message would not actually land where it appears), `patch` is folded into +/// the leading system message instead: its `content` is appended, its +/// `sections` are merged key-by-key, and its `tools_added`/`tools_removed` +/// are merged into the leading message's own lists. This *does* change the +/// leading message's content and therefore the prefix fingerprint — expected, +/// since the provider has no mid-transcript system slot for the patch to +/// occupy without moving it there on the wire anyway. A transcript with no +/// leading `Message::System` gets one inserted at the front. +pub(super) fn apply_tool_change_patch( + messages: &mut Vec, + patch: SystemMessage, + mid_conversation: bool, +) { + if mid_conversation { + messages.push(Message::System(patch)); + return; + } + match messages.first_mut() { + Some(Message::System(leading)) => fold_patch(leading, patch), + _ => messages.insert(0, Message::System(patch)), + } +} + +/// Merges `patch` into `leading` in place: content is appended, sections are +/// merged key-by-key (`None` removes), and tool deltas are merged so the +/// leading message's own `tools_added`/`tools_removed` reflect the net +/// effect of both the original declaration and the patch. +fn fold_patch(leading: &mut SystemMessage, patch: SystemMessage) { + leading.content.extend(patch.content); + for (name, value) in patch.sections { + match value { + Some(text) => { + leading.sections.insert(name, Some(text)); + } + None => { + leading.sections.remove(&name); + } + } + } + for tool in patch.tools_added { + leading.tools_removed.retain(|name| *name != tool.name); + leading + .tools_added + .retain(|existing| existing.name != tool.name); + leading.tools_added.push(tool); + } + for name in patch.tools_removed { + leading.tools_added.retain(|existing| existing.name != name); + if !leading.tools_removed.contains(&name) { + leading.tools_removed.push(name); + } + } +} + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-harness/src/agent_loop/tool_changes/test.rs b/crates/tinyagents-harness/src/agent_loop/tool_changes/test.rs new file mode 100644 index 00000000..f7da7d86 --- /dev/null +++ b/crates/tinyagents-harness/src/agent_loop/tool_changes/test.rs @@ -0,0 +1,209 @@ +//! Unit tests for the tool-change patch mechanism (gap B6). +//! +//! Covers: exactly-one-patch diffing, transcript round-tripping through +//! [`replay_system_state`], and that the mid-conversation insert path keeps +//! the leading-prefix [`PromptBuilder`] fingerprint stable while the folded +//! path (correctly) does not. + +use serde_json::json; +use tinyinference_llm::message::{Message, replay_system_state}; +use tinyinference_llm::tool::ToolSchema; + +use super::*; +use crate::prompt::PromptBuilder; + +fn tool(name: &str) -> ToolSchema { + ToolSchema::new(name, format!("{name} tool"), json!({"type": "object"})) +} + +#[test] +fn no_diff_when_tool_sets_match() { + let set = vec![tool("search"), tool("read_file")]; + assert!(diff_tool_set(&set, &set).is_none()); +} + +#[test] +fn diff_reports_additions_and_removals() { + let previous = vec![tool("search"), tool("read_file")]; + let current = vec![tool("search"), tool("browse")]; + + let patch = diff_tool_set(&previous, ¤t).expect("tool set changed"); + let added: Vec<&str> = patch.tools_added.iter().map(|t| t.name.as_str()).collect(); + assert_eq!(added, vec!["browse"]); + assert_eq!(patch.tools_removed, vec!["read_file".to_string()]); + assert!(patch.sections.contains_key(TOOL_CHANGES_SECTION)); + let summary = patch.sections[TOOL_CHANGES_SECTION].as_ref().unwrap(); + assert!(summary.contains("browse")); + assert!(summary.contains("read_file")); +} + +#[test] +fn diff_reports_a_changed_schema_as_an_addition_only() { + let previous = vec![ToolSchema::new("search", "v1", json!({"type": "object"}))]; + let current = vec![ToolSchema::new("search", "v2", json!({"type": "object"}))]; + + let patch = diff_tool_set(&previous, ¤t).expect("schema changed"); + assert_eq!(patch.tools_added.len(), 1); + assert_eq!(patch.tools_added[0].description, "v2"); + assert!(patch.tools_removed.is_empty()); +} + +/// Test 1: a tool-set change mid-run results in **exactly one** patch system +/// message being appended when the profile supports mid-conversation system +/// messages. +#[test] +fn mid_conversation_patch_appends_exactly_one_system_message() { + let mut messages = vec![ + Message::system("baseline persona"), + Message::user("hi"), + Message::assistant("hello"), + ]; + let system_count_before = messages + .iter() + .filter(|m| matches!(m, Message::System(_))) + .count(); + + let previous = vec![tool("search")]; + let current = vec![tool("search"), tool("browse")]; + let patch = diff_tool_set(&previous, ¤t).expect("tool set changed"); + apply_tool_change_patch(&mut messages, patch, true); + + let system_count_after = messages + .iter() + .filter(|m| matches!(m, Message::System(_))) + .count(); + assert_eq!(system_count_after, system_count_before + 1); + // The patch landed at the tail, after the existing conversation. + assert!(matches!(messages.last(), Some(Message::System(_)))); +} + +/// Test 1 (fold variant): when the profile does not support mid-conversation +/// system messages, the patch is folded into the leading system message +/// instead of appending a new one — the system-message *count* does not +/// grow, but the delta is still fully recorded. +#[test] +fn folded_patch_does_not_add_a_new_system_message() { + let mut messages = vec![Message::system("baseline persona"), Message::user("hi")]; + let previous = vec![tool("search")]; + let current = vec![tool("search"), tool("browse")]; + let patch = diff_tool_set(&previous, ¤t).expect("tool set changed"); + apply_tool_change_patch(&mut messages, patch, false); + + let system_count = messages + .iter() + .filter(|m| matches!(m, Message::System(_))) + .count(); + assert_eq!(system_count, 1); + let Message::System(leading) = &messages[0] else { + panic!("expected a leading system message"); + }; + assert_eq!( + leading + .tools_added + .iter() + .map(|t| t.name.as_str()) + .collect::>(), + vec!["browse"] + ); +} + +/// Test 2: `replay_system_state` on a transcript carrying multiple patches +/// reconstructs the same effective tool set that was live when each patch was +/// produced. +#[test] +fn replay_round_trips_a_sequence_of_mid_conversation_patches() { + let mut messages = vec![Message::system("baseline persona"), Message::user("hi")]; + + // Turn 1 -> 2: the live set grows from {} to {search, read_file}. + let live_1 = vec![tool("search"), tool("read_file")]; + if let Some(patch) = diff_tool_set(&[], &live_1) { + apply_tool_change_patch(&mut messages, patch, true); + } + messages.push(Message::assistant("using search")); + + // Turn 2 -> 3: read_file drops, browse appears. + let live_2 = vec![tool("search"), tool("browse")]; + if let Some(patch) = diff_tool_set(&live_1, &live_2) { + apply_tool_change_patch(&mut messages, patch, true); + } + messages.push(Message::assistant("using browse")); + + let (_, replayed_tools) = replay_system_state(&messages); + let mut replayed_names: Vec<&str> = replayed_tools.iter().map(|t| t.name.as_str()).collect(); + replayed_names.sort(); + let mut expected_names: Vec<&str> = live_2.iter().map(|t| t.name.as_str()).collect(); + expected_names.sort(); + assert_eq!(replayed_names, expected_names); +} + +/// Test 3: with a profile that supports mid-conversation system messages, +/// inserting the patch does not change the fingerprint of the unchanged +/// leading system-segment prefix. +#[test] +fn mid_conversation_patch_keeps_the_leading_prefix_fingerprint_stable() { + let leading = vec![Message::system("baseline persona")]; + + let mut before = PromptBuilder::new(); + before.push_system("system", leading.clone()); + let fingerprint_before = before.fingerprint(); + + // Simulate the patched transcript: the leading system run is untouched; + // the patch lands after the tail (a user turn, then the patch). + let mut messages = leading.clone(); + messages.push(Message::user("hi")); + let previous = vec![tool("search")]; + let current = vec![tool("search"), tool("browse")]; + let patch = diff_tool_set(&previous, ¤t).expect("tool set changed"); + apply_tool_change_patch(&mut messages, patch, true); + + // Recompute `system_end` the way `run_loop.rs` does: the leading run of + // `Message::System` messages only. + let system_end = messages + .iter() + .take_while(|m| matches!(m, Message::System(_))) + .count(); + let mut after = PromptBuilder::new(); + after.push_system("system", messages[..system_end].to_vec()); + let fingerprint_after = after.fingerprint(); + + assert_eq!(fingerprint_before, fingerprint_after); +} + +/// Test 3 (companion): when the profile does *not* support mid-conversation +/// system messages, the patch is folded into the leading system message, so +/// the prefix fingerprint is expected to change — this asserts correctness +/// of the folded content, not cache stability. +#[test] +fn folded_patch_changes_the_leading_prefix_fingerprint_but_stays_correct() { + let leading = vec![Message::system("baseline persona")]; + + let mut before = PromptBuilder::new(); + before.push_system("system", leading.clone()); + let fingerprint_before = before.fingerprint(); + + let mut messages = leading.clone(); + let previous = vec![tool("search")]; + let current = vec![tool("search"), tool("browse")]; + let patch = diff_tool_set(&previous, ¤t).expect("tool set changed"); + apply_tool_change_patch(&mut messages, patch, false); + + let mut after = PromptBuilder::new(); + after.push_system("system", messages.clone()); + let fingerprint_after = after.fingerprint(); + + // The fold path is expected to recompute the fingerprint... + assert_ne!(fingerprint_before, fingerprint_after); + // ...but the fold itself is correct: the leading message now declares the + // full current tool set. + let Message::System(leading_after) = &messages[0] else { + panic!("expected a leading system message"); + }; + assert_eq!( + leading_after + .tools_added + .iter() + .map(|t| t.name.as_str()) + .collect::>(), + vec!["browse"] + ); +} diff --git a/crates/tinyagents-harness/src/agent_loop/tools.rs b/crates/tinyagents-harness/src/agent_loop/tools.rs index 984e5084..44209118 100644 --- a/crates/tinyagents-harness/src/agent_loop/tools.rs +++ b/crates/tinyagents-harness/src/agent_loop/tools.rs @@ -14,8 +14,14 @@ //! "Started/terminal pairing" below. //! 2. **Execution**: when the turn requests **two or more** tools and **no //! tool-wrap middleware** ([`crate::middleware::ToolMiddleware`]) -//! is registered, the admitted calls run **concurrently** -//! (`join_all`), so turn latency is the slowest tool instead of the sum. +//! is registered, the admitted calls run **concurrently**, so turn latency +//! is the slowest tool instead of the sum (bounded by +//! [`RunLimits::max_tool_concurrency`][crate::limits::RunLimits::max_tool_concurrency] +//! when set — see I-8; unbounded, i.e. every eligible call starts at once, +//! when unset). Lifecycle middleware does **not** force the serial path: +//! admission (phase 1) already ran every `before_tool` hook to completion, +//! serially, before any concurrent future is built, so there is nothing +//! left for a lifecycle middleware to mutate once execution starts. //! Otherwise execution is serial, preserving the historical semantics. //! [`AgentEvent::ToolStarted`] is emitted here, once every admission has //! succeeded, so a call that is announced always runs. @@ -70,16 +76,22 @@ //! - **Cancellation**: observed between admissions (before each call starts), //! matching the serial path, which also never interrupts a mid-flight tool. //! - **Errors**: an `Err` fails the turn at the first call in original order. -//! Difference: in -//! serial mode later calls never start after a failure; in concurrent mode -//! they were already in flight and run to completion (their results are -//! discarded). Tools that must not observe a sibling's failure should be run -//! under a tool-wrap middleware (serial) or a harness without -//! parallel-capable turns. +//! Difference: in serial mode later calls never start after a failure; in +//! concurrent mode they were already in flight and run to completion, but +//! their results are discarded — each already-started sibling still gets +//! exactly one terminal event, [`AgentEvent::ToolFailed`] with +//! `"aborted: sibling tool call failed"`, so the started/terminal invariant +//! above holds even on this path. Tools that must not observe a sibling's +//! failure should be run under a tool-wrap middleware (serial) or a harness +//! without parallel-capable turns. //! use super::model_call::ToolCallBase; use super::*; -use crate::tool::{ToolDispatch, provider_schema}; +use crate::tool::{ + DeferredToolRequests, LedgerFailure, ToolDispatch, ToolEffectSettle, ToolEffectStart, + ToolEffectStatus, provider_schema, +}; +use sha2::{Digest, Sha256}; use tinyinference_llm::message::ContentBlock; use tinytools::{ToolCall as CanonicalToolCall, ToolCallId, ToolCallOptions}; @@ -95,6 +107,54 @@ enum ResolvedToolCall { /// tool, invalid arguments); a success result for an intrinsic answer /// (`tool_search`). Answered(tinytools::ToolResult), + /// No tool runs *yet*: the call needs a human decision or host-side + /// execution first (A2). The loop finishes the batch's other calls and + /// then hands every deferred request to the caller (or the inline + /// `DeferredToolHandler`). + Deferred(DeferredRequest), +} + +/// Which [`DeferredToolRequests`] list a deferred call belongs to. +#[derive(Clone, Copy, Debug)] +pub(super) enum DeferredKind { + /// Needs an [`crate::tool::ToolApprovalDecision`]; the harness runs the tool + /// on approval. + Approval, + /// Needs a [`crate::tool::DeferredCallResult`]; the host runs the tool. + External, +} + +/// One call the loop is handing back instead of answering. +#[derive(Clone, Debug)] +pub(super) struct DeferredRequest { + /// The call as the model made it (original arguments, so an approver + /// sees — and may edit — exactly what the model asked for). + pub(super) call: ToolCall, + pub(super) kind: DeferredKind, + /// Stable label for the `ToolDeferred` event. + pub(super) reason: &'static str, + /// Host-only payload from `ApprovalRequired`/`CallDeferred`, if any. + pub(super) metadata: Option, +} + +impl DeferredRequest { + fn approval(call: ToolCall, reason: &'static str, metadata: Option) -> Self { + Self { + call, + kind: DeferredKind::Approval, + reason, + metadata, + } + } + + fn external(call: ToolCall, reason: &'static str, metadata: Option) -> Self { + Self { + call, + kind: DeferredKind::External, + reason, + metadata, + } + } } /// One requested call after admission, in original order. @@ -116,10 +176,21 @@ enum AdmittedCall { call: ToolCall, result: tinytools::ToolResult, }, + /// Deferred at admission: nothing runs, nothing is announced; folded into + /// the batch's [`DeferredToolRequests`] in original order. + Deferred(DeferredRequest), } /// One transcript slot per requested call, in original order, used by the /// concurrent path to reassemble results deterministically. +/// +/// `Recovered` carries a full `ToolResult` (now noticeably larger than +/// `Execute`'s no-op payload since the vendor `tinytools::ToolControl`/ +/// `follow_up`/`metadata` fields landed); boxing it would touch every +/// construction and pattern-match site in this file for a one-shot, +/// short-lived per-call value, so the size difference is accepted here +/// rather than threaded through as indirection. +#[allow(clippy::large_enum_variant)] enum ToolSlot { /// An executed call: consumes the next prepared/result pair in order. Execute, @@ -128,6 +199,8 @@ enum ToolSlot { call: ToolCall, result: tinytools::ToolResult, }, + /// A call deferred at admission (see [`AdmittedCall::Deferred`]). + Deferred(DeferredRequest), } /// Admission metadata for one executable call, paired 1:1 (in order) with its @@ -135,6 +208,10 @@ enum ToolSlot { struct PreparedToolCall { call_id: CallId, tool_name: String, + /// The admitted call, kept so an execution-time deferral + /// (`ApprovalRequired`/`CallDeferred` raised by the tool) can hand the + /// original request back through [`DeferredToolRequests`]. + call: ToolCall, options: ToolCallOptions, captured_input: Option, started_at_ms: u64, @@ -142,7 +219,61 @@ struct PreparedToolCall { output_origin: crate::host::ContentOrigin, } +/// Derives a best-effort deduplication key for one tool call from its name +/// and arguments (B5). +/// +/// `tinytools::ToolPolicy` does not currently declare an explicit +/// idempotency key field, so this hashes `(tool name, arguments)` with +/// SHA-256: two calls to the same tool with identical arguments derive the +/// same key, which is exactly what a host wants to notice when deciding +/// whether an orphaned effect might have already landed. This is content +/// equality, not a cryptographic guarantee — a tool whose "same effect" notion +/// differs from "identical arguments" (e.g. one that reads a clock) should +/// not rely on this key alone. +fn tool_call_idempotency_key(tool_name: &str, arguments: &Value) -> String { + let mut hasher = Sha256::new(); + hasher.update(tool_name.as_bytes()); + hasher.update([0u8]); + hasher.update(serde_json::to_vec(arguments).unwrap_or_default()); + let digest = hasher.finalize(); + digest.iter().map(|byte| format!("{byte:02x}")).collect() +} + impl AgentHarness { + /// Resolves the effective host tool allow-list for `ctx`, or `Ok(None)` + /// when nothing should be restricted. + /// + /// Three cases: + /// - Not a hosted run at all (`host_invocation_binding` returns `None`): + /// no host allow-list concept applies, so this returns `None` (allow + /// every registered tool, same as an explicit-model run). + /// - Hosted, and the resolved [`crate::host::AgentDefinition`] declared a + /// non-empty tool list: returns that set. Only those names are + /// dispatchable, checked with plain set membership — no empty-set + /// bypass (that bypass was I-9: an empty `HashSet` used to mean + /// "unrestricted" instead of "nothing"). + /// - Hosted, but the definition declared no tools at all (an empty or + /// absent list): fails closed by default — returns `Some(HashSet::new())`, + /// which allows nothing — unless + /// [`crate::host::HostCapabilities::fail_closed_tool_allowlist`] was + /// explicitly turned off on this host, in which case it returns `None` + /// (legacy unrestricted behavior, opt-in only). + pub(super) fn resolve_tool_allowlist( + &self, + ctx: &RunContext, + ) -> Result>> { + let Some(binding) = crate::runtime::host_invocation_binding::(ctx)? else { + return Ok(None); + }; + Ok(match &binding.allowed_tools { + Some(declared) => Some(declared.clone()), + None if binding.host.fail_closed_tool_allowlist => { + Some(std::collections::HashSet::new()) + } + None => None, + }) + } + /// Builds the run's deferred-tool catalogue: every /// [`tinytools::ToolExposure::Deferred`] registration the host allow-list /// admits, or an empty catalogue when discovery is disabled. @@ -184,12 +315,17 @@ impl AgentHarness { { return Ok(None); } - let allowed_tools = crate::runtime::host_invocation_binding::(ctx)? - .map(|binding| binding.allowed_tools); + // Reuses `resolve_tool_allowlist` (I-9's fail-closed allow-list + // resolution) rather than reading `binding.allowed_tools` directly, + // so the discovery bridge honors the same + // `fail_closed_tool_allowlist` policy as the direct tool set built in + // `run_loop_body` — an empty declared list never falls back to + // "unrestricted" here either. + let allowed_tools = self.resolve_tool_allowlist(ctx)?; let host_allows = |name: &str| { allowed_tools .as_ref() - .is_none_or(|allowed| allowed.is_empty() || allowed.contains(name)) + .is_none_or(|allowed| allowed.contains(name)) }; let catalog = self.deferred_catalog(&host_allows); if catalog.is_empty() { @@ -291,6 +427,11 @@ impl AgentHarness { /// Dispatches to the concurrent path when it is safe (see the module docs /// for the exact conditions and preserved semantics); otherwise runs the /// historical serial path. + /// + /// Returns the calls the batch **deferred** (A2) — empty for the common + /// case. A deferred call gets no tool-result row; every other call in the + /// batch is still executed and answered, so the caller only has to decide + /// what to do with the pending ones (exit, or resolve inline). pub(super) async fn execute_tools( &self, state: &State, @@ -299,7 +440,7 @@ impl AgentHarness { status: &mut HarnessRunStatus, messages: &mut Vec, tool_calls: Vec, - ) -> Result<()> { + ) -> Result { // Injection and argument normalization change the model payload before // execution. Until admission has produced those authoritative values, // a declaration cannot safely make a parallel decision from raw model @@ -312,7 +453,6 @@ impl AgentHarness { if should_execute_tools_concurrently( tool_calls.len(), canonical_parallel_safe, - self.middleware.len(), self.middleware.tool_middleware_len(), ) { self.execute_tools_concurrently(state, ctx, run, status, messages, tool_calls) @@ -392,14 +532,61 @@ impl AgentHarness { // they answer the model and let it try again, so counting them is what // bounds the correction loop. if let Err(err) = self.middleware.run_before_tool(ctx, state, call).await { - tinyagents_tracing::debug!( + tracing::debug!( "[agent_loop::tools] `before_tool` refused `{}` (call `{}`); \ releasing its tool-call slot: {err}", call.name, call.id ); ctx.limits.rollback_tool_calls(1); - return Err(err); + // A2/A3 signals from a `before_tool` hook are decisions about + // *this call*, not failures of the run: a deferral hands the + // call back to the host, and the retry/failed vocabulary answers + // the model without running the tool (the `HumanApprovalMiddleware` + // `Deny` outcome, for one). + return match err { + TinyAgentsError::ApprovalRequired { metadata } => Ok(ResolvedToolCall::Deferred( + DeferredRequest::approval(call.clone(), "approval_required", Some(metadata)), + )), + TinyAgentsError::CallDeferred { metadata } => Ok(ResolvedToolCall::Deferred( + DeferredRequest::external(call.clone(), "call_deferred", Some(metadata)), + )), + TinyAgentsError::ToolFailed(message) => Ok(ResolvedToolCall::Answered( + tinytools::ToolResult::failed(message), + )), + TinyAgentsError::ModelRetry(message) => Ok(ResolvedToolCall::Answered( + tinytools::ToolResult::retry(message), + )), + other => Err(other), + }; + } + + // Before giving up on provider-unparseable arguments (below), try the + // conservative, meaning-preserving repairs in the protocol crate's + // `tinytools_agent::repair::json` (unquoted keys, redundant wrapping + // braces, leaked chat-template quote tokens — see that module's doc + // comment for the exact defects it targets). + // This is the one place I-13 asked for it applied: admission was + // short-circuiting straight to a tool error without ever trying the + // repair the module exists for. On success the call proceeds through + // normal (schema) validation below as if the provider had sent it + // clean, rather than round-tripping a "fix your JSON" error the model + // often cannot actually act on. + if call.invalid.is_some() + && let Some(raw) = call.arguments.as_str() + && let Some(repaired) = tinytools_agent::repair::json::recover_object(raw) + { + let call_id = CallId::new(call.id.clone()); + let record = ctx.emit(AgentEvent::InvalidToolArgs { + call_id, + tool_name: call.name.clone(), + arguments: call.arguments.clone(), + error: call.invalid.clone().unwrap_or_default(), + recovery: "repaired".to_string(), + }); + status.set_last_event(record.id); + call.arguments = repaired; + call.invalid = None; } // The provider marked this call's arguments unparseable (a small local @@ -431,11 +618,10 @@ impl AgentHarness { // Hosted turns carry an explicit definition allowlist. Do not merely // hide disallowed schemas: a model can still fabricate a name, so the // dispatch boundary must reject it too. - let allowed_tools = crate::runtime::host_invocation_binding::(ctx)? - .map(|binding| binding.allowed_tools); + let allowed_tools = self.resolve_tool_allowlist(ctx)?; let is_allowed = allowed_tools .as_ref() - .is_none_or(|allowed| allowed.is_empty() || allowed.contains(&call.name)); + .is_none_or(|allowed| allowed.contains(&call.name)); let (dispatch, tool) = match is_allowed .then(|| self.tools.model_dispatch(&call.name)) .flatten() @@ -458,9 +644,9 @@ impl AgentHarness { .tools .dispatch(tool_name) .filter(|_| { - allowed_tools.as_ref().is_none_or(|allowed| { - allowed.is_empty() || allowed.contains(tool_name) - }) + allowed_tools + .as_ref() + .is_none_or(|allowed| allowed.contains(tool_name)) }) .map(|dispatch| (tool_name.clone(), dispatch)), _ => None, @@ -492,7 +678,7 @@ impl AgentHarness { .filter(|name| { allowed_tools .as_ref() - .is_none_or(|allowed| allowed.is_empty() || allowed.contains(name)) + .is_none_or(|allowed| allowed.contains(name)) }) .collect::>() .join(", "); @@ -608,6 +794,33 @@ impl AgentHarness { message, ))); } + // Deferral (A2), after validation so an approver only ever sees a + // call the tool would actually accept, and before host authorization + // so a host's own gate is not consulted for a call a human has not + // yet approved. A call the resume path already approved + // (`RunContext::is_call_approved`) goes straight through. + if !ctx.is_call_approved(&call.id) { + let original = + ToolCall::new(call.id.clone(), call.name.clone(), model_arguments.clone()); + if crate::tool::is_external_tool(tool.as_ref()) { + ctx.limits.rollback_tool_calls(1); + return Ok(ResolvedToolCall::Deferred(DeferredRequest::external( + original, "external", None, + ))); + } + let policy = tool.policy(); + if policy.access.approval_required { + ctx.limits.rollback_tool_calls(1); + let metadata = serde_json::to_value(&policy.display) + .ok() + .filter(|value| value.as_object().is_some_and(|map| !map.is_empty())); + return Ok(ResolvedToolCall::Deferred(DeferredRequest::approval( + original, + "approval_required", + metadata, + ))); + } + } // Host authorization is deliberately last in admission: the gate sees // the raw provider arguments (including any forged hidden fields), // while execution receives the prepared trusted arguments. A hosted @@ -617,26 +830,18 @@ impl AgentHarness { let request = crate::host::ToolCallRequest::new( call.name.clone(), model_arguments, - binding.agent_id, + binding.agent_id.clone(), ) .with_call_id(CallId::new(call.id.clone())); - let cancellation = ctx.cancellation.clone(); let authorization = binding.host.security.authorize_tool(&request); - let decision = match self.call_budget(ctx) { - Some(remaining) => tokio::select! { - biased; - _ = cancellation.cancelled() => return Err(TinyAgentsError::Cancelled), - result = tokio::time::timeout(remaining, authorization) => result.map_err(|_| TinyAgentsError::Timeout(format!( + let decision = ctx + .bounded(self.call_budget(ctx), authorization, || { + format!( "tool authorization for run `{}` exceeded its remaining wall-clock budget", ctx.run_id() - )))?, - }, - None => tokio::select! { - biased; - _ = cancellation.cancelled() => return Err(TinyAgentsError::Cancelled), - result = authorization => result, - }, - }?; + ) + }) + .await?; if !decision.is_allowed() { let reason = decision .denial_reason() @@ -692,6 +897,7 @@ impl AgentHarness { PreparedToolCall { call_id, tool_name, + call: call.clone(), options, captured_input, started_at_ms, @@ -700,6 +906,86 @@ impl AgentHarness { } } + /// Records a tool-effect-ledger `started` row for `prepared` (B5), if a + /// ledger is attached to `ctx`. A no-op when [`RunContext::tool_effect_ledger`] + /// is `None`. + /// + /// Must be called *before* the tool actually executes, so a crash between + /// this write and the call settling is observable on resume. When the + /// write itself fails, [`RunContext::tool_effect_ledger_failure`] decides + /// whether that is fatal ([`LedgerFailure::Abort`], the default — the + /// caller must fail the call and propagate the error) or merely logged + /// ([`LedgerFailure::Continue`] — the call proceeds unrecorded). + async fn record_tool_effect_started( + &self, + ctx: &RunContext, + arguments: &Value, + prepared: &PreparedToolCall, + ) -> Result<()> { + let Some(ledger) = ctx.tool_effect_ledger.clone() else { + return Ok(()); + }; + let idempotency_key = tool_call_idempotency_key(&prepared.tool_name, arguments); + let start = ToolEffectStart { + run_id: ctx.run_id().clone(), + call_id: prepared.call_id.clone(), + tool: prepared.tool_name.clone(), + idempotency_key, + effect_summary: None, + }; + if let Err(err) = ledger.started(start).await { + return match ctx.tool_effect_ledger_failure { + LedgerFailure::Abort => Err(err), + LedgerFailure::Continue => { + tracing::warn!( + "[agent_loop::tools] tool-effect ledger `started` write failed for \ + call `{}` (tool `{}`): {err} — continuing per \ + LedgerFailure::Continue", + prepared.call_id.as_str(), + prepared.tool_name + ); + Ok(()) + } + }; + } + Ok(()) + } + + /// Records a tool-effect-ledger terminal row for `prepared` (B5), if a + /// ledger is attached to `ctx`. A no-op when [`RunContext::tool_effect_ledger`] + /// is `None`. + /// + /// Deliberately best-effort and never fatal: by the time this is called + /// the tool has already executed (or its execution future has already + /// failed), so aborting the run over a *settle* write failure would + /// discard a real result rather than merely skip recording one. A failed + /// settle write is logged; the row stays `started` and will surface again + /// from [`crate::tool::ToolEffectLedger::unresolved`] on the next resume. + async fn record_tool_effect_settled( + &self, + ctx: &RunContext, + prepared: &PreparedToolCall, + status: ToolEffectStatus, + ) { + let Some(ledger) = ctx.tool_effect_ledger.clone() else { + return; + }; + let settle = ToolEffectSettle { + run_id: ctx.run_id().clone(), + call_id: prepared.call_id.clone(), + status, + effect_summary: None, + }; + if let Err(err) = ledger.settled(settle).await { + tracing::warn!( + "[agent_loop::tools] tool-effect ledger `settled` write failed for call `{}` \ + (tool `{}`): {err}", + prepared.call_id.as_str(), + prepared.tool_name + ); + } + } + /// Terminal partner of [`AgentEvent::ToolStarted`] on the abort path: /// emits [`AgentEvent::ToolFailed`] and closes the call's `active_tool_calls` /// entry. @@ -720,7 +1006,7 @@ impl AgentHarness { ) { release_active_tool_call(status, call_id); let duration_ms = crate::ids::now_ms().saturating_sub(started_at_ms); - tinyagents_tracing::debug!( + tracing::debug!( "[agent_loop::tools] tool `{tool_name}` call `{}` failed after {duration_ms} ms: \ {error}", call_id.as_str() @@ -737,6 +1023,12 @@ impl AgentHarness { /// Fold phase for one completed call: the lifecycle `after_tool` hooks, /// accounting, the `ToolCompleted` emission, and the transcript append. + /// + /// Returns the result's `follow_up` content as a user message (B2), or + /// `None` when there is none. It is **not** appended here: a provider + /// requires every tool row of a batch to sit directly after the assistant + /// row that requested it, so the batch driver appends the follow-ups + /// only after its last tool row (see [`append_follow_ups`]). #[allow(clippy::too_many_arguments)] async fn finish_tool_call( &self, @@ -747,7 +1039,7 @@ impl AgentHarness { messages: &mut Vec, prepared: PreparedToolCall, mut result: tinytools::ToolResult, - ) -> Result<()> { + ) -> Result> { // Canonical ToolResult is intentionally correlation-free. The harness // owns `PreparedToolCall` and uses it below for transcript pairing, // events, and elapsed time; a tool cannot forge any of those fields. @@ -782,25 +1074,18 @@ impl AgentHarness { // value can reach the provider. if let Some(binding) = crate::runtime::host_invocation_binding::(ctx)? { let rendered = result.output_for_llm(prepared.options.prefer_markdown); - let cancellation = ctx.cancellation.clone(); let screening = binding .host .security .screen_input(&rendered, prepared.output_origin); - let screened = match self.call_budget(ctx) { - Some(remaining) => tokio::select! { - biased; - _ = cancellation.cancelled() => return Err(TinyAgentsError::Cancelled), - result = tokio::time::timeout(remaining, screening) => result.map_err(|_| TinyAgentsError::Timeout(format!( - "tool-output screening for run `{}` exceeded its remaining wall-clock budget", ctx.run_id() - )))?, - }, - None => tokio::select! { - biased; - _ = cancellation.cancelled() => return Err(TinyAgentsError::Cancelled), - result = screening => result, - }, - }; + let screened = ctx + .bounded(self.call_budget(ctx), screening, || { + format!( + "tool-output screening for run `{}` exceeded its remaining wall-clock budget", + ctx.run_id() + ) + }) + .await; match screened { Ok(crate::host::ScreenOutcome::Pass) => {} Ok(crate::host::ScreenOutcome::Redacted(text)) => { @@ -840,7 +1125,7 @@ impl AgentHarness { result.markdown_formatted = None; } } - tinyagents_tracing::debug!( + tracing::debug!( tool = %prepared.tool_name, agent = %binding.agent_id, ?outcome, @@ -848,10 +1133,64 @@ impl AgentHarness { ); } + // A tool's own `ToolControl` (`return_direct`/`terminate`/`goto`/ + // `state_update`) is the tool-vocabulary half of A1: it is *data* the + // tool returned, not a middleware decision, so it is translated into + // the same `MiddlewareControl` request a `Middleware` would make + // rather than a separate mechanism. `return_direct` and `terminate` + // both mean "the model never gets another turn": this call's own + // output becomes the run's final response, which — unlike + // `MiddlewareControl::StopWithFinal` — `JumpTo(End)` alone cannot + // express (it falls back to the *last assistant message*, which is + // one turn too early here), so the final response is set directly. + if let Some(control) = result.control.clone() { + // `return_direct` is now a per-call override (`Option`): + // `None` means "no opinion", so it falls back to the tool's own + // static `Tool::return_direct` default rather than being treated + // as `false`. + let return_direct = control.return_direct.unwrap_or_else(|| { + self.tools + .dispatch(&prepared.tool_name) + .map(|dispatch| dispatch.tool().return_direct()) + .unwrap_or(false) + }); + if return_direct || control.terminate { + run.final_response = Some(ModelResponse::assistant( + result.output_for_llm(prepared.options.prefer_markdown), + )); + ctx.request_control(MiddlewareControl::JumpTo(LoopTarget::End)); + } else if let Some(goto) = &control.goto { + match goto.as_str() { + "model" => ctx.request_control(MiddlewareControl::JumpTo(LoopTarget::Model)), + "tools" => ctx.request_control(MiddlewareControl::JumpTo(LoopTarget::Tools)), + "end" => ctx.request_control(MiddlewareControl::JumpTo(LoopTarget::End)), + other => tracing::debug!( + target: "tinyagents::agent_loop", + tool = %prepared.tool_name, + goto = other, + "[agent_loop] tool requested an unrecognized `goto` target; ignoring" + ), + } + } + if let Some(update) = control.state_update.clone() { + ctx.push_tool_state_update(update); + } + } + run.tool_calls += 1; if prepared.executed { run.executed_tools.push(prepared.tool_name.clone()); } + // Host-only metadata (B2): recorded on the run and on the event + // below, never rendered into the transcript row. + if let Some(metadata) = result.metadata.clone() { + run.tool_metadata + .push(crate::middleware::ToolResultMetadata { + call_id: prepared.call_id.clone(), + tool_name: prepared.tool_name.clone(), + metadata, + }); + } status.tool_calls = run.tool_calls; release_active_tool_call(status, &prepared.call_id); let model_output = result.output_for_llm(prepared.options.prefer_markdown); @@ -882,6 +1221,7 @@ impl AgentHarness { duration_ms: Some(duration_ms), output_bytes: Some(output_bytes), error, + metadata: result.metadata.clone(), }); crate::runtime::emit_host_progress::( ctx, @@ -903,7 +1243,7 @@ impl AgentHarness { &result, prepared.options, ))); - Ok(()) + Ok(follow_up_message(&result.follow_up)) } /// Executes requested tools one at a time (the historical semantics; used @@ -916,67 +1256,196 @@ impl AgentHarness { status: &mut HarnessRunStatus, messages: &mut Vec, tool_calls: Vec, - ) -> Result<()> { - for mut call in tool_calls { - let dispatch = match self.admit_tool_call(state, ctx, status, &mut call).await? { - ResolvedToolCall::Tool { dispatch, .. } => dispatch, - ResolvedToolCall::Answered(result) => { - self.recover_tool_call(state, ctx, run, status, messages, &call, result) - .await?; - continue; + ) -> Result { + let mut deferred = DeferredToolRequests::default(); + let mut follow_ups = Vec::new(); + for call in tool_calls { + follow_ups.extend( + self.execute_tool_serially(state, ctx, run, status, messages, call, &mut deferred) + .await?, + ); + } + append_follow_ups(messages, follow_ups); + Ok(deferred) + } + + /// Admits, executes, and folds **one** call on the serial path, recording + /// a deferral into `deferred` instead of answering it. + /// + /// Shared by [`Self::execute_tools_serially`] and the resume path + /// (`apply_deferred_results`), which re-runs an approved call through + /// exactly this pipeline so admission, the wrap onion, and the fold are + /// never duplicated. + /// + /// Returns the call's follow-up user message, if any, for the batch + /// driver to append after its last tool row (see [`append_follow_ups`]). + #[allow(clippy::too_many_arguments)] + pub(super) async fn execute_tool_serially( + &self, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + messages: &mut Vec, + mut call: ToolCall, + deferred: &mut DeferredToolRequests, + ) -> Result> { + let dispatch = match self.admit_tool_call(state, ctx, status, &mut call).await? { + ResolvedToolCall::Tool { dispatch, .. } => dispatch, + ResolvedToolCall::Answered(result) => { + return self + .recover_tool_call(state, ctx, run, status, messages, &call, result) + .await; + } + ResolvedToolCall::Deferred(request) => { + self.defer_tool_call(ctx, status, request, deferred); + return Ok(None); + } + }; + + let options = dispatch.call_options(&call.arguments); + let prepared = + self.start_tool_call(ctx, status, &call, options, true, dispatch.output_origin()); + if let Err(err) = self + .record_tool_effect_started(ctx, &call.arguments, &prepared) + .await + { + self.fail_tool_call( + ctx, + status, + &prepared.call_id, + &prepared.tool_name, + prepared.started_at_ms, + &err, + ); + return Err(err); + } + + // The real tool call is the innermost base of the tool-wrap + // onion (same before -> wrap -> after ordering as the model + // path): lifecycle `before_tool` ran in admission, the wrap onion + // runs here, and lifecycle `after_tool` runs in the fold. The + // crate-owned tool policy returns a recoverable tool error; the + // outer run budget still aborts when the whole run is exhausted. + let run_budget = self.call_budget(ctx); + let base = ToolCallBase { + dispatch, + options, + timeout_settings: self.tool_timeouts.clone(), + }; + let run_id = ctx.run_id().as_str().to_string(); + let fut = self.middleware.run_wrapped_tool(ctx, state, call, &base); + // TinyTools distinguishes a fatal execution `Err` from a + // recoverable `ToolResult::error`; no harness error-policy facade + // rewrites that canonical distinction. + let guarded = futures::FutureExt::map(fut, |result| { + result.map(|wrapped| wrapped.into_result_with_control()) + }); + let outcome = Self::with_call_budget( + run_budget, + &run_id, + "tool call", + super::model_call::RUN_BOUND_LABEL, + guarded, + ) + .await; + let (result, wrap_control) = match outcome { + Ok(pair) => pair, + Err(err) => { + if let Some(request) = execution_deferral(&prepared.call, &err) { + // Settled `Deferred` in the ledger (not left `started`): + // the call is genuinely paused pending external + // resolution, and `resume_deferred` is what eventually + // answers it — see `defer_started_tool_call`'s doc + // comment. + self.defer_started_tool_call(ctx, status, &prepared, request, deferred) + .await; + return Ok(None); } - }; + self.record_tool_effect_settled(ctx, &prepared, ToolEffectStatus::Failed) + .await; + self.fail_tool_call( + ctx, + status, + &prepared.call_id, + &prepared.tool_name, + prepared.started_at_ms, + &err, + ); + return Err(err); + } + }; + // A `ToolMiddleware::wrap_tool` that short-circuited with + // `MiddlewareToolOutcome::Command` carries no real result; queue + // its control the same way `run_wrapped_model`'s call site does + // (see the comment there). + if let Some(control) = wrap_control { + ctx.request_control(control); + } - let options = dispatch.call_options(&call.arguments); - let prepared = - self.start_tool_call(ctx, status, &call, options, true, dispatch.output_origin()); - - // The real tool call is the innermost base of the tool-wrap - // onion (same before -> wrap -> after ordering as the model - // path): lifecycle `before_tool` ran in admission, the wrap onion - // runs here, and lifecycle `after_tool` runs in the fold. The - // crate-owned tool policy returns a recoverable tool error; the - // outer run budget still aborts when the whole run is exhausted. - let run_budget = self.call_budget(ctx); - let base = ToolCallBase { - dispatch, - options, - timeout_settings: self.tool_timeouts.clone(), - }; - let run_id = ctx.run_id().as_str().to_string(); - let fut = self.middleware.run_wrapped_tool(ctx, state, call, &base); - // TinyTools distinguishes a fatal execution `Err` from a - // recoverable `ToolResult::error`; no harness error-policy facade - // rewrites that canonical distinction. - let guarded = - futures::FutureExt::map(fut, |result| result.map(|wrapped| wrapped.into_result())); - let outcome = Self::with_call_budget( - run_budget, - &run_id, - "tool call", - super::model_call::RUN_BOUND_LABEL, - guarded, - ) + self.record_tool_effect_settled(ctx, &prepared, ToolEffectStatus::Completed) .await; - let result = match outcome { - Ok(result) => result, - Err(err) => { - self.fail_tool_call( - ctx, - status, - &prepared.call_id, - &prepared.tool_name, - prepared.started_at_ms, - &err, - ); - return Err(err); - } - }; + self.finish_tool_call(state, ctx, run, status, messages, prepared, result) + .await + } - self.finish_tool_call(state, ctx, run, status, messages, prepared, result) - .await?; + /// Records a call deferred at admission (A2): emits `ToolDeferred` and + /// files the request under the right [`DeferredToolRequests`] list. The + /// admission slot was already released by `admit_tool_call`; the call is + /// re-admitted (and re-counted) if it is later approved. + fn defer_tool_call( + &self, + ctx: &RunContext, + status: &mut HarnessRunStatus, + request: DeferredRequest, + deferred: &mut DeferredToolRequests, + ) { + let call_id = CallId::new(request.call.id.clone()); + tracing::debug!( + "[agent_loop::tools] deferring call `{}` for `{}` ({})", + request.call.id, + request.call.name, + request.reason + ); + let record = ctx.emit(AgentEvent::ToolDeferred { + call_id: call_id.clone(), + reason: request.reason.to_string(), + }); + status.set_last_event(record.id); + if let Some(metadata) = request.metadata { + deferred.metadata.insert(call_id, metadata); } - Ok(()) + match request.kind { + DeferredKind::Approval => deferred.approvals.push(request.call), + DeferredKind::External => deferred.calls.push(request.call), + } + } + + /// Terminal partner of [`AgentEvent::ToolStarted`] for a call the *tool + /// itself* deferred mid-execution by raising `ApprovalRequired` / + /// `CallDeferred`: closes the in-flight entry, releases the tool-call + /// slot (the call never produced a result), settles its tool-effect-ledger + /// row as [`ToolEffectStatus::Deferred`], and files the request. + /// + /// Settling to `Deferred` (rather than leaving the row `started`) is + /// what keeps [`Self::reconcile_tool_effects`] — which only reconciles + /// rows still `started` — from mistaking this deliberate pause for a + /// crash artifact on a later resume. See that method's doc comment and + /// [`AgentHarness::resume_deferred`][crate::agent_loop::AgentHarness::resume_deferred] + /// for the full picture. + async fn defer_started_tool_call( + &self, + ctx: &mut RunContext, + status: &mut HarnessRunStatus, + prepared: &PreparedToolCall, + request: DeferredRequest, + deferred: &mut DeferredToolRequests, + ) { + release_active_tool_call(status, &prepared.call_id); + ctx.limits.rollback_tool_calls(1); + self.record_tool_effect_settled(ctx, prepared, ToolEffectStatus::Deferred) + .await; + self.defer_tool_call(ctx, status, request, deferred); } /// Answers a call that no tool ran — unknown tool, schema-invalid @@ -992,7 +1461,7 @@ impl AgentHarness { /// /// [err]: crate::tool::ToolResult::error #[allow(clippy::too_many_arguments)] - async fn recover_tool_call( + pub(super) async fn recover_tool_call( &self, state: &State, ctx: &mut RunContext, @@ -1001,8 +1470,8 @@ impl AgentHarness { messages: &mut Vec, call: &ToolCall, result: tinytools::ToolResult, - ) -> Result<()> { - tinyagents_tracing::debug!( + ) -> Result> { + tracing::debug!( "[agent_loop::tools] answering call `{}` for `{}` without executing a tool", call.id, call.name @@ -1032,7 +1501,8 @@ impl AgentHarness { status: &mut HarnessRunStatus, messages: &mut Vec, tool_calls: Vec, - ) -> Result<()> { + ) -> Result { + let mut deferred = DeferredToolRequests::default(); // Phase 1 — admission, serial, in call order. Nothing is announced and // nothing is queued here: an admission failure at call *k* must not // leave calls `0..k` with a `ToolStarted` they will never answer, nor @@ -1050,6 +1520,9 @@ impl AgentHarness { ResolvedToolCall::Answered(result) => { admitted.push(AdmittedCall::Recovered { call, result }) } + ResolvedToolCall::Deferred(request) => { + admitted.push(AdmittedCall::Deferred(request)) + } } } @@ -1072,6 +1545,10 @@ impl AgentHarness { slots.push(ToolSlot::Recovered { call, result }); continue; } + AdmittedCall::Deferred(request) => { + slots.push(ToolSlot::Deferred(request)); + continue; + } }; let options = dispatch.call_options(&call.arguments); @@ -1083,6 +1560,26 @@ impl AgentHarness { true, dispatch.output_origin(), )); + let just_prepared = prepared.last().expect("just pushed"); + if let Err(err) = self + .record_tool_effect_started(ctx, &call.arguments, just_prepared) + .await + { + // Every call in `prepared` so far (including this one) already + // emitted `ToolStarted`; give each one a terminal event before + // bailing, mirroring the sibling-abort handling in phase 4. + for sibling in &prepared { + self.fail_tool_call( + ctx, + status, + &sibling.call_id, + &sibling.tool_name, + sibling.started_at_ms, + &err, + ); + } + return Err(err); + } slots.push(ToolSlot::Execute); // Each call is bounded by its recoverable tool policy inside the @@ -1095,12 +1592,13 @@ impl AgentHarness { let run_budget = self.call_budget(ctx); let run_id = ctx.run_id().as_str().to_string(); futures.push(async move { - let fut = async move { - dispatch - .execute(state, call.arguments, options, parent_ctx) - .await - .map_err(map_tool_dispatch_error) - }; + let fut = execute_tool_recovering_model_retry(dispatch.execute( + state, + CallId::new(call.id), + call.arguments, + options, + parent_ctx, + )); let fut = Self::with_tool_policy_timeout(tool_timeout, timeout_result, fut); // As in serial mode, canonical execution errors remain fatal; // reported tool errors travel in `ToolResult::is_error`. @@ -1115,19 +1613,36 @@ impl AgentHarness { }); } - // Phase 3 — run all admitted calls concurrently. `join_all` preserves - // input order, so results pair 1:1 with `prepared`. - let results = futures::future::join_all(futures).await; + // Phase 3 — run all admitted calls concurrently, bounded by + // `RunLimits::max_tool_concurrency` when set (I-8). `buffered(n)` + // polls up to `n` futures at once and yields them **in input order** + // (unlike `buffer_unordered`), so results still pair 1:1 with + // `prepared` exactly as `join_all` (the unbounded case) did. + let concurrency = self + .policy + .limits + .max_tool_concurrency + .unwrap_or(futures.len().max(1)); + let results: Vec<_> = futures::stream::iter(futures) + .buffered(concurrency) + .collect() + .await; // Phase 4 — fold in original call order: the first call whose policy // kept its failure fatal (in that order) fails the turn; siblings // already ran to completion. let mut executed = prepared.into_iter().zip(results); + let mut follow_ups = Vec::new(); for slot in slots { match slot { ToolSlot::Recovered { call, result } => { - self.recover_tool_call(state, ctx, run, status, messages, &call, result) - .await?; + follow_ups.extend( + self.recover_tool_call(state, ctx, run, status, messages, &call, result) + .await?, + ); + } + ToolSlot::Deferred(request) => { + self.defer_tool_call(ctx, status, request, &mut deferred); } ToolSlot::Execute => { let (prepared, result) = executed @@ -1136,6 +1651,23 @@ impl AgentHarness { let result = match result { Ok(result) => result, Err(err) => { + if let Some(request) = execution_deferral(&prepared.call, &err) { + self.defer_started_tool_call( + ctx, + status, + &prepared, + request, + &mut deferred, + ) + .await; + continue; + } + self.record_tool_effect_settled( + ctx, + &prepared, + ToolEffectStatus::Failed, + ) + .await; self.fail_tool_call( ctx, status, @@ -1144,30 +1676,310 @@ impl AgentHarness { prepared.started_at_ms, &err, ); + // Every remaining `Execute` slot already emitted + // `ToolStarted` (phase 2) and is registered in + // `status.active_tool_calls`, but its future + // already resolved (phase 3 ran every future to + // completion via `join_all`) without ever getting + // a terminal event, because this fold stopped + // here. Give each of them one now so every + // `ToolStarted` still has exactly one terminal + // partner and no tool call is reported in-flight + // after the run has already failed. + let aborted = TinyAgentsError::Tool( + "aborted: sibling tool call failed".to_string(), + ); + for (sibling_prepared, _) in executed { + self.record_tool_effect_settled( + ctx, + &sibling_prepared, + ToolEffectStatus::Failed, + ) + .await; + self.fail_tool_call( + ctx, + status, + &sibling_prepared.call_id, + &sibling_prepared.tool_name, + sibling_prepared.started_at_ms, + &aborted, + ); + } return Err(err); } }; - self.finish_tool_call(state, ctx, run, status, messages, prepared, result) - .await?; + self.record_tool_effect_settled(ctx, &prepared, ToolEffectStatus::Completed) + .await; + follow_ups.extend( + self.finish_tool_call(state, ctx, run, status, messages, prepared, result) + .await?, + ); } } } - Ok(()) + append_follow_ups(messages, follow_ups); + Ok(deferred) + } +} + +/// Appends a batch's follow-up user messages (B2) after its last tool row, +/// in the calls' original order. +/// +/// One user message per call that returned `follow_up` content, rather than +/// one merged message: each keeps its own block list, and a provider that +/// merges adjacent user turns does so on the wire anyway. +pub(super) fn append_follow_ups(messages: &mut Vec, follow_ups: Vec) { + messages.extend(follow_ups); +} + +/// Builds the user message a result's `follow_up` blocks become (B2), or +/// `None` when the result has none. +/// +/// Block mapping — the same as the tool row's, except an image gets a real +/// [`ContentBlock::Image`] because a *user* message may carry one: +/// - `Text` → `Text`; `Json` → `Json`. +/// - `Image` → `Image(ImageRef)`: a URL as-is, inline bytes as a +/// `data:;base64,` URI; `mime_type` set from the block. +/// - `File` → `Text("[file ()]")` (the vendor +/// `ToolContent::render` placeholder), because the message model has no +/// file block yet; the host still has the full block on the event side if +/// it needs the bytes. +fn follow_up_message(follow_up: &[tinytools::ToolContent]) -> Option { + if follow_up.is_empty() { + return None; + } + let content = follow_up + .iter() + .map(|block| match block { + tinytools::ToolContent::Text { text } => ContentBlock::Text(text.clone()), + tinytools::ToolContent::Json { data } => ContentBlock::Json(data.clone()), + tinytools::ToolContent::Image { media_type, data } => { + let url = match data { + tinytools::ImageData::Url(url) => url.clone(), + tinytools::ImageData::Base64(bytes) => { + format!("data:{media_type};base64,{bytes}") + } + }; + ContentBlock::Image(tinyinference_llm::message::ImageRef { + url, + mime_type: Some(media_type.clone()), + }) + } + file @ tinytools::ToolContent::File { .. } => ContentBlock::Text(file.render()), + }) + .collect(); + Some(Message::User(tinyinference_llm::message::UserMessage { + content, + })) +} + +/// Turns an execution-time `ApprovalRequired`/`CallDeferred` (raised by the +/// tool through `Err`, and passed through [`execute_tool_recovering_model_retry`] +/// untouched) into the request the loop hands back, or `None` for any other +/// error. +fn execution_deferral(call: &ToolCall, error: &TinyAgentsError) -> Option { + match error { + TinyAgentsError::ApprovalRequired { metadata } => Some(DeferredRequest::approval( + call.clone(), + "approval_required", + Some(metadata.clone()), + )), + TinyAgentsError::CallDeferred { metadata } => Some(DeferredRequest::external( + call.clone(), + "call_deferred", + Some(metadata.clone()), + )), + _ => None, + } +} + +impl AgentHarness { + /// Reconciles unresolved tool-effect-ledger rows (B5) before resuming a + /// run from durable transcript `messages`. + /// + /// A crash between [`Self::record_tool_effect_started`] and the matching + /// settle write (or between the settle write and the tool result being + /// folded into `messages`) leaves a `started` row a resumed run must + /// resolve one way or another before it can safely feed `messages` back + /// into the loop: the last assistant turn may still carry a tool call + /// with no matching [`Message::Tool`] answer. + /// + /// For every tool call on the *last* assistant message that has no + /// [`Message::Tool`] answer yet **and** an unresolved (`started`) ledger + /// row, this consults the tool's declared + /// [`tinytools::ToolReplay`][tinytools::ToolPolicy::runtime]: + /// + /// - [`tinytools::ToolReplay::Safe`]: the call is left unanswered. + /// `messages` is not appended to for that call, so the normal loop + /// re-executes it exactly as it would a fresh call — the tool declared + /// this safe. + /// - [`tinytools::ToolReplay::Never`] (the default): a synthesized + /// tool-error result ("interrupted before settlement") is appended in + /// place of a real answer, the ledger row is settled as + /// [`crate::tool::ToolEffectStatus::Interrupted`], and the loop never + /// re-attempts the call. + /// + /// A call whose tool is no longer registered on this harness (renamed, + /// removed since the interrupted run) is treated as [`ToolReplay::Never`] + /// — fail closed rather than blindly re-run an unknown effect. + /// + /// Only ledger rows still in [`crate::tool::ToolEffectStatus::Started`] + /// are candidates: a call deferred mid-execution + /// (`ApprovalRequired`/`CallDeferred`) is settled as + /// [`crate::tool::ToolEffectStatus::Deferred`] by `defer_started_tool_call` + /// the moment it pauses, so [`crate::tool::ToolEffectLedger::unresolved`] + /// — which lists only `started` rows — never surfaces it here; a `Deferred` + /// row is exactly what [`AgentHarness::resume_deferred`] settles to + /// `Completed`/`Failed` once its answer runs. + /// + /// `excluded` is a second, defense-in-depth guard against the same + /// mistake: any call id in it is skipped even if its ledger row is + /// (unexpectedly) still `started` — e.g. the `Deferred` settle write + /// above failed and was only logged (settle writes are best-effort, see + /// [`Self::record_tool_effect_settled`]). [`AgentHarness::resume_deferred`] + /// passes the ids `results` is about to answer; any other caller — a host + /// reconciling a genuine crash, where no `results` exists at all — passes + /// an empty set. + /// + /// Returns the messages synthesized for `Never`-classified calls (already + /// appended to `messages` as well), so a caller that journals messages + /// separately from the in-memory transcript knows what changed. Returns + /// an empty `Vec` immediately, without any ledger I/O, when `ctx` has no + /// [`crate::tool::ToolEffectLedger`] attached or the transcript has no + /// pending tool calls. + pub async fn reconcile_tool_effects( + &self, + ctx: &RunContext, + run_id: &str, + messages: &mut Vec, + excluded: &std::collections::HashSet, + ) -> Result> { + let mut synthesized = Vec::new(); + let Some(ledger) = ctx.tool_effect_ledger.clone() else { + return Ok(synthesized); + }; + + // The calls a resumed run must judge are exactly the tool calls on + // the *last* assistant turn — any earlier assistant tool-call turn + // already has its answers folded in by definition, since the loop + // never advances past an unanswered turn. + let Some(pending_calls) = messages.iter().rev().find_map(|message| match message { + Message::Assistant(assistant) if !assistant.tool_calls.is_empty() => { + Some(assistant.tool_calls.clone()) + } + _ => None, + }) else { + return Ok(synthesized); + }; + let already_answered: std::collections::HashSet<&str> = messages + .iter() + .filter_map(|message| match message { + Message::Tool(tool_message) => Some(tool_message.tool_call_id.as_str()), + _ => None, + }) + .collect(); + let unanswered: Vec<&ToolCall> = pending_calls + .iter() + .filter(|call| { + !already_answered.contains(call.id.as_str()) + && !excluded.contains(&CallId::new(call.id.clone())) + }) + .collect(); + if unanswered.is_empty() { + return Ok(synthesized); + } + + let unresolved = ledger.unresolved(run_id).await?; + for call in unanswered { + let Some(effect) = unresolved.iter().find(|effect| effect.call_id == call.id) else { + // No ledger row for this call: nothing was ever journaled as + // started for it (e.g. a ledger was attached only after the + // interrupted attempt began), so there is nothing to + // reconcile — leave it for the loop to handle as it always + // has. + continue; + }; + let replay = self + .tools + .dispatch(&call.name) + .map(|dispatch| dispatch.tool().policy().runtime.replay) + .unwrap_or(tinytools::ToolReplay::Never); + match replay { + tinytools::ToolReplay::Safe => { + ctx.emit(AgentEvent::ToolEffectReconciled { + call_id: CallId::new(call.id.clone()), + action: "re_execute".to_string(), + }); + tracing::info!( + "[agent_loop::tools] reconciling unresolved tool effect for call `{}` \ + (tool `{}`) as ToolReplay::Safe — leaving unanswered for re-execution", + call.id, + call.name + ); + } + tinytools::ToolReplay::Never => { + let result = tinytools::ToolResult::error("interrupted before settlement"); + let tool_message = tool_message_from_result( + call.id.clone(), + &result, + ToolCallOptions::default(), + ); + messages.push(Message::Tool(tool_message.clone())); + synthesized.push(Message::Tool(tool_message)); + if let Err(err) = ledger + .settled(ToolEffectSettle { + run_id: crate::ids::RunId::new(run_id), + call_id: CallId::new(call.id.clone()), + status: ToolEffectStatus::Interrupted, + effect_summary: Some(effect.tool.clone()), + }) + .await + { + tracing::warn!( + "[agent_loop::tools] failed to settle interrupted tool effect for \ + call `{}` (tool `{}`): {err}", + call.id, + call.name + ); + } + ctx.emit(AgentEvent::ToolEffectReconciled { + call_id: CallId::new(call.id.clone()), + action: "interrupted".to_string(), + }); + tracing::info!( + "[agent_loop::tools] reconciled unresolved tool effect for call `{}` \ + (tool `{}`) as ToolReplay::Never — synthesized an interrupted result", + call.id, + call.name + ); + } + } + } + Ok(synthesized) } } /// Decides whether a batch may leave the serial path. /// -/// Lifecycle middleware runs during admission and can rewrite a call's name or -/// arguments. Until that mutable admission phase is made a separate completed -/// batch, any lifecycle middleware conservatively forces serial execution. +/// Lifecycle middleware used to force serial execution unconditionally +/// (`lifecycle_middleware == 0`), but that precondition never actually +/// applied: lifecycle `before_tool` hooks that can rewrite a call's name or +/// arguments run during **admission** (`admit_tool_call`, phase 1 of +/// [`AgentHarness::execute_tools_concurrently`]), which is already serial and +/// completes in full — for every call in the batch — before any concurrent +/// future is built. By the time phase 3 runs the futures, every call has its +/// final, lifecycle-rewritten name and arguments; there is nothing left for a +/// lifecycle middleware to still mutate concurrently (I-8). Tool-*wrap* +/// middleware (`tool_wrap_middleware`) is a separate concern: the concurrent +/// path drives each tool directly, bypassing the wrap onion entirely (see +/// that method's docs), so a registered `ToolMiddleware` still forces serial +/// execution — dropping it silently would skip the middleware. fn should_execute_tools_concurrently( calls: usize, canonical_parallel_safe: bool, - lifecycle_middleware: usize, tool_wrap_middleware: usize, ) -> bool { - calls > 1 && canonical_parallel_safe && lifecycle_middleware == 0 && tool_wrap_middleware == 0 + calls > 1 && canonical_parallel_safe && tool_wrap_middleware == 0 } /// A batch may leave the serial path only when every registered declaration @@ -1229,6 +2041,12 @@ fn tool_message_from_result( .map(|block| match block { tinytools::ToolContent::Text { text } => ContentBlock::Text(text.clone()), tinytools::ToolContent::Json { data } => ContentBlock::Json(data.clone()), + // Image/File blocks have no provider-neutral `ContentBlock` + // representation yet (see `docs/sdk-gaps/tools.md`); render the same + // short placeholder `ToolContent::render()` uses so a model + // still sees *something* rather than the block vanishing. + other @ (tinytools::ToolContent::Image { .. } + | tinytools::ToolContent::File { .. }) => ContentBlock::Text(other.render()), }) .collect() }; @@ -1252,13 +2070,76 @@ fn tool_message_from_result( /// Maps a canonical-dispatch failure back to the harness error surface. /// -/// Only cancellation and timeout retain their safe typed classifications. -/// Every other typed or foreign error is collapsed because message-bearing -/// errors can include credentials or user data exposed to model/event consumers. +/// Cancellation, timeout, and the structural errors that can escape a nested +/// sub-agent call ([`TinyAgentsError::SubAgentDepth`], +/// [`TinyAgentsError::LimitExceeded`]) keep their own typed classification. +/// Every other typed or foreign error is collapsed to a generic +/// [`TinyAgentsError::Tool`] because message-bearing errors from arbitrary +/// tool code can include credentials or user data exposed to model/event +/// consumers. +/// +/// Preserving the structural variants matters for retry correctness, not just +/// diagnostics: [`crate::retry::is_retryable`] treats every +/// [`TinyAgentsError::Tool`] as unconditionally retryable (arbitrary +/// tool-authored text has no shared vocabulary to classify against), but a +/// depth cap or run-limit violation is deterministic and will never succeed +/// on retry. Flattening `SubAgentDepth`/`LimitExceeded` into `Tool` made a +/// `RetryMiddleware` around tools re-run a permanently failing sub-agent call +/// until its attempt budget was exhausted (M-3). +/// Runs a dispatch call, folding a [`TinyAgentsError::ModelRetry`]/ +/// [`TinyAgentsError::ToolFailed`] the tool raised as `Err` into a +/// recoverable [`tinytools::ToolResult`] instead of aborting the run. +/// +/// This is A3's unified retry/failure vocabulary for tool errors: a tool that +/// wants "ask the model to try again" (the common case — a transient or +/// correctable failure) returns `Err(TinyAgentsError::ModelRetry(..).into())` +/// instead of `Ok(ToolResult::error(..))`, so it reads the same as any other +/// `?`-propagated failure in the tool's implementation while the harness +/// still folds it into the ordinary "tool ran, told the model to fix it" +/// transcript path (via [`tinytools::ToolResult::retry`]) rather than ending +/// the run. `ToolFailed` is the permanent counterpart +/// ([`tinytools::ToolResult::failed`]); every other error still maps through +/// [`map_tool_dispatch_error`] unchanged, preserving TinyTools' "`Err` aborts +/// the run" contract for genuine dispatch failures. +pub(super) async fn execute_tool_recovering_model_retry( + fut: Fut, +) -> Result +where + Fut: std::future::Future>, +{ + match fut.await { + Ok(result) => Ok(result), + Err(error) => match error.downcast::() { + Ok(TinyAgentsError::ModelRetry(message)) => Ok(tinytools::ToolResult::retry(message)), + Ok(TinyAgentsError::ToolFailed(message)) => Ok(tinytools::ToolResult::failed(message)), + // A2: a deferral is a typed signal for the loop, not a failure to + // redact. The metadata is host-only (never model-visible), so it + // is safe to carry through the wrap onion to the fold. + Ok( + deferral @ (TinyAgentsError::ApprovalRequired { .. } + | TinyAgentsError::CallDeferred { .. }), + ) => Err(deferral), + Ok(other) => Err(map_tool_dispatch_error(anyhow::Error::from(other))), + Err(error) => Err(map_tool_dispatch_error(error)), + }, + } +} + pub(super) fn map_tool_dispatch_error(error: anyhow::Error) -> TinyAgentsError { match error.downcast::() { Ok(TinyAgentsError::Cancelled) => TinyAgentsError::Cancelled, Ok(TinyAgentsError::Timeout(message)) => TinyAgentsError::Timeout(message), + Ok(TinyAgentsError::CallTimeout(message)) => TinyAgentsError::CallTimeout(message), + // `usize` carries no free-form content, so it is always safe to keep. + Ok(TinyAgentsError::SubAgentDepth(depth)) => TinyAgentsError::SubAgentDepth(depth), + // The message is harness-generated (a limit description), not + // attacker/tool-controlled, but is redacted anyway for the same + // "never assume a message is safe" posture as every other variant + // here; only the *classification* needs to survive for retry + // purposes. + Ok(TinyAgentsError::LimitExceeded(_)) => { + TinyAgentsError::LimitExceeded("tool dispatch hit a run limit".to_string()) + } Ok(_) => TinyAgentsError::Tool("tool dispatch failed".to_string()), Err(_) => TinyAgentsError::Tool("tool dispatch failed".to_string()), } @@ -1428,6 +2309,7 @@ mod canonical_result_tests { ], is_error: true, markdown_formatted: Some("## compact failure".to_string()), + ..ToolResult::default() } } @@ -1552,12 +2434,69 @@ mod canonical_result_tests { } #[test] - fn lifecycle_rewrite_of_a_safe_call_forces_the_serial_route() { - // `before_tool` receives `&mut ToolCall`, so a middleware may rewrite - // a raw-safe call into an unsafe tool/action. The loop consequently - // never selects its concurrent path while any lifecycle middleware is - // present, regardless of the pre-admission declaration result. - assert!(!should_execute_tools_concurrently(2, true, 1, 0)); - assert!(should_execute_tools_concurrently(2, true, 0, 0)); + fn lifecycle_middleware_no_longer_forces_the_serial_route() { + // Regression test (I-8): lifecycle middleware used to force the + // serial path unconditionally, on the theory that `before_tool` can + // rewrite a call (`&mut ToolCall`) while execution is concurrently in + // flight. That never actually applied: admission (including every + // `before_tool` hook) is serial and completes in full, for every call + // in the batch, before any concurrent future is built — so a + // lifecycle middleware has nothing left to mutate once execution + // starts. Only tool-*wrap* middleware (bypassed entirely by the + // concurrent path) still forces serial execution. + assert!(should_execute_tools_concurrently(2, true, 0)); + } + + #[test] + fn tool_wrap_middleware_still_forces_the_serial_route() { + // The concurrent path drives each tool directly, skipping the + // tool-wrap onion; a registered `ToolMiddleware` must still force + // serial execution or it would silently never run. + assert!(!should_execute_tools_concurrently(2, true, 1)); + } + + #[test] + fn map_tool_dispatch_error_preserves_sub_agent_depth_and_limit_exceeded() { + // M-3 regression: every non-cancel/timeout error used to collapse to + // a generic `Tool("tool dispatch failed")`, which `is_retryable` + // treats as unconditionally retryable. A `SubAgentDepth`/ + // `LimitExceeded` escaping a nested sub-agent tool call is + // deterministic and will never succeed on retry, so it must keep its + // own classification instead of masquerading as a retryable tool + // error. + let depth_err = anyhow::Error::from(TinyAgentsError::SubAgentDepth(4)); + assert!(matches!( + map_tool_dispatch_error(depth_err), + TinyAgentsError::SubAgentDepth(4) + )); + + let limit_err = anyhow::Error::from(TinyAgentsError::LimitExceeded( + "some sensitive detail".to_string(), + )); + match map_tool_dispatch_error(limit_err) { + TinyAgentsError::LimitExceeded(message) => { + assert!( + !message.contains("sensitive"), + "the original message must still be redacted: {message}" + ); + } + other => panic!("expected LimitExceeded, got {other:?}"), + } + } + + #[test] + fn map_tool_dispatch_error_still_redacts_a_genuine_tool_error() { + // An ordinary tool-authored error (arbitrary text, possibly carrying + // secrets or user data) must still be collapsed to a generic message, + // unlike the structural errors above. + let tool_err = anyhow::Error::from(TinyAgentsError::Model( + "leaked api key sk-secret".to_string(), + )); + match map_tool_dispatch_error(tool_err) { + TinyAgentsError::Tool(message) => { + assert!(!message.contains("sk-secret")); + } + other => panic!("expected Tool, got {other:?}"), + } } } diff --git a/crates/tinyagents-harness/src/agent_loop/types.rs b/crates/tinyagents-harness/src/agent_loop/types.rs index 0c115318..c3349d58 100644 --- a/crates/tinyagents-harness/src/agent_loop/types.rs +++ b/crates/tinyagents-harness/src/agent_loop/types.rs @@ -50,6 +50,30 @@ pub(crate) enum LoopExit { LimitStop(LimitKind), /// Steering latched a pause; the run is resumable, not finished. Paused(PauseState), + /// One or more tool calls in the last batch need a human decision or + /// host-side execution before the run can continue (A2). The transcript + /// keeps the assistant's tool-call row and every non-deferred sibling's + /// result; resume with + /// [`crate::runtime::AgentHarness::resume_deferred`]. + Deferred(crate::tool::DeferredToolRequests), +} + +/// The effect of draining a pending [`crate::context::MiddlewareControl`] at +/// one of the loop's safe checkpoints. +/// +/// Kept distinct from [`LoopExit`] because not every drained control ends the +/// run: [`crate::context::MiddlewareControl::JumpTo`]`(`[`crate::context::LoopTarget::Model`]`)` +/// must abandon the current turn (skip whatever the checkpoint's caller was +/// about to do next) without exiting the loop body, which a plain +/// `Option` cannot express. +#[derive(Clone, Debug)] +pub(crate) enum ControlEffect { + /// Nothing to do; the checkpoint's caller proceeds as it otherwise would. + None, + /// Abandon the rest of this turn and restart the loop body from the top. + ContinueLoop, + /// The run is done; propagate this [`LoopExit`] to the caller. + Exit(LoopExit), } /// The full result of an agent-loop invocation: the accumulated [`AgentRun`] diff --git a/crates/tinyagents-harness/src/artifacts/ops.rs b/crates/tinyagents-harness/src/artifacts/ops.rs index 7beec69b..8b73bb3f 100644 --- a/crates/tinyagents-harness/src/artifacts/ops.rs +++ b/crates/tinyagents-harness/src/artifacts/ops.rs @@ -125,7 +125,7 @@ pub fn note_artifact_handoff( paths: &[String], ) -> usize { for path in paths { - tinyagents_tracing::info!( + tracing::info!( stage = %stage, agent_id = %agent_id, task_id = %task_id, @@ -304,7 +304,7 @@ impl ArtifactOffload { redacted: stored.changed, }; - tinyagents_tracing::info!( + tracing::info!( agent_id = %self.agent_id, task_id = %self.task_id, kind = artifact.kind.as_str(), @@ -395,7 +395,7 @@ pub async fn offload_oversized_result( // credentials `write` just scrubbed out of the file. let abstract_text = build_abstract(&stored, ABSTRACT_BUDGET_CHARS); let pointer = render_artifact_pointer(&artifact, &abstract_text, read_tool); - tinyagents_tracing::info!( + tracing::info!( path = %artifact.relative_path, inline_bytes = output.len(), pointer_bytes = pointer.len(), @@ -405,7 +405,7 @@ pub async fn offload_oversized_result( (pointer, Some(artifact)) } Err(err) => { - tinyagents_tracing::warn!( + tracing::warn!( error = %err, inline_bytes = output.len(), threshold_bytes, diff --git a/crates/tinyagents-harness/src/blocking.rs b/crates/tinyagents-harness/src/blocking.rs new file mode 100644 index 00000000..097fd83c --- /dev/null +++ b/crates/tinyagents-harness/src/blocking.rs @@ -0,0 +1,70 @@ +//! Shared helper for running blocking (synchronous file/DB) work off the +//! tokio runtime. +//! +//! Several backends — [`crate::store::FileStore`], the JSONL append store, and +//! [`crate::cache::SqliteResponseCache`] under the `sqlite` feature — perform +//! blocking I/O (`std::fs::*`, `rusqlite` calls) that must never run directly +//! inside an `async fn` body, since that stalls whichever tokio worker thread +//! happens to poll it. [`run_blocking`] offloads the work via +//! `tokio::task::spawn_blocking` when a runtime is present, and falls back to +//! running it inline when there is none (e.g. a synchronous caller outside any +//! runtime, such as some test harnesses). + +use crate::error::{Result, TinyAgentsError}; + +/// Runs `work` off the async runtime via `spawn_blocking`, falling back to +/// running it inline when no tokio runtime is currently entered. +pub(crate) async fn run_blocking(work: F) -> Result +where + F: FnOnce() -> Result + Send + 'static, + T: Send + 'static, +{ + match tokio::runtime::Handle::try_current() { + Ok(handle) => handle + .spawn_blocking(work) + .await + .map_err(|e| TinyAgentsError::Validation(format!("blocking task error: {e}")))?, + Err(_) => work(), + } +} + +#[cfg(test)] +mod test { + use super::run_blocking; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::Duration; + + /// A slow, synchronous "store" body run through `run_blocking`. On a + /// `current_thread` runtime, a call that blocks the worker thread inline + /// (e.g. `std::thread::sleep` called directly in an `async fn`) would + /// starve every other task, including a concurrent timer. Routing it + /// through `run_blocking` must let the timer still fire while the slow + /// work is in flight (I-4). + #[tokio::test(flavor = "current_thread")] + async fn run_blocking_does_not_stall_the_runtime() { + let timer_fired = Arc::new(AtomicBool::new(false)); + let timer_fired_task = Arc::clone(&timer_fired); + + let timer = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(20)).await; + timer_fired_task.store(true, Ordering::SeqCst); + }); + + let slow_work = run_blocking(move || -> crate::error::Result<()> { + std::thread::sleep(Duration::from_millis(200)); + Ok(()) + }); + + // The slow blocking work and the short timer run concurrently; if + // blocking I/O were run inline on the current_thread runtime, the + // timer would never fire before `slow_work` completes because the + // single worker would be parked in `std::thread::sleep`. + slow_work.await.unwrap(); + assert!( + timer_fired.load(Ordering::SeqCst), + "concurrent timer should have fired while the blocking work ran off-thread" + ); + timer.await.unwrap(); + } +} diff --git a/crates/tinyagents-harness/src/cache/key.rs b/crates/tinyagents-harness/src/cache/key.rs index 6141ff7a..e0ee0135 100644 --- a/crates/tinyagents-harness/src/cache/key.rs +++ b/crates/tinyagents-harness/src/cache/key.rs @@ -297,13 +297,13 @@ pub fn apply_prompt_cache_breakpoints(request: &mut ModelRequest) -> bool { .get(PROMPT_CACHE_KEY_OPTION) .is_some() { - tinyagents_tracing::debug!( + tracing::debug!( "[cache] prompt_cache_key already set by caller; leaving provider_options untouched" ); return false; } let Some(derived) = prompt_cache_key(request) else { - tinyagents_tracing::debug!( + tracing::debug!( "[cache] protect_prompt_prefix is on but the request declares no cacheable prefix; \ no prompt_cache_key derived" ); @@ -318,6 +318,6 @@ pub fn apply_prompt_cache_breakpoints(request: &mut ModelRequest) -> bool { Value::String(derived.clone()), ); } - tinyagents_tracing::debug!(prompt_cache_key = %derived, "[cache] injected provider prompt-cache breakpoint"); + tracing::debug!(prompt_cache_key = %derived, "[cache] injected provider prompt-cache breakpoint"); true } diff --git a/crates/tinyagents-harness/src/cache/layout.rs b/crates/tinyagents-harness/src/cache/layout.rs index fef97334..c1c80c51 100644 --- a/crates/tinyagents-harness/src/cache/layout.rs +++ b/crates/tinyagents-harness/src/cache/layout.rs @@ -155,7 +155,7 @@ impl CacheLayoutEvent { } event.violates_policy = policy.protect_prompt_prefix; if event.violates_policy { - tinyagents_tracing::warn!( + tracing::warn!( content_only_change = event.content_only_change, before = ?event.segment_ids_before, after = ?event.segment_ids_after, diff --git a/crates/tinyagents-harness/src/cache/memory.rs b/crates/tinyagents-harness/src/cache/memory.rs index 38ffbfa2..48ae4bb5 100644 --- a/crates/tinyagents-harness/src/cache/memory.rs +++ b/crates/tinyagents-harness/src/cache/memory.rs @@ -110,7 +110,7 @@ impl LruResponseMap { }; self.remove(&victim); self.stats.evictions = self.stats.evictions.saturating_add(1); - tinyagents_tracing::trace!(key = %victim, "[cache] evicted least-recently-used entry"); + tracing::trace!(key = %victim, "[cache] evicted least-recently-used entry"); } } @@ -136,7 +136,7 @@ impl ResponseCache for InMemoryResponseCache { inner.stats.expirations = inner.stats.expirations.saturating_add(1); inner.stats.misses = inner.stats.misses.saturating_add(1); inner.sync_size_stats(); - tinyagents_tracing::debug!(key = %key, "[cache] entry expired; treating as miss"); + tracing::debug!(key = %key, "[cache] entry expired; treating as miss"); return Ok(None); } let hit = inner.data.get(key).map(|entry| entry.value.clone()); @@ -187,7 +187,7 @@ impl ResponseCache for InMemoryResponseCache { inner.order.clear(); inner.bytes = 0; inner.sync_size_stats(); - tinyagents_tracing::debug!(dropped, "[cache] cleared every in-memory response entry"); + tracing::debug!(dropped, "[cache] cleared every in-memory response entry"); Ok(()) } diff --git a/crates/tinyagents-harness/src/cache/singleflight.rs b/crates/tinyagents-harness/src/cache/singleflight.rs index ad39b52c..17f54e7c 100644 --- a/crates/tinyagents-harness/src/cache/singleflight.rs +++ b/crates/tinyagents-harness/src/cache/singleflight.rs @@ -126,22 +126,20 @@ impl SingleFlight { let Some(claim) = claim else { // A poisoned map must never take the run down: fall back to simply // making the call, which is the un-collapsed behaviour. - tinyagents_tracing::warn!( - "[cache] single-flight map poisoned; issuing the model call directly" - ); + tracing::warn!("[cache] single-flight map poisoned; issuing the model call directly"); return call().await.map(|response| (response, false)); }; let mut receiver = claim; // Follower: wait for the leader rather than duplicating the call. if let Some(receiver) = receiver.as_mut() { - tinyagents_tracing::debug!(key = %key, "[cache] joining an in-flight identical model call"); + tracing::debug!(key = %key, "[cache] joining an in-flight identical model call"); match receiver.recv().await { Ok(Outcome::Ready(response)) => return Ok((*response, true)), // Leader failed, or dropped the channel without sending (a // cancelled or panicking leader). Either way, run it ourselves. Ok(Outcome::Failed) | Err(_) => { - tinyagents_tracing::debug!( + tracing::debug!( key = %key, "[cache] in-flight leader did not produce a response; issuing our own call" ); diff --git a/crates/tinyagents-harness/src/cache/sqlite.rs b/crates/tinyagents-harness/src/cache/sqlite.rs index 27de7fee..e6318d02 100644 --- a/crates/tinyagents-harness/src/cache/sqlite.rs +++ b/crates/tinyagents-harness/src/cache/sqlite.rs @@ -61,8 +61,9 @@ fn sqlite_err(context: &str, err: impl std::fmt::Display) -> TinyAgentsError { fn now_ms() -> i64 { SystemTime::now() .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as i64) - .unwrap_or(0) + .ok() + .and_then(|d| i64::try_from(d.as_millis()).ok()) + .unwrap_or(i64::MAX) } impl SqliteResponseCache { @@ -94,7 +95,7 @@ impl SqliteResponseCache { /// not a correctness one. pub fn from_connection(conn: Connection) -> Result { if let Err(error) = conn.pragma_update(None, "journal_mode", "WAL") { - tinyagents_tracing::debug!(%error, "[cache] sqlite WAL unavailable; continuing with the default journal mode"); + tracing::debug!(%error, "[cache] sqlite WAL unavailable; continuing with the default journal mode"); } conn.execute_batch(SCHEMA) .map_err(|e| sqlite_err("create schema", e))?; @@ -135,32 +136,43 @@ impl SqliteResponseCache { #[async_trait] impl ResponseCache for SqliteResponseCache { async fn get(&self, key: &str) -> Result> { - let conn = self.lock()?; - let row: Option<(String, Option)> = conn - .query_row( - "SELECT value, expiry FROM response_cache WHERE ns = ?1 AND key = ?2", - params![self.namespace, key], - |row| Ok((row.get(0)?, row.get(1)?)), - ) - .optional() - .map_err(|e| sqlite_err("read entry", e))?; - let Some((value, expiry)) = row else { - return Ok(None); - }; - // Lazy expiry purge: a stale row is deleted on the way past, so a cache - // that is read but never written still sheds expired entries. - if expiry.is_some_and(|at| at <= now_ms()) { - conn.execute( - "DELETE FROM response_cache WHERE ns = ?1 AND key = ?2", - params![self.namespace, key], - ) - .map_err(|e| sqlite_err("purge expired entry", e))?; - tinyagents_tracing::debug!(key = %key, "[cache] sqlite entry expired; treating as miss"); - return Ok(None); - } - let response: ModelResponse = - serde_json::from_str(&value).map_err(|e| sqlite_err("decode entry", e))?; - Ok(Some(response)) + let conn = Arc::clone(&self.conn); + let namespace = self.namespace.clone(); + let key = key.to_string(); + // rusqlite is synchronous; run it off the tokio worker so a cache hit + // on the model hot path never stalls the runtime (see I-4). + crate::blocking::run_blocking(move || -> Result> { + let conn = conn + .lock() + .map_err(|_| sqlite_err("connection lock", "poisoned"))?; + let row: Option<(String, Option)> = conn + .query_row( + "SELECT value, expiry FROM response_cache WHERE ns = ?1 AND key = ?2", + params![namespace, key], + |row| Ok((row.get(0)?, row.get(1)?)), + ) + .optional() + .map_err(|e| sqlite_err("read entry", e))?; + let Some((value, expiry)) = row else { + return Ok(None); + }; + // Lazy expiry purge: a stale row is deleted on the way past, so a + // cache that is read but never written still sheds expired + // entries. + if expiry.is_some_and(|at| at <= now_ms()) { + conn.execute( + "DELETE FROM response_cache WHERE ns = ?1 AND key = ?2", + params![namespace, key], + ) + .map_err(|e| sqlite_err("purge expired entry", e))?; + tracing::debug!(key = %key, "[cache] sqlite entry expired; treating as miss"); + return Ok(None); + } + let response: ModelResponse = + serde_json::from_str(&value).map_err(|e| sqlite_err("decode entry", e))?; + Ok(Some(response)) + }) + .await } async fn put(&self, key: &str, value: ModelResponse) -> Result<()> { @@ -174,31 +186,49 @@ impl ResponseCache for SqliteResponseCache { ttl: Option, ) -> Result<()> { let encoded = serde_json::to_string(&value).map_err(|e| sqlite_err("encode entry", e))?; - let expiry = ttl.map(|ttl| now_ms().saturating_add(ttl.as_millis() as i64)); - let conn = self.lock()?; - conn.execute( - "INSERT OR REPLACE INTO response_cache (ns, key, value, expiry) \ - VALUES (?1, ?2, ?3, ?4)", - params![self.namespace, key, encoded, expiry], - ) - .map_err(|e| sqlite_err("write entry", e))?; - Ok(()) + let expiry = ttl.map(|ttl| { + let millis = i64::try_from(ttl.as_millis()).unwrap_or(i64::MAX); + now_ms().saturating_add(millis) + }); + let conn = Arc::clone(&self.conn); + let namespace = self.namespace.clone(); + let key = key.to_string(); + crate::blocking::run_blocking(move || -> Result<()> { + let conn = conn + .lock() + .map_err(|_| sqlite_err("connection lock", "poisoned"))?; + conn.execute( + "INSERT OR REPLACE INTO response_cache (ns, key, value, expiry) \ + VALUES (?1, ?2, ?3, ?4)", + params![namespace, key, encoded, expiry], + ) + .map_err(|e| sqlite_err("write entry", e))?; + Ok(()) + }) + .await } async fn clear(&self) -> Result<()> { - let conn = self.lock()?; - let dropped = conn - .execute( - "DELETE FROM response_cache WHERE ns = ?1", - params![self.namespace], - ) - .map_err(|e| sqlite_err("clear namespace", e))?; - tinyagents_tracing::debug!( - namespace = %self.namespace, - dropped, - "[cache] cleared the sqlite response cache namespace" - ); - Ok(()) + let conn = Arc::clone(&self.conn); + let namespace = self.namespace.clone(); + crate::blocking::run_blocking(move || -> Result<()> { + let conn = conn + .lock() + .map_err(|_| sqlite_err("connection lock", "poisoned"))?; + let dropped = conn + .execute( + "DELETE FROM response_cache WHERE ns = ?1", + params![namespace], + ) + .map_err(|e| sqlite_err("clear namespace", e))?; + tracing::debug!( + namespace = %namespace, + dropped, + "[cache] cleared the sqlite response cache namespace" + ); + Ok(()) + }) + .await } fn stats(&self) -> CacheStats { @@ -212,8 +242,8 @@ impl ResponseCache for SqliteResponseCache { ); match row { Ok((entries, bytes)) => CacheStats { - entries: entries.max(0) as u64, - bytes: bytes.max(0) as u64, + entries: u64::try_from(entries).unwrap_or(0), + bytes: u64::try_from(bytes).unwrap_or(0), ..CacheStats::default() }, Err(_) => CacheStats::default(), diff --git a/crates/tinyagents-harness/src/capability/mod.rs b/crates/tinyagents-harness/src/capability/mod.rs new file mode 100644 index 00000000..e7652b2e --- /dev/null +++ b/crates/tinyagents-harness/src/capability/mod.rs @@ -0,0 +1,190 @@ +//! Capability bundles (gap G3, `docs/runtime-comparison/plan.md`): +//! instructions, toolset, middleware, model defaults, exposure, and +//! `defer_loading` composed as one named unit instead of wired separately. +//! +//! `docs/runtime-comparison/pydantic-ai.md` §4 "Capabilities as the unit of +//! composition" is the design source: Pydantic AI's v2 `AbstractCapability` is +//! a bigger idea than middleware alone — it is what a "skill" or "plugin" is. +//! [`Capability`] is that bundle for TinyAgents, +//! [`crate::runtime::AgentHarness::with_capability`] is the harness-side +//! consumer, and [`CapabilityToolSet`]/[`LoadCapabilityTool`] implement the +//! `defer_loading` / `load_capability` on-demand loading mechanic. +//! +//! # Where this type lives, and why +//! +//! `Capability` composes [`crate::tool::toolset::ToolSet`] and [`crate::middleware::Middleware`] trait objects, both +//! native to this crate. It cannot live in `tinyagents-definition` (the +//! lower crate both `tinyagents-harness` and `tinyagents-registry` already +//! depend on): `tinyagents-definition` has zero dependency on +//! `tinyagents-harness` by design — that is what keeps the dependency graph +//! acyclic — so a definition-crate `Capability` could not name a `ToolSet` or +//! `Middleware` type without creating one. `tinyagents-registry` *does* +//! already depend on `tinyagents-harness`, so it can (and does) reference this +//! type — see `tinyagents_registry::CapabilityRegistry::register_capability` +//! — but `CapabilityRegistry` itself carries no `Ctx` type parameter +//! (none of its other stored kinds — models, tools, graphs, agents — need +//! one either, since they are all `Ctx`-free harness/tinytools types), so +//! registry storage for this specific, `Ctx`-generic bundle goes through +//! type-erased `Box` storage instead of adding a `Ctx` parameter to +//! the whole registry for one feature. See that method's doc comment for the +//! erasure mechanics. + +mod types; + +#[cfg(test)] +mod test; + +use std::sync::Arc; + +use serde_json::Value; +use tinytools::ToolExposure; + +use crate::error::{Result, TinyAgentsError}; +use crate::middleware::Middleware; +use crate::tool::toolset::ToolSet; + +pub use types::{ + Capability, CapabilityToolSet, LOAD_CAPABILITY_TOOL_NAME, LoadCapabilityTool, + ModelRequestDefaults, +}; +pub(crate) use types::{CapabilitySpec, ModelDefaultsSpec}; + +impl Capability { + /// Creates a capability with `name` and every optional field unset: + /// no instructions, no toolset, no middleware, no model defaults, + /// [`ToolExposure::Direct`] exposure, and `defer_loading: false`. + pub fn new(name: impl Into) -> Self { + Self { + name: name.into(), + instructions: None, + toolset: None, + middleware: Vec::new(), + model_defaults: None, + exposure: ToolExposure::Direct, + defer_loading: false, + } + } + + /// Sets the instructions contributed to the system prompt while this + /// capability is loaded. Returns `self` for chaining. + #[must_use] + pub fn with_instructions(mut self, instructions: impl Into) -> Self { + self.instructions = Some(instructions.into()); + self + } + + /// Sets the toolset this capability contributes. Returns `self` for + /// chaining. + #[must_use] + pub fn with_toolset(mut self, toolset: Arc>) -> Self { + self.toolset = Some(toolset); + self + } + + /// Appends one middleware instance, applied in declaration order. + /// Returns `self` for chaining. + #[must_use] + pub fn with_middleware(mut self, middleware: Arc>) -> Self { + self.middleware.push(middleware); + self + } + + /// Sets the model-request defaults applied when this capability is + /// installed. Returns `self` for chaining. + #[must_use] + pub fn with_model_defaults(mut self, defaults: ModelRequestDefaults) -> Self { + self.model_defaults = Some(defaults); + self + } + + /// Sets the [`ToolExposure`] applied to every tool this capability + /// contributes. Returns `self` for chaining. + #[must_use] + pub fn with_exposure(mut self, exposure: ToolExposure) -> Self { + self.exposure = exposure; + self + } + + /// Marks this capability as loaded on demand only, via + /// [`LOAD_CAPABILITY_TOOL_NAME`]. Returns `self` for chaining. + #[must_use] + pub fn with_defer_loading(mut self, defer_loading: bool) -> Self { + self.defer_loading = defer_loading; + self + } + + /// Builds a capability from a JSON spec: `{"name", "instructions"?, + /// "exposure"? ("direct"|"deferred"|"hidden", default "direct"), + /// "defer_loading"? (default `false`), "model_defaults"? + /// {"response_format"?, "fallback_models"?}}`. + /// + /// Only the declarative fields round-trip through JSON (see + /// [`Self::to_spec`]): the built capability's `toolset` and `middleware` + /// are always empty, since neither can be represented in JSON. A host + /// parsing a `.rag` `capability "name"` reference (or any other + /// JSON-declared capability) wires those in afterward with + /// [`Self::with_toolset`]/[`Self::with_middleware`] before installing it + /// via [`crate::runtime::AgentHarness::with_capability`]. + /// + /// # Errors + /// + /// Returns [`TinyAgentsError::Capability`] if `value` does not match the + /// spec shape, or if `name` is missing or blank. + pub fn from_spec(value: Value) -> Result { + let spec: CapabilitySpec = serde_json::from_value(value).map_err(|err| { + TinyAgentsError::Capability(format!("invalid capability spec: {err}")) + })?; + if spec.name.trim().is_empty() { + return Err(TinyAgentsError::Capability( + "capability spec is missing a non-blank `name`".to_string(), + )); + } + let model_defaults = spec.model_defaults.map(|defaults| ModelRequestDefaults { + default_response_format: defaults.response_format, + fallback: if defaults.fallback_models.is_empty() { + None + } else { + Some(crate::retry::FallbackPolicy { + models: defaults.fallback_models, + }) + }, + }); + Ok(Self { + name: spec.name, + instructions: spec.instructions, + toolset: None, + middleware: Vec::new(), + model_defaults, + exposure: spec.exposure.into(), + defer_loading: spec.defer_loading, + }) + } + + /// Renders this capability's declarative fields (name, instructions, + /// exposure, defer_loading, model defaults) as the JSON shape + /// [`Self::from_spec`] parses. The toolset and middleware are not + /// representable in JSON and are omitted; round-tripping a capability + /// through `to_spec`/`from_spec` therefore preserves every field except + /// those two. + pub fn to_spec(&self) -> Value { + let spec = CapabilitySpec { + name: self.name.clone(), + instructions: self.instructions.clone(), + exposure: self.exposure.into(), + defer_loading: self.defer_loading, + model_defaults: self + .model_defaults + .as_ref() + .map(|defaults| ModelDefaultsSpec { + response_format: defaults.default_response_format.clone(), + fallback_models: defaults + .fallback + .as_ref() + .map(|fallback| fallback.models.clone()) + .unwrap_or_default(), + }), + }; + serde_json::to_value(spec).expect("CapabilitySpec always serializes") + } +} diff --git a/crates/tinyagents-harness/src/capability/test.rs b/crates/tinyagents-harness/src/capability/test.rs new file mode 100644 index 00000000..8a1ad4ea --- /dev/null +++ b/crates/tinyagents-harness/src/capability/test.rs @@ -0,0 +1,265 @@ +//! Unit tests for the capability bundle (gap G3). + +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::json; +use tinytools::{Tool, ToolExposure, ToolResult}; + +use super::*; +use crate::context::{RunConfig, RunContext}; +use crate::middleware::Middleware; +use crate::runtime::AgentHarness; +use crate::tool::toolset::ToolSet; + +fn ctx() -> RunContext<()> { + RunContext::new(RunConfig::new("run-capability"), ()) +} + +/// A minimal deterministic tool, mirroring `tool::toolset::test::EchoTool`. +struct StubTool { + name: String, +} + +impl StubTool { + fn new(name: impl Into) -> Self { + Self { name: name.into() } + } +} + +#[async_trait] +impl Tool for StubTool { + fn name(&self) -> &str { + &self.name + } + + fn description(&self) -> &str { + "A stub tool for capability tests." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({"type": "object", "properties": {}}) + } + + async fn execute(&self, _args: serde_json::Value) -> anyhow::Result { + Ok(ToolResult::success(format!("{}-result", self.name))) + } +} + +/// A single-tool `ToolSet` wrapping one [`StubTool`]. +struct StubToolSet { + tool: Arc, +} + +impl StubToolSet { + fn new(name: impl Into) -> Self { + Self { + tool: Arc::new(StubTool::new(name)), + } + } +} + +#[async_trait] +impl ToolSet<(), ()> for StubToolSet { + async fn tools(&self, _ctx: &RunContext<()>) -> Result>> { + Ok(vec![self.tool.clone()]) + } + + async fn call( + &self, + name: &str, + args: serde_json::Value, + _ctx: &RunContext<()>, + ) -> Result { + if name == self.tool.name() { + self.tool + .execute(args) + .await + .map_err(|err| TinyAgentsError::Tool(err.to_string())) + } else { + Err(TinyAgentsError::ToolNotFound(name.to_string())) + } + } +} + +/// A no-op middleware; the tests only assert it reached the harness's stack. +struct StubMiddleware; + +#[async_trait] +impl Middleware<(), ()> for StubMiddleware { + fn name(&self) -> &str { + "stub-middleware" + } +} + +#[tokio::test] +async fn with_capability_installs_toolset_middleware_and_model_defaults() { + let capability: Capability<(), ()> = Capability::new("research") + .with_instructions("Use the research tool for lookups.") + .with_toolset(Arc::new(StubToolSet::new("lookup"))) + .with_middleware(Arc::new(StubMiddleware)) + .with_model_defaults(ModelRequestDefaults { + default_response_format: Some(tinyinference_llm::model::ResponseFormat::JsonObject), + fallback: None, + }); + + let mut harness: AgentHarness<(), ()> = AgentHarness::new(); + harness.with_capability(capability); + + // Toolset composition: the capability's tool is advertised. + let toolset = harness.toolset().expect("toolset installed").clone(); + let tools = toolset.tools(&ctx()).await.expect("tools resolve"); + assert!(tools.iter().any(|tool| tool.name() == "lookup")); + assert_eq!( + toolset.instructions().as_deref(), + Some("Use the research tool for lookups.") + ); + + // Middleware appended. + assert_eq!(harness.middleware().len(), 1); + + // Model defaults applied to the policy. + assert_eq!( + harness.policy().default_response_format, + Some(tinyinference_llm::model::ResponseFormat::JsonObject) + ); +} + +#[tokio::test] +async fn with_capability_composes_with_an_existing_toolset() { + let mut harness: AgentHarness<(), ()> = AgentHarness::new(); + harness.with_toolset(Arc::new(StubToolSet::new("base"))); + harness.with_capability( + Capability::new("extra").with_toolset(Arc::new(StubToolSet::new("extra-tool"))), + ); + + let toolset = harness.toolset().expect("toolset installed").clone(); + let names: Vec = toolset + .tools(&ctx()) + .await + .expect("tools resolve") + .iter() + .map(|tool| tool.name().to_string()) + .collect(); + assert!(names.iter().any(|name| name == "base")); + assert!(names.iter().any(|name| name == "extra-tool")); +} + +#[tokio::test] +async fn defer_loading_capability_hides_tools_and_instructions_until_loaded() { + let capability: Capability<(), ()> = Capability::new("advanced") + .with_instructions("Advanced instructions.") + .with_toolset(Arc::new(StubToolSet::new("advanced-tool"))) + .with_defer_loading(true); + let toolset = CapabilityToolSet::new(vec![capability]); + + let before = toolset.tools(&ctx()).await.expect("tools resolve"); + let before_names: Vec<&str> = before.iter().map(|tool| tool.name()).collect(); + assert!(!before_names.contains(&"advanced-tool")); + assert!(before_names.contains(&LOAD_CAPABILITY_TOOL_NAME)); + assert!(toolset.instructions().unwrap().contains("advanced")); + assert!(toolset.loaded_names().is_empty()); + + let result = toolset + .call( + LOAD_CAPABILITY_TOOL_NAME, + json!({"capability": "advanced"}), + &ctx(), + ) + .await + .expect("load_capability call succeeds"); + assert!(!result.is_error); + assert_eq!(toolset.loaded_names(), vec!["advanced".to_string()]); + + let after = toolset.tools(&ctx()).await.expect("tools resolve"); + let after_names: Vec<&str> = after.iter().map(|tool| tool.name()).collect(); + assert!(after_names.contains(&"advanced-tool")); + assert_eq!( + toolset.instructions().as_deref(), + Some("Advanced instructions.") + ); + + // The now-loaded tool is dispatchable through the composed toolset too. + let call_result = toolset + .call("advanced-tool", json!({}), &ctx()) + .await + .expect("dispatch succeeds"); + assert_eq!(call_result.text(), "advanced-tool-result"); +} + +#[tokio::test] +async fn load_capability_rejects_an_unknown_name() { + let capability: Capability<(), ()> = Capability::new("advanced").with_defer_loading(true); + let toolset = CapabilityToolSet::new(vec![capability]); + + let result = toolset + .call( + LOAD_CAPABILITY_TOOL_NAME, + json!({"capability": "nope"}), + &ctx(), + ) + .await + .expect("call resolves (a reported tool error, not a dispatch failure)"); + assert!(result.is_error); + assert!(toolset.loaded_names().is_empty()); +} + +#[tokio::test] +async fn no_load_capability_tool_when_nothing_defers() { + let capability: Capability<(), ()> = + Capability::new("eager").with_toolset(Arc::new(StubToolSet::new("eager-tool"))); + let toolset = CapabilityToolSet::new(vec![capability]); + + let tools = toolset.tools(&ctx()).await.expect("tools resolve"); + assert!( + !tools + .iter() + .any(|tool| tool.name() == LOAD_CAPABILITY_TOOL_NAME) + ); +} + +#[test] +fn capability_from_spec_round_trips_declarative_fields() { + let original: Capability<(), ()> = Capability::new("research") + .with_instructions("Use research tools.") + .with_exposure(ToolExposure::Deferred) + .with_defer_loading(true) + .with_model_defaults(ModelRequestDefaults { + default_response_format: Some(tinyinference_llm::model::ResponseFormat::JsonObject), + fallback: Some(crate::retry::FallbackPolicy { + models: vec!["primary".to_string(), "secondary".to_string()], + }), + }); + + let spec = original.to_spec(); + assert_eq!(spec["name"], "research"); + assert_eq!(spec["exposure"], "deferred"); + assert_eq!(spec["defer_loading"], true); + + let rebuilt: Capability<(), ()> = Capability::from_spec(spec).expect("spec parses"); + assert_eq!(rebuilt.name, original.name); + assert_eq!(rebuilt.instructions, original.instructions); + assert_eq!(rebuilt.exposure, original.exposure); + assert_eq!(rebuilt.defer_loading, original.defer_loading); + assert_eq!(rebuilt.model_defaults, original.model_defaults); + // Not representable in JSON: always empty on a spec-built capability. + assert!(rebuilt.toolset.is_none()); + assert!(rebuilt.middleware.is_empty()); +} + +#[test] +fn capability_from_spec_defaults_exposure_and_defer_loading() { + let capability: Capability<(), ()> = + Capability::from_spec(json!({"name": "minimal"})).expect("spec parses"); + assert_eq!(capability.name, "minimal"); + assert_eq!(capability.instructions, None); + assert_eq!(capability.exposure, ToolExposure::Direct); + assert!(!capability.defer_loading); + assert_eq!(capability.model_defaults, None); +} + +#[test] +fn capability_from_spec_rejects_a_blank_name() { + let err = Capability::<(), ()>::from_spec(json!({"name": " "})).unwrap_err(); + assert!(matches!(err, TinyAgentsError::Capability(_))); +} diff --git a/crates/tinyagents-harness/src/capability/types.rs b/crates/tinyagents-harness/src/capability/types.rs new file mode 100644 index 00000000..36894f97 --- /dev/null +++ b/crates/tinyagents-harness/src/capability/types.rs @@ -0,0 +1,450 @@ +//! Type definitions for the capability bundle module (gap G3, +//! `docs/runtime-comparison/plan.md`, `docs/runtime-comparison/pydantic-ai.md` +//! §4 "Capabilities as the unit of composition"). +//! +//! [`Capability`] is the bundle: instructions + a composable +//! [`crate::tool::toolset::ToolSet`] + middleware + model-request defaults + +//! [`ToolExposure`] + `defer_loading`, mirroring Pydantic AI's +//! `AbstractCapability`. [`CapabilityToolSet`] is the [`ToolSet`] adaptor that +//! makes a list of capabilities composable like any other toolset — +//! [`crate::runtime::AgentHarness::with_capability`] installs one — and +//! [`LoadCapabilityTool`] is the synthetic tool it auto-registers whenever any +//! capability declares `defer_loading: true`, so a model can bring a deferred +//! capability's tools and instructions into scope mid-run (Pydantic AI's +//! `load_capability`). + +use std::collections::HashSet; +use std::sync::{Arc, RwLock}; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use tinytools::{Tool, ToolExposure, ToolResult}; + +use crate::context::RunContext; +use crate::error::{Result, TinyAgentsError}; +use crate::middleware::Middleware; +use crate::retry::FallbackPolicy; +use crate::tool::toolset::{OverrideTool, ToolSet}; +use tinyinference_llm::model::ResponseFormat; + +/// Name of the synthetic tool [`CapabilityToolSet`] auto-registers whenever +/// at least one of its capabilities declares [`Capability::defer_loading`]. +pub const LOAD_CAPABILITY_TOOL_NAME: &str = "load_capability"; + +/// Default per-capability overrides applied to the harness's +/// [`crate::runtime::RunPolicy`] by [`crate::runtime::AgentHarness::with_capability`]. +/// +/// Deliberately narrower than the full [`crate::runtime::RunPolicy`]: only +/// the fields that are both meaningfully "this capability's preference" and +/// cheaply serializable (for [`Capability::to_spec`]/[`Capability::from_spec`]) +/// are included. [`crate::retry::RetryPolicy`] is not — it carries a +/// non-serializable predicate closure — so a capability wanting a custom +/// retry policy must be built programmatically with +/// [`Capability::with_middleware`]/[`crate::runtime::AgentHarness::with_policy`] +/// instead. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct ModelRequestDefaults { + /// Overrides [`crate::runtime::RunPolicy::default_response_format`] when + /// set. + pub default_response_format: Option, + /// Overrides [`crate::runtime::RunPolicy::fallback`] when set. + pub fallback: Option, +} + +impl ModelRequestDefaults { + /// Applies every set field onto `policy` in place, leaving the fields + /// this bundle leaves `None` untouched. + pub fn apply_to(&self, policy: &mut crate::runtime::RunPolicy) { + if let Some(format) = &self.default_response_format { + policy.default_response_format = Some(format.clone()); + } + if let Some(fallback) = &self.fallback { + policy.fallback = Some(fallback.clone()); + } + } +} + +/// A composable capability bundle (gap G3): the unit a `.rag` blueprint, +/// `AgentDefinition`, or host session references as one named thing instead +/// of wiring a toolset, middleware, and model defaults separately. +/// +/// Generic over the same `State`/`Ctx` pair as +/// [`crate::runtime::AgentHarness`] and +/// [`crate::tool::toolset::ToolSet`] — a `Capability` is meant to +/// be installed directly onto a harness via +/// [`crate::runtime::AgentHarness::with_capability`], not stored inside a +/// registry that has no `Ctx` dimension of its own (see +/// `tinyagents-registry`'s `CapabilityRegistry::register_capability`, which +/// type-erases this value through `Box` for exactly that reason). +pub struct Capability { + /// The capability's stable, unique name. + pub name: String, + /// Instructions this capability contributes to the system prompt while + /// loaded (immediately, unless [`Self::defer_loading`] is set). + pub instructions: Option, + /// The tools this capability contributes, if any. + pub toolset: Option>>, + /// Middleware appended to the harness's stack when this capability is + /// installed. + pub middleware: Vec>>, + /// Model-request defaults applied to the harness's [`crate::runtime::RunPolicy`] + /// when this capability is installed. + pub model_defaults: Option, + /// The [`ToolExposure`] applied uniformly to every tool + /// [`Self::toolset`] contributes, overriding each tool's own declared + /// exposure. Defaults to [`ToolExposure::Direct`]. + pub exposure: ToolExposure, + /// When `true`, this capability's tools and instructions are withheld + /// until a model calls [`LOAD_CAPABILITY_TOOL_NAME`] with this + /// capability's [`Self::name`] (Pydantic AI's `defer_loading`). + pub defer_loading: bool, +} + +impl Clone for Capability { + /// Manual `Clone`, not `#[derive(Clone)]`: a derive would add spurious + /// `State: Clone, Ctx: Clone` bounds even though neither type parameter + /// is stored by value here (only inside already-`Clone` `Arc`s). + fn clone(&self) -> Self { + Self { + name: self.name.clone(), + instructions: self.instructions.clone(), + toolset: self.toolset.clone(), + middleware: self.middleware.clone(), + model_defaults: self.model_defaults.clone(), + exposure: self.exposure, + defer_loading: self.defer_loading, + } + } +} + +impl std::fmt::Debug for Capability { + /// Renders the declarative fields; the toolset and middleware are opaque + /// trait objects, so only their presence/count is shown. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Capability") + .field("name", &self.name) + .field("instructions", &self.instructions) + .field("has_toolset", &self.toolset.is_some()) + .field("middleware_count", &self.middleware.len()) + .field("model_defaults", &self.model_defaults) + .field("exposure", &self.exposure) + .field("defer_loading", &self.defer_loading) + .finish() + } +} + +/// The JSON-facing shape [`Capability::from_spec`]/[`Capability::to_spec`] +/// round-trip. Only the declarative fields survive: a `toolset` and +/// `middleware` are live trait objects and cannot be represented in JSON, so +/// a capability built from a spec always has an empty toolset/middleware — +/// a host wanting either wires them in afterward with +/// [`Capability::with_toolset`]/[`Capability::with_middleware`]. +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub(crate) struct CapabilitySpec { + pub(crate) name: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) instructions: Option, + #[serde(default)] + pub(crate) exposure: ExposureSpec, + #[serde(default)] + pub(crate) defer_loading: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) model_defaults: Option, +} + +/// JSON-serializable mirror of [`ToolExposure`], which does not itself +/// derive `Serialize`/`Deserialize`. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum ExposureSpec { + #[default] + Direct, + Deferred, + Hidden, +} + +impl From for ToolExposure { + fn from(value: ExposureSpec) -> Self { + match value { + ExposureSpec::Direct => ToolExposure::Direct, + ExposureSpec::Deferred => ToolExposure::Deferred, + ExposureSpec::Hidden => ToolExposure::Hidden, + } + } +} + +impl From for ExposureSpec { + fn from(value: ToolExposure) -> Self { + match value { + ToolExposure::Direct => ExposureSpec::Direct, + ToolExposure::Deferred => ExposureSpec::Deferred, + ToolExposure::Hidden => ExposureSpec::Hidden, + } + } +} + +/// JSON-serializable mirror of [`ModelRequestDefaults`]: `fallback_models` +/// stands in for [`FallbackPolicy`] (which does not derive +/// `Serialize`/`Deserialize`). +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] +pub(crate) struct ModelDefaultsSpec { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) response_format: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub(crate) fallback_models: Vec, +} + +/// [`ToolSet`] adaptor composing a list of [`Capability`] bundles, gating a +/// `defer_loading` capability's tools and instructions behind a +/// [`LoadCapabilityTool`] call (gap G3). +/// +/// Installed by [`crate::runtime::AgentHarness::with_capability`], which +/// composes it with any toolset already installed via +/// [`crate::runtime::AgentHarness::with_toolset`]. The load state is shared +/// interior-mutable state (`Arc>>`) so a call to +/// [`LOAD_CAPABILITY_TOOL_NAME`] on one turn is visible to +/// [`ToolSet::tools`]/[`ToolSet::instructions`] on the very next turn — which +/// is exactly what lets the agent loop's existing tool-change diff (gap B6, +/// `crate::agent_loop::tool_changes`) pick up the change automatically: no +/// patch message needs to be hand-constructed here. +pub struct CapabilityToolSet { + pub(crate) capabilities: Vec>, + pub(crate) loaded: Arc>>, + pub(crate) load_tool: Option>, +} + +impl CapabilityToolSet { + /// Builds a toolset over `capabilities`. Every non-`defer_loading` + /// capability starts loaded; every `defer_loading` capability starts + /// unloaded and — since at least one is present — a + /// [`LoadCapabilityTool`] is auto-registered to bring it into scope. + pub fn new(capabilities: Vec>) -> Self { + let mut loaded_names = HashSet::new(); + let mut deferred_names = Vec::new(); + for capability in &capabilities { + if capability.defer_loading { + deferred_names.push(capability.name.clone()); + } else { + loaded_names.insert(capability.name.clone()); + } + } + let loaded = Arc::new(RwLock::new(loaded_names)); + let load_tool = if deferred_names.is_empty() { + None + } else { + deferred_names.sort(); + Some(Arc::new(LoadCapabilityTool::new( + deferred_names, + loaded.clone(), + ))) + }; + Self { + capabilities, + loaded, + load_tool, + } + } + + /// Names of every capability currently loaded (every non-deferred + /// capability, plus every deferred one a [`LoadCapabilityTool`] call has + /// loaded), sorted for deterministic assertions. + pub fn loaded_names(&self) -> Vec { + let mut names: Vec = self + .loaded + .read() + .expect("capability load state lock poisoned") + .iter() + .cloned() + .collect(); + names.sort(); + names + } + + fn is_loaded(&self, capability: &Capability, loaded: &HashSet) -> bool { + !capability.defer_loading || loaded.contains(&capability.name) + } + + /// The synthetic [`LoadCapabilityTool`] this toolset auto-registered, if + /// any capability declared `defer_loading: true`. `None` when every + /// capability loads eagerly. + pub fn load_tool(&self) -> Option> { + self.load_tool + .as_ref() + .map(|tool| tool.clone() as Arc) + } +} + +#[async_trait] +impl ToolSet for CapabilityToolSet { + async fn tools(&self, ctx: &RunContext) -> Result>> { + let loaded = self + .loaded + .read() + .expect("capability load state lock poisoned") + .clone(); + let mut out = Vec::new(); + for capability in &self.capabilities { + if !self.is_loaded(capability, &loaded) { + continue; + } + if let Some(toolset) = &capability.toolset { + for tool in toolset.tools(ctx).await? { + out.push( + Arc::new(OverrideTool::new(tool).with_exposure(capability.exposure)) + as Arc, + ); + } + } + } + if let Some(load_tool) = &self.load_tool { + out.push(load_tool.clone() as Arc); + } + Ok(out) + } + + async fn call(&self, name: &str, args: Value, ctx: &RunContext) -> Result { + if let Some(load_tool) = &self.load_tool + && name == LOAD_CAPABILITY_TOOL_NAME + { + return load_tool + .execute(args) + .await + .map_err(|err| TinyAgentsError::Tool(err.to_string())); + } + let loaded = self + .loaded + .read() + .expect("capability load state lock poisoned") + .clone(); + for capability in &self.capabilities { + if !self.is_loaded(capability, &loaded) { + continue; + } + let Some(toolset) = &capability.toolset else { + continue; + }; + let owns = toolset + .tools(ctx) + .await? + .iter() + .any(|tool| tool.name() == name); + if owns { + return toolset.call(name, args, ctx).await; + } + } + Err(TinyAgentsError::ToolNotFound(name.to_string())) + } + + fn instructions(&self) -> Option { + let loaded = self + .loaded + .read() + .expect("capability load state lock poisoned") + .clone(); + let mut parts = Vec::new(); + let mut pending = Vec::new(); + for capability in &self.capabilities { + if self.is_loaded(capability, &loaded) { + if let Some(instructions) = &capability.instructions { + parts.push(instructions.clone()); + } + } else { + pending.push(capability.name.clone()); + } + } + if !pending.is_empty() { + pending.sort(); + parts.push(format!( + "Additional capabilities are available via `{LOAD_CAPABILITY_TOOL_NAME}`: {}.", + pending.join(", ") + )); + } + if parts.is_empty() { + None + } else { + Some(parts.join("\n\n")) + } + } + + async fn for_run(&self, ctx: &RunContext) -> Result<()> { + let loaded = self + .loaded + .read() + .expect("capability load state lock poisoned") + .clone(); + for capability in &self.capabilities { + if !self.is_loaded(capability, &loaded) { + continue; + } + if let Some(toolset) = &capability.toolset { + toolset.for_run(ctx).await?; + } + } + Ok(()) + } +} + +/// The synthetic tool [`CapabilityToolSet`] auto-registers whenever at least +/// one of its capabilities declares `defer_loading: true`. +/// +/// Calling it with a known deferred capability name marks that capability +/// loaded in the shared state every [`CapabilityToolSet`] method reads, so +/// the very next turn's [`ToolSet::tools`]/[`ToolSet::instructions`] reflect +/// it — which is what lets the agent loop's existing per-turn tool-change +/// diff (gap B6) record the change as an ordinary transcript patch, with no +/// bespoke wiring needed here. +pub struct LoadCapabilityTool { + pub(crate) available: Vec, + pub(crate) loaded: Arc>>, +} + +impl LoadCapabilityTool { + pub(crate) fn new(available: Vec, loaded: Arc>>) -> Self { + Self { available, loaded } + } +} + +#[async_trait] +impl Tool for LoadCapabilityTool { + fn name(&self) -> &str { + LOAD_CAPABILITY_TOOL_NAME + } + + fn description(&self) -> &str { + "Loads a deferred capability bundle by name, making its tools and \ + instructions available for the rest of this run." + } + + fn parameters_schema(&self) -> Value { + json!({ + "type": "object", + "properties": { + "capability": { + "type": "string", + "description": "The deferred capability name to load.", + "enum": self.available, + } + }, + "required": ["capability"], + }) + } + + async fn execute(&self, args: Value) -> anyhow::Result { + let Some(name) = args.get("capability").and_then(Value::as_str) else { + return Ok(ToolResult::error( + "`capability` argument is required".to_string(), + )); + }; + if !self.available.iter().any(|available| available == name) { + return Ok(ToolResult::error(format!( + "unknown deferred capability `{name}`" + ))); + } + self.loaded + .write() + .expect("capability load state lock poisoned") + .insert(name.to_string()); + Ok(ToolResult::success(format!("capability `{name}` loaded"))) + } +} diff --git a/crates/tinyagents-harness/src/context/mod.rs b/crates/tinyagents-harness/src/context/mod.rs index 88877464..c501c272 100644 --- a/crates/tinyagents-harness/src/context/mod.rs +++ b/crates/tinyagents-harness/src/context/mod.rs @@ -302,46 +302,153 @@ impl RunContext { config, data, stores: StoreRegistry::new(), + namespaced_store: None, + state_view: None, events, limits, steering: None, + run_queue: None, cancellation: CancellationToken::new(), control: std::sync::Arc::new(std::sync::Mutex::new(None)), + state_updates: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())), + tool_state_updates: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())), workspace: None, on_error_dispatched: false, streaming: false, host_agent_id: None, host_authority: None, terminal_observer: None, + active_model_call: None, + deferred_results: None, + approved_calls: std::collections::HashSet::new(), + child_ordinal: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)), + tool_effect_ledger: None, + tool_effect_ledger_failure: crate::tool::LedgerFailure::default(), + compaction_sink: None, } } - /// Builds an isolated child context from this live parent context. + /// Attaches the resolutions for the deferred tool calls this run resumes + /// (A2). The agent loop applies them to the unanswered tool calls on the + /// transcript's last assistant row before making its next model call. + /// Prefer [`crate::runtime::AgentHarness::resume_deferred`], which does + /// this for you. + #[must_use] + pub fn with_deferred_results(mut self, results: crate::tool::DeferredToolResults) -> Self { + self.deferred_results = Some(results); + self + } + + /// Takes the pending deferred-call resolutions, if any (A2). + pub(crate) fn take_deferred_results(&mut self) -> Option { + self.deferred_results.take() + } + + /// Whether a human approved the tool call `call_id` on resume (A2). + /// + /// The agent loop consults this to skip its own deferral checks for an + /// approved call; an approval gate implemented as a `before_tool` + /// middleware should consult it too so it does not re-defer a call the + /// human already decided on. + pub fn is_call_approved(&self, call_id: &str) -> bool { + self.approved_calls.contains(call_id) + } + + /// Marks `call_id` as approved for this run (A2). + pub(crate) fn mark_call_approved(&mut self, call_id: impl Into) { + self.approved_calls.insert(call_id.into()); + } + + /// Returns the next value from this context's own child-ordinal counter + /// (starting at `0`), advancing it. + /// + /// The counter is per-context, not process-global: a freshly constructed + /// context (including a child context, which never inherits its parent's + /// counter) always starts at `0`. Callers that spawn deterministically + /// named children — [`crate::subagent::SubAgent`], for one — use this + /// instead of a process-wide sequence so two processes calling the same + /// parent context's child spawner in the same order derive identical + /// ordinals, and therefore identical child run ids (M-2). + pub fn next_child_ordinal(&self) -> u64 { + self.child_ordinal + .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + } + + /// Builds an isolated child context from this live parent context, + /// propagating the parent's host authority. /// /// A child gets a new run id, lineage record, [`LimitTracker`], control /// slot, and instance id. It deliberately shares the capabilities that /// describe one recursive operation: cancellation, events, stores, /// workspace policy, steering, streaming mode, thread identity, output - /// cap, and depth cap. Metadata is shallow-merged automatically: any key + /// cap, and depth cap. Metadata is shallow-merged automatically: any key /// set on `child_config.metadata` overlays the parent's metadata object /// (see `shallow_merge_metadata`), so callers only need to pass the /// child-specific keys. - pub fn child( + /// + /// This keeps the child's `Ctx` type identical to the parent's, which is + /// what makes propagating [`Self::host_authority`] sound: the type-erased + /// authority installed by a hosted invocation is keyed to the exact + /// `(State, Ctx)` pair it was constructed for, and this method is the only + /// place that carries it forward. A recursive call that needs a + /// *different* `Ctx` type must go through [`Self::child_with_data`] + /// instead, which never propagates host authority. + pub fn child(&self, child_config: RunConfig, data: Ctx) -> Result> { + let mut child = self.child_without_authority(child_config, data)?; + child.host_authority = self.host_authority.clone(); + Ok(child) + } + + /// Builds an isolated child context whose user data type may differ from + /// this context's, deliberately *not* propagating host authority. + /// + /// Use this whenever the child's `Ctx` differs from the parent's (for + /// example, a differently-typed sub-harness). Because [`RunContext`] does + /// not track its `State` type parameter at all, and the erased host + /// authority is keyed to a specific `(State, Ctx)` pair, there is no sound + /// way to check at this boundary whether the parent's authority would + /// still apply to the child's types. Rather than guess, the child simply + /// starts unhosted; a caller that legitimately needs to delegate hosted + /// authority across a `Ctx` change must do so explicitly through the + /// hosted subagent entry points, which re-derive authority from the live + /// host capability bundle rather than reinterpreting the parent's. + pub fn child_with_data( + &self, + child_config: RunConfig, + data: ChildCtx, + ) -> Result> { + self.child_without_authority(child_config, data) + } + + fn child_without_authority( &self, child_config: RunConfig, data: ChildCtx, ) -> Result> { let mut config = self.config.child(child_config)?; config.metadata = shallow_merge_metadata(&self.config.metadata, config.metadata); + let child_run_id = config.run_id.clone(); + // Derive a per-child handle (not a bare clone): it shares the parent's + // queue/policy but only drains commands addressed to *this* child's + // run id, `SteeringTarget::Root`-addressed commands stay with the + // parent, and its pause/checkpoint state is its own (I-5). + let steering = self + .steering + .as_ref() + .map(|handle| handle.for_child(child_run_id)); let mut child = RunContext::new(config, data) .with_stores(self.stores.clone()) + .with_optional_namespaced_store(self.namespaced_store.clone()) + .with_optional_state_view(self.state_view.clone()) .with_events(self.events.clone()) .with_cancellation(self.cancellation.clone()) - .with_optional_steering(self.steering.clone()) + .with_optional_steering(steering) .with_optional_workspace(self.workspace.clone()) .with_streaming(self.streaming); child.host_agent_id = self.host_agent_id.clone(); - child.host_authority = self.host_authority.clone(); + child.tool_effect_ledger = self.tool_effect_ledger.clone(); + child.tool_effect_ledger_failure = self.tool_effect_ledger_failure; + child.compaction_sink = self.compaction_sink.clone(); Ok(child) } @@ -417,7 +524,14 @@ impl RunContext { /// request. This gives competing middleware layers a deterministic outcome /// instead of last-writer-wins — e.g. a pause request is never downgraded to /// a stop by a later, weaker request. + /// + /// [`MiddlewareControl::Continue`] is never installed: it carries no + /// instruction, so requesting it is a no-op regardless of what (if + /// anything) is already pending. pub fn request_control(&self, control: MiddlewareControl) { + if matches!(control, MiddlewareControl::Continue) { + return; + } if let Ok(mut guard) = self.control.lock() { let replace = match guard.as_ref() { Some(existing) => control.precedence() > existing.precedence(), @@ -434,6 +548,56 @@ impl RunContext { self.control.lock().ok().and_then(|mut guard| guard.take()) } + /// Queues a [`StateUpdate`] for the host to apply. + /// + /// The agent loop only ever holds `state: &State` (a shared reference), so + /// [`MiddlewareControl::UpdateState`] cannot be applied in place; the loop + /// pushes it here instead of discarding it. Called by + /// [`crate::agent_loop`]'s control-checkpoint handling; a host drains the + /// queue with [`Self::take_state_updates`] and applies each update against + /// its own `&mut State` between runs (or between turns, via its own + /// checkpoint). + pub fn push_state_update(&self, update: StateUpdate) { + if let Ok(mut guard) = self.state_updates.lock() { + guard.push(update); + } + } + + /// Drains every [`StateUpdate`] queued so far, in request order. + pub fn take_state_updates(&self) -> Vec { + self.state_updates + .lock() + .ok() + .map(|mut guard| std::mem::take(&mut *guard)) + .unwrap_or_default() + } + + /// Queues a raw JSON state update a tool requested via + /// [`tinytools::ToolControl::state_update`][tc]. + /// + /// A canonical tool has no access to the harness's typed `State`, so its + /// state update travels as `serde_json::Value` rather than a + /// [`StateUpdate`] closure. Kept as a separate queue (not merged into + /// [`Self::push_state_update`]) so a host can tell a middleware-originated + /// typed update from a tool-originated JSON one without downcasting. + /// + /// [tc]: tinytools::ToolControl::state_update + pub fn push_tool_state_update(&self, update: serde_json::Value) { + if let Ok(mut guard) = self.tool_state_updates.lock() { + guard.push(update); + } + } + + /// Drains every raw JSON tool state update queued so far, in request + /// order. See [`Self::push_tool_state_update`]. + pub fn take_tool_state_updates(&self) -> Vec { + self.tool_state_updates + .lock() + .ok() + .map(|mut guard| std::mem::take(&mut *guard)) + .unwrap_or_default() + } + /// Attaches a [`CancellationToken`] so an orchestrator can request that this /// run stop cooperatively at its next safe checkpoint. /// @@ -454,6 +618,48 @@ impl RunContext { self } + /// Attaches the hierarchical store every tool in this run receives as + /// [`ToolExecutionContext::store`][crate::tool::ToolExecutionContext::store] + /// (B1). See [`RunContext::namespaced_store`]. + #[must_use] + pub fn with_namespaced_store( + mut self, + store: std::sync::Arc, + ) -> Self { + self.namespaced_store = Some(store); + self + } + + fn with_optional_namespaced_store( + mut self, + store: Option>, + ) -> Self { + self.namespaced_store = store; + self + } + + /// Attaches a read-only snapshot of the application state every tool in + /// this run can recover with + /// [`ToolExecutionContext::state::()`][crate::tool::ToolExecutionContext::state] + /// (B1). See [`RunContext::state_view`] for why this is an owned `Arc` + /// the host supplies rather than the loop's own `&State`. + #[must_use] + pub fn with_state_view( + mut self, + state: std::sync::Arc, + ) -> Self { + self.state_view = Some(state); + self + } + + fn with_optional_state_view( + mut self, + state: Option>, + ) -> Self { + self.state_view = state; + self + } + /// Replaces the event sink with a (possibly shared) `events`. pub fn with_events(mut self, events: EventSink) -> Self { self.events = events; @@ -466,8 +672,65 @@ impl RunContext { /// The agent loop drains the handle before each model call via /// [`crate::steering::apply_pending_steering`]. Without this the /// run accepts no steering. + /// + /// Binds the handle to this run's id as the **root** of its steering tree + /// (see [`crate::steering::SteeringTarget::Root`]); a child run created + /// from this context via [`Self::child`]/[`Self::child_with_data`] gets a + /// derived handle scoped to its own id instead of sharing this binding + /// (I-5). pub fn with_steering(mut self, steering: crate::steering::SteeringHandle) -> Self { - self.steering = Some(steering); + let root_run_id = self.lineage().root_run_id.clone(); + self.steering = Some(steering.bind_root(root_run_id)); + self + } + + /// Attaches a [`crate::run_queue::RunQueueHandle`] so messages pushed + /// from outside the run reach the transcript at the loop's safe turn + /// boundaries (A4). See [`RunContext::run_queue`] for the drain points + /// and [`crate::runtime::RunPolicy::queue_mode`] for how many items each + /// boundary takes. Without this the loop consumes no queued messages. + pub fn with_run_queue(mut self, queue: crate::run_queue::RunQueueHandle) -> Self { + self.run_queue = Some(queue); + self + } + + /// Attaches a durable tool-effect ledger (B5), so the agent loop writes a + /// `started` row before each tool call executes and a `completed`/ + /// `failed` row after it settles. `None` (the default) disables all + /// ledger writes. + /// + /// See [`crate::tool::ToolEffectLedger`] and + /// [`RunContext::with_tool_effect_ledger_failure`] for how a `started` + /// write failure is handled. + #[must_use] + pub fn with_tool_effect_ledger( + mut self, + ledger: std::sync::Arc, + ) -> Self { + self.tool_effect_ledger = Some(ledger); + self + } + + /// Sets how the agent loop reacts when [`crate::tool::ToolEffectLedger::started`] + /// itself fails. Defaults to [`crate::tool::LedgerFailure::Abort`]. + #[must_use] + pub fn with_tool_effect_ledger_failure(mut self, failure: crate::tool::LedgerFailure) -> Self { + self.tool_effect_ledger_failure = failure; + self + } + + /// Attaches a durable [`crate::summarization::CompactionSink`] so every + /// [`crate::summarization::CompactionRecord`] a compaction produces on + /// this run is persisted (typically into a session's entry tree), + /// instead of only living in + /// [`crate::middleware::ContextCompressionMiddleware::records`]'s + /// in-process buffer. `None` (the default) disables persistence. + #[must_use] + pub fn with_compaction_sink( + mut self, + sink: std::sync::Arc, + ) -> Self { + self.compaction_sink = Some(sink); self } @@ -491,6 +754,48 @@ impl RunContext { self.config.depth() } + /// Races `fut` against this run's cooperative cancellation and, when + /// `deadline` is `Some`, a wall-clock timeout — the one home for the + /// `tokio::select! { biased; _ = cancelled() => .., _ = timeout(remaining, + /// fut) => .. }` pattern that used to be copied at every host/provider I-O + /// boundary in the agent loop (R-1). + /// + /// `timeout_message` is only invoked when the timeout branch actually + /// fires, so callers can build a call-specific message (which fields it + /// names, which deadline it blames) without paying for the `format!` on + /// the hot, non-timeout path. `fut`'s own error type must convert from + /// [`TinyAgentsError`] so `Cancelled`/`Timeout` can be returned through + /// the same `Result` the callee already returns. + pub(crate) async fn bounded( + &self, + deadline: Option, + fut: impl std::future::Future>, + timeout_message: impl FnOnce() -> String, + ) -> std::result::Result + where + E: From, + { + match deadline { + Some(remaining) => tokio::select! { + biased; + _ = self.cancellation.cancelled() => { + Err(crate::error::TinyAgentsError::Cancelled.into()) + } + result = tokio::time::timeout(remaining, fut) => match result { + Ok(inner) => inner, + Err(_) => Err(crate::error::TinyAgentsError::Timeout(timeout_message()).into()), + }, + }, + None => tokio::select! { + biased; + _ = self.cancellation.cancelled() => { + Err(crate::error::TinyAgentsError::Cancelled.into()) + } + result = fut => result, + }, + } + } + /// Returns the maximum sub-agent / recursion depth permitted for this run /// tree. pub fn max_depth(&self) -> usize { diff --git a/crates/tinyagents-harness/src/context/stats.rs b/crates/tinyagents-harness/src/context/stats.rs index 67bcce0e..36c728fb 100644 --- a/crates/tinyagents-harness/src/context/stats.rs +++ b/crates/tinyagents-harness/src/context/stats.rs @@ -20,6 +20,8 @@ pub struct ContextStatistics { pub text_chars: usize, /// Image blocks across every role. pub images: usize, + /// Audio, video, and document blocks across every role. + pub media: usize, /// Tool calls requested by assistant messages. pub tool_calls: usize, /// Tool result messages. @@ -36,7 +38,7 @@ pub fn context_statistics(messages: &[Message]) -> ContextStatistics { }; let mut requested = std::collections::HashSet::new(); for message in messages { - let content = match message { + let content: &[ContentBlock] = match message { Message::System(message) => &message.content, Message::User(message) => &message.content, Message::Assistant(message) => { @@ -51,6 +53,8 @@ pub fn context_statistics(messages: &[Message]) -> ContextStatistics { } &message.content } + // Host-side out-of-band record; carries no content blocks. + Message::Custom(_) => &[], }; for block in content { match block { @@ -61,6 +65,9 @@ pub fn context_statistics(messages: &[Message]) -> ContextStatistics { stats.text_chars += value.to_string().chars().count(); } ContentBlock::Image(_) => stats.images += 1, + ContentBlock::Audio(_) | ContentBlock::Video(_) | ContentBlock::Document(_) => { + stats.media += 1; + } ContentBlock::RedactedThinking { .. } => {} } } @@ -77,11 +84,13 @@ pub fn estimate_context_tokens(messages: &[Message], tokenize: impl Fn(&str) -> messages .iter() .map(|message| { - let content = match message { + let content: &[ContentBlock] = match message { Message::System(message) => &message.content, Message::User(message) => &message.content, Message::Assistant(message) => &message.content, Message::Tool(message) => &message.content, + // Host-side out-of-band record; carries no content blocks. + Message::Custom(_) => &[], }; let mut visible = content .iter() diff --git a/crates/tinyagents-harness/src/context/test.rs b/crates/tinyagents-harness/src/context/test.rs index e4dfea70..8a586b63 100644 --- a/crates/tinyagents-harness/src/context/test.rs +++ b/crates/tinyagents-harness/src/context/test.rs @@ -186,8 +186,12 @@ fn child_carries_explicit_lineage_and_rejects_the_depth_cap() { .with_max_turn_output_tokens(123), (), ); - let child = parent.child(RunConfig::new("child"), "child-data").unwrap(); - let grandchild = child.child(RunConfig::new("grandchild"), ()).unwrap(); + let child = parent + .child_with_data(RunConfig::new("child"), "child-data") + .unwrap(); + let grandchild = child + .child_with_data(RunConfig::new("grandchild"), ()) + .unwrap(); assert_eq!(parent.lineage().root_run_id.as_str(), "root"); assert_eq!(parent.lineage().parent_run_id, None); @@ -211,7 +215,7 @@ fn child_carries_explicit_lineage_and_rejects_the_depth_cap() { assert_eq!(grandchild.thread_id().unwrap().as_str(), "thread"); assert_eq!(grandchild.config.max_turn_output_tokens, Some(123)); assert!(matches!( - grandchild.child(RunConfig::new("too-deep"), ()), + grandchild.child_with_data(RunConfig::new("too-deep"), ()), Err(crate::TinyAgentsError::SubAgentDepth(2)) )); } @@ -373,6 +377,7 @@ fn context_statistics_preserve_tool_request_result_pairing_and_image_counts() { content: vec![ContentBlock::Text("call it".into())], tool_calls: vec![ToolCall::new("call-1", "lookup", serde_json::json!({}))], usage: None, + origin: None, }), Message::Tool(tinyinference_llm::message::ToolMessage { tool_call_id: "call-1".into(), @@ -394,6 +399,7 @@ fn context_statistics_preserve_tool_request_result_pairing_and_image_counts() { messages: 3, text_chars: 13, images: 1, + media: 0, tool_calls: 1, tool_results: 1, paired_tool_results: 1, @@ -401,6 +407,23 @@ fn context_statistics_preserve_tool_request_result_pairing_and_image_counts() { ); } +#[test] +fn context_statistics_counts_audio_video_and_document_blocks_as_media() { + use tinyinference_llm::message::{ContentBlock, MediaRef, Message, UserMessage}; + + let messages = vec![Message::User(UserMessage { + content: vec![ + ContentBlock::Audio(MediaRef::url("https://example.com/a.wav")), + ContentBlock::Video(MediaRef::base64("AAAA", "video/mp4")), + ContentBlock::Document(MediaRef::path("/tmp/doc.pdf")), + ], + })]; + + let stats = context_statistics(&messages); + assert_eq!(stats.media, 3); + assert_eq!(stats.images, 0); +} + #[test] fn token_estimation_uses_the_callers_tokenizer() { use tinyinference_llm::message::Message; @@ -435,6 +458,7 @@ fn token_estimation_includes_structured_blocks_for_every_role() { let messages = vec![ Message::System(SystemMessage { content: vec![ContentBlock::ProviderExtension(json.clone())], + ..Default::default() }), Message::User(UserMessage { content: vec![ContentBlock::Json(json.clone())], @@ -444,6 +468,7 @@ fn token_estimation_includes_structured_blocks_for_every_role() { content: vec![ContentBlock::ProviderExtension(json.clone())], tool_calls: vec![], usage: None, + origin: None, }), Message::Tool(ToolMessage { tool_call_id: "call".into(), @@ -472,6 +497,7 @@ fn token_estimation_includes_assistant_tool_names_and_arguments() { serde_json::json!({"query": "one two three"}), )], usage: None, + origin: None, })]; let rendered = std::cell::RefCell::new(String::new()); @@ -486,3 +512,90 @@ fn token_estimation_includes_assistant_tool_names_and_arguments() { assert!(rendered.contains("search_docs")); assert!(rendered.contains("one two three")); } + +// ── RunContext::bounded (R-1) ─────────────────────────────────────────────── + +#[tokio::test] +async fn bounded_returns_the_futures_ok_value_with_no_deadline() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-bounded-ok"), ()); + let result: Result = ctx + .bounded(None, async { Ok(42) }, || "unused".to_string()) + .await; + assert_eq!(result.unwrap(), 42); +} + +#[tokio::test] +async fn bounded_passes_through_the_futures_own_error() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-bounded-err"), ()); + let result: Result = ctx + .bounded( + None, + async { Err(crate::error::TinyAgentsError::Model("boom".to_string())) }, + || "unused".to_string(), + ) + .await; + match result { + Err(crate::error::TinyAgentsError::Model(message)) => assert_eq!(message, "boom"), + other => panic!("expected a passthrough Model error, got {other:?}"), + } +} + +#[tokio::test] +async fn bounded_fires_the_timeout_message_only_when_the_deadline_elapses() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-bounded-timeout"), ()); + let mut message_built = false; + let result: Result = ctx + .bounded( + Some(std::time::Duration::from_millis(5)), + async { + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + Ok(1) + }, + || { + message_built = true; + "call-specific timeout message".to_string() + }, + ) + .await; + assert!(message_built); + match result { + Err(crate::error::TinyAgentsError::Timeout(message)) => { + assert_eq!(message, "call-specific timeout message"); + } + other => panic!("expected Timeout, got {other:?}"), + } +} + +#[tokio::test] +async fn bounded_does_not_build_the_timeout_message_on_the_success_path() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-bounded-no-timeout"), ()); + let result: Result = ctx + .bounded( + Some(std::time::Duration::from_secs(60)), + async { Ok(7) }, + || panic!("timeout_message must not be called when the future finishes first"), + ) + .await; + assert_eq!(result.unwrap(), 7); +} + +#[tokio::test] +async fn bounded_returns_cancelled_when_the_run_is_cancelled_before_the_future_resolves() { + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-bounded-cancel"), ()); + let cancellation = ctx.cancellation.clone(); + cancellation.cancel(); + let result: Result = ctx + .bounded( + None, + async { + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + Ok(1) + }, + || "unused".to_string(), + ) + .await; + assert!(matches!( + result, + Err(crate::error::TinyAgentsError::Cancelled) + )); +} diff --git a/crates/tinyagents-harness/src/context/types.rs b/crates/tinyagents-harness/src/context/types.rs index ac61dd84..4ac48fe3 100644 --- a/crates/tinyagents-harness/src/context/types.rs +++ b/crates/tinyagents-harness/src/context/types.rs @@ -14,22 +14,55 @@ //! `crate::context` directly. Implementations and tests live in the //! sibling `mod.rs` and `test.rs`. -use std::any::Any; - use serde::{Deserialize, Serialize}; use crate::cancel::CancellationToken; use crate::events::EventSink; -use crate::ids::{RunId, ThreadId}; +use crate::ids::{CallId, RunId, ThreadId}; use crate::limits::LimitTracker; use crate::steering::SteeringHandle; use crate::store::StoreRegistry; -/// One-shot observer invoked with the exact accumulated run when a driver -/// completes or is dropped. Kept crate-private: it is runtime lifecycle glue, -/// not a host policy extension point. +/// One-shot observer invoked with a cheap summary of the accumulated run when +/// a driver completes or is dropped. Kept crate-private: it is runtime +/// lifecycle glue, not a host policy extension point. +/// +/// Takes [`TerminalRunSummary`], not the full [`crate::middleware::AgentRun`] +/// (M-6): every installed observer only ever reads the final text, usage, and +/// executed-tool names, never the full transcript, and the observer needs an +/// *owned* value (the hosted path moves it into a spawned task that can +/// outlive the caller's stack frame) — so `&AgentRun` will not do either. The +/// summary is `Clone` and carries none of `AgentRun::messages`, which can be +/// the largest field by far on a long-running conversation. pub(crate) type TerminalObserver = - Box) + Send + Sync + 'static>; + Box) + Send + Sync + 'static>; + +/// Cheap, owned summary of an [`crate::middleware::AgentRun`] for +/// [`TerminalObserver`] — see that type's docs for why this exists instead of +/// the full run. +#[derive(Clone, Debug, Default)] +pub(crate) struct TerminalRunSummary { + /// The final response text, if the run produced one. Mirrors + /// [`crate::middleware::AgentRun::text`]. + pub(crate) text: Option, + /// Cumulative token usage across the run. `Copy`, so cloning this summary + /// is not where any cost lives. + pub(crate) usage: tinyinference_llm::usage::UsageTotals, + /// Names of calls that reached a tool executor, in execution order. + /// Mirrors [`crate::middleware::AgentRun::executed_tools`]. + pub(crate) executed_tools: Vec, +} + +impl TerminalRunSummary { + /// Builds a summary from a live run without cloning its transcript. + pub(crate) fn from_run(run: &crate::middleware::AgentRun) -> Self { + Self { + text: run.text(), + usage: run.usage, + executed_tools: run.executed_tools.clone(), + } + } +} /// The immutable ancestry of a run in a recursive harness invocation tree. /// @@ -134,6 +167,76 @@ pub struct RunConfig { pub lineage: RunLineage, } +/// Where [`MiddlewareControl::JumpTo`] sends the agent loop next. +/// +/// Modelled on LangChain's `jump_to: "model" | "tools" | "end"`. See +/// `docs/modules/harness/middleware.md` for exactly how each target is +/// realized against the loop's checkpoint structure. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LoopTarget { + /// Skip any remaining tool execution for this turn and go straight to the + /// next model call. + Model, + /// Proceed to (or continue) tool execution for this turn. A no-op when the + /// turn has no tool calls to run — there is nothing to jump to. + Tools, + /// Stop the loop now, finishing the run with the transcript as it stands. + End, +} + +/// A typed hook that mutates application state, carried by +/// [`MiddlewareControl::UpdateState`]. +/// +/// `State` is type-erased on construction (`RunContext` is not generic over +/// it) and recovered by [`Self::apply`] via a runtime check. Built with +/// [`StateUpdate::new`], which captures an `Fn(&mut State)` closure in an +/// `Arc` so [`MiddlewareControl`] (and therefore `StateUpdate`) stays +/// [`Clone`] — required because [`RunContext::request_control`] may compare +/// and replace a pending request. +/// +/// The agent loop only ever sees `state: &State` (a shared reference), so it +/// cannot apply this itself. [`RunContext::take_state_updates`] queues every +/// requested update instead; a host that owns `&mut State` between runs (or +/// between turns, via its own checkpoint) drains and applies them. See +/// `docs/modules/harness/middleware.md` for the full contract. +/// The type-erased closure a [`StateUpdate`] wraps. +type ErasedStateUpdateFn = std::sync::Arc; + +#[derive(Clone)] +pub struct StateUpdate { + apply: ErasedStateUpdateFn, +} + +impl StateUpdate { + /// Captures `f` as a state update for the concrete application state type + /// `S`. Applying the update against any other type is a documented no-op + /// (see [`Self::apply`]). + pub fn new(f: impl Fn(&mut S) + Send + Sync + 'static) -> Self { + Self { + apply: std::sync::Arc::new(move |state: &mut dyn std::any::Any| { + if let Some(state) = state.downcast_mut::() { + f(state); + } + }), + } + } + + /// Applies this update to `state` when `state`'s concrete type matches the + /// type this update was constructed for. A mismatched type is a silent + /// no-op: the update was requested by middleware generic over a different + /// `State`, which a host wiring several harnesses together can otherwise + /// hit legitimately. + pub fn apply(&self, state: &mut S) { + (self.apply)(state as &mut dyn std::any::Any); + } +} + +impl std::fmt::Debug for StateUpdate { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("StateUpdate(..)") + } +} + /// A structured control outcome a middleware (or any step) can request on the /// [`RunContext`] to steer the agent loop from outside its `Result<()>` return /// channel. @@ -144,8 +247,28 @@ pub struct RunConfig { /// "stop after an early-exit tool" or "pause on budget" no longer need a /// bespoke side channel. Requests are visible via /// [`RunContext::take_control`]. -#[derive(Clone, Debug, PartialEq, Eq)] +/// +/// A [`Middleware`][crate::middleware::Middleware] hook may also *return* one +/// of these directly from its `_control`-suffixed variant (for example +/// [`before_model_control`][crate::middleware::Middleware::before_model_control]); +/// the [`MiddlewareStack`][crate::middleware::MiddlewareStack] resolves a +/// non-[`Continue`](Self::Continue) return into exactly the same +/// [`RunContext::request_control`] call a hook could have made explicitly — +/// returning control is sugar over the side channel, not a second mechanism. +#[derive(Clone, Debug)] pub enum MiddlewareControl { + /// No control requested. The default a `_control` hook returns when it has + /// nothing to say; never itself installed as a pending request (see + /// [`RunContext::request_control`]). + Continue, + /// Route the loop to `target` at the next safe checkpoint. See + /// [`LoopTarget`] for what each target does. + JumpTo(LoopTarget), + /// Queue a typed state mutation for the host to apply. The loop itself + /// only ever holds `&State`, so this is queued on + /// [`RunContext::take_state_updates`] rather than applied in place; see + /// [`StateUpdate`]. + UpdateState(StateUpdate), /// Stop the loop now and use this text as the final assistant response. StopWithFinal(String), /// Pause the run at the next safe checkpoint, surfacing @@ -163,6 +286,11 @@ impl MiddlewareControl { /// A stable label for this control outcome, used in audit events. pub fn kind(&self) -> &'static str { match self { + MiddlewareControl::Continue => "continue", + MiddlewareControl::JumpTo(LoopTarget::Model) => "jump_to:model", + MiddlewareControl::JumpTo(LoopTarget::Tools) => "jump_to:tools", + MiddlewareControl::JumpTo(LoopTarget::End) => "jump_to:end", + MiddlewareControl::UpdateState(_) => "update_state", MiddlewareControl::StopWithFinal(_) => "stop_with_final", MiddlewareControl::Interrupt { .. } => "interrupt", } @@ -173,10 +301,19 @@ impl MiddlewareControl { /// [`StopWithFinal`](Self::StopWithFinal) because pausing to preserve state /// for a later resume is stronger than terminating with a final answer, so /// a pause request is never silently downgraded to a stop. + /// [`Continue`](Self::Continue) is the lowest rank: it carries no + /// instruction and is never itself installed as a pending request (see + /// [`RunContext::request_control`]). [`UpdateState`](Self::UpdateState) + /// and [`JumpTo`](Self::JumpTo) sit below the two run-ending outcomes so a + /// state patch or a soft reroute never displaces a stop or an interrupt + /// that a later hook in the same phase also requested. pub fn precedence(&self) -> u8 { match self { - MiddlewareControl::StopWithFinal(_) => 1, - MiddlewareControl::Interrupt { .. } => 2, + MiddlewareControl::Continue => 0, + MiddlewareControl::UpdateState(_) => 1, + MiddlewareControl::JumpTo(_) => 2, + MiddlewareControl::StopWithFinal(_) => 3, + MiddlewareControl::Interrupt { .. } => 4, } } } @@ -208,6 +345,26 @@ pub struct RunContext { pub data: Ctx, /// Registry of named long-term stores. pub stores: StoreRegistry, + /// Optional hierarchical long-term store handed to every tool this run + /// invokes as + /// [`ToolExecutionContext::store`][crate::tool::ToolExecutionContext::store] + /// (B1). Distinct from [`Self::stores`], the flat named registry: this + /// is the one [`NamespacedStore`][crate::store::namespaced::NamespacedStore] + /// a tool may read and write directly — memories, scratch state, a + /// per-user cache — without the harness minting a name for it. `None` + /// means tools get no store. Attach one with + /// [`RunContext::with_namespaced_store`]; shared with child contexts + /// exactly like `stores`. + pub namespaced_store: Option>, + /// Optional type-erased, read-only view of the application state handed + /// to every tool this run invokes, recovered by + /// [`ToolExecutionContext::state`][crate::tool::ToolExecutionContext::state] + /// (B1). Erased because `RunContext` is not generic over `State` and the + /// agent loop only ever holds a borrowed `&State` it cannot lend to a + /// concurrent tool future; the host attaches an owned `Arc` snapshot + /// with [`RunContext::with_state_view`] instead. Shared with child + /// contexts, which run against the same application state. + pub state_view: Option>, /// Event fan-out bus for observability. pub events: EventSink, /// Live limit tracker derived from `config`. @@ -220,6 +377,15 @@ pub struct RunContext { /// model call via /// [`crate::steering::apply_pending_steering`]. pub steering: Option, + /// Optional multi-lane message queue the agent loop drains at its turn + /// boundaries (A4): `Steer` after each tool batch and at a natural + /// finish, `Followup` at a natural finish only, `Collect` once at run + /// end onto [`crate::middleware::AgentRun::collected`]. `None` means the + /// loop consumes no queued messages. Attach one with + /// [`RunContext::with_run_queue`]; never inherited by a child context, + /// because a queue has no per-run addressing and a child draining its + /// parent's queue would steal the parent's messages. + pub run_queue: Option, /// Cooperative cancellation token for this run. /// /// Defaults to a fresh, never-cancelled [`CancellationToken`], so a run is @@ -235,6 +401,15 @@ pub struct RunContext { /// loop (stop with a final response, or interrupt). Drained by the agent /// loop at its safe checkpoints via [`RunContext::take_control`]. pub control: std::sync::Arc>>, + /// Queued [`StateUpdate`]s a middleware or tool requested via + /// [`MiddlewareControl::UpdateState`], drained by a host through + /// [`RunContext::take_state_updates`]. See that method's docs for why the + /// loop cannot apply these itself. + pub(crate) state_updates: std::sync::Arc>>, + /// Queued raw JSON state updates a tool requested via + /// [`tinytools::ToolControl::state_update`]. See + /// [`RunContext::push_tool_state_update`]. + pub(crate) tool_state_updates: std::sync::Arc>>, /// An optional host-supplied workspace/sandbox descriptor threaded into every /// [`ToolExecutionContext`][crate::tool::ToolExecutionContext] this /// run creates, so tools discover their allowed root from context rather @@ -263,8 +438,82 @@ pub struct RunContext { /// is deliberately not serializable or public: it keeps a hosted parent /// from accidentally delegating through a child's unrelated (or absent) /// capability bundle. - pub(crate) host_authority: Option>, + /// + /// Erased through [`crate::runtime::ErasedHostAuthority`] rather than + /// `dyn Any`: the generic explicit-model loop must stay callable with a + /// borrowed (non-`'static`) `State`/`Ctx`, and `Any::downcast_ref` + /// requires `'static` at the *read* site, which such a caller can never + /// prove. The custom trait instead exposes a type-name check that needs + /// no `'static` bound on either side; see + /// [`crate::runtime::host_invocation_binding`] for how the read side + /// uses it to fail closed on a mismatch. + pub(crate) host_authority: Option>, /// Runtime-owned terminal lifecycle callback, consumed exactly once by the /// agent-loop guard even when the driving future is cancelled or dropped. pub(crate) terminal_observer: Option, + /// The [`CallId`] the agent loop minted for the model call currently in + /// flight through the model-wrap middleware onion, mirroring + /// [`crate::events::HarnessRunStatus::active_model_call`]. + /// + /// Set by the loop immediately before invoking + /// [`crate::middleware::MiddlewareStack::run_wrapped_model`] and cleared + /// right after, so a `ModelMiddleware` such as + /// [`crate::middleware::library::RetryMiddleware`] can correlate its own + /// `RetryScheduled` events with the same call id the loop uses, instead of + /// deriving an uncorrelated one from `ctx.run_id()` alone (see I-7). + /// `None` outside that window, and always `None` for a caller that never + /// goes through the agent loop. + pub active_model_call: Option, + /// Resolutions for the deferred tool calls left pending on the transcript + /// this run is resuming (A2). Taken by the agent loop before its first + /// model call and applied to the unanswered tool calls on the last + /// assistant row; see + /// [`crate::runtime::AgentHarness::resume_deferred`]. Never inherited by + /// a child context. + pub(crate) deferred_results: Option, + /// Tool-call ids a human approved on resume (A2). Admission skips the + /// deferral checks for these, and a `before_tool` hook's + /// `ApprovalRequired` is ignored for them, so an approved call cannot be + /// deferred a second time by the same gate. Read with + /// [`RunContext::is_call_approved`]. + pub(crate) approved_calls: std::collections::HashSet, + /// Monotonic, per-context (not process-global) counter handed out by + /// [`RunContext::next_child_ordinal`], used to derive deterministic child + /// run ids (e.g. [`crate::subagent::SubAgent`]'s `{name}-d{depth}-{parent + /// run id}-{ordinal}`) instead of a process-global sequence (M-2). Starts + /// at `0` for every freshly constructed context — including a child + /// context, which gets its own fresh counter rather than inheriting the + /// parent's — so two processes that call the same parent context's child + /// spawner in the same order derive identical ordinals, and therefore + /// identical child run ids. + pub(crate) child_ordinal: std::sync::Arc, + /// Durable tool-effect ledger for this run, when a host wants crash-safe + /// bookkeeping of tool-call side effects (B5). `None` (the default) means + /// no ledger writes happen and [`crate::tool::ToolPolicy`]'s + /// `runtime.replay` declaration has nothing to guard resume against — the + /// agent loop behaves exactly as it did before this existed. Attach one + /// with [`RunContext::with_tool_effect_ledger`]; a child context inherits + /// its parent's ledger, matching how `stores`/`events` propagate. + pub tool_effect_ledger: Option>, + /// How the agent loop reacts when a `started` write to + /// [`Self::tool_effect_ledger`] itself fails, before the tool call it was + /// about to journal has executed. See + /// [`crate::tool::LedgerFailure`] for the two modes; defaults to + /// [`crate::tool::LedgerFailure::Abort`]. + pub tool_effect_ledger_failure: crate::tool::LedgerFailure, + /// Durable sink for [`crate::summarization::CompactionRecord`]s this run + /// produces, when a host wants every compaction persisted somewhere + /// durable rather than only kept in + /// [`crate::middleware::ContextCompressionMiddleware::records`]'s + /// in-process buffer. + /// + /// `None` (the default) means compaction runs exactly as it did before + /// this existed — no persistence side effect. Attach one with + /// [`RunContext::with_compaction_sink`]; a child context inherits its + /// parent's sink, matching how `stores`/`events`/`tool_effect_ledger` + /// propagate. `tinyagents-harness` cannot depend on + /// `tinyagents-session` (the dependency runs the other way), so this is + /// a trait object rather than a concrete `Arc` — see + /// [`crate::summarization::CompactionSink`]'s docs. + pub compaction_sink: Option>, } diff --git a/crates/tinyagents-harness/src/cost/mod.rs b/crates/tinyagents-harness/src/cost/mod.rs index 29133b7a..842cf118 100644 --- a/crates/tinyagents-harness/src/cost/mod.rs +++ b/crates/tinyagents-harness/src/cost/mod.rs @@ -48,6 +48,24 @@ impl AddAssign for CostTotals { } } +/// Selects the [`PriceTier`] matching `input_tokens`, when `pricing` declares +/// any. The match is the tier with the smallest `up_to_tokens` that is still +/// `>= input_tokens`, or the unlimited tier (`up_to_tokens: None`) when +/// `input_tokens` exceeds every capped tier. Returns `None` when `pricing` +/// declares no tiers. +/// +/// See [`ModelPricing::tiers`]. +fn select_tier(pricing: &ModelPricing, input_tokens: u64) -> Option<&PriceTier> { + if pricing.tiers.is_empty() { + return None; + } + pricing + .tiers + .iter() + .filter(|tier| tier.up_to_tokens.is_none_or(|cap| input_tokens <= cap)) + .min_by_key(|tier| tier.up_to_tokens.unwrap_or(u64::MAX)) +} + /// Estimates the cost of a [`Usage`] record using per-token [`ModelPricing`]. /// /// Missing prices contribute zero. Cache read and cache creation tokens are @@ -60,20 +78,35 @@ impl AddAssign for CostTotals { /// standard rate *and* separately pricing the cached/reasoning subset would /// double-charge those tokens, so the standard-rate cost is computed on the /// non-cached/non-reasoning remainder only. +/// +/// When `pricing` declares [`ModelPricing::tiers`], the tier matching this +/// call's `usage.input_tokens` (see [`select_tier`]) supplies the input, +/// output, cache-read, and cache-write rates, falling back to the flat +/// [`ModelPricing`] fields for any rate the tier leaves unset. Reasoning +/// tokens are always priced at the flat +/// [`ModelPricing::output_reasoning_per_token`] rate; tiers do not currently +/// carry a reasoning override. pub fn estimate_cost(pricing: &ModelPricing, usage: &Usage) -> CostTotals { let price = |rate: Option, tokens: u64| rate.unwrap_or(0.0) * tokens as f64; + let tier = select_tier(pricing, usage.input_tokens); + let input_rate = tier.and_then(|t| t.input).or(pricing.input_per_token); + let output_rate = tier.and_then(|t| t.output).or(pricing.output_per_token); + let cache_read_rate = tier + .and_then(|t| t.cache_read) + .or(pricing.cache_read_input_per_token); + let cache_write_rate = tier + .and_then(|t| t.cache_write) + .or(pricing.cache_creation_input_per_token); + let billable_input_tokens = usage.input_tokens.saturating_sub(usage.cache_read_tokens); let billable_output_tokens = usage.output_tokens.saturating_sub(usage.reasoning_tokens); let mut totals = CostTotals { - input_cost: price(pricing.input_per_token, billable_input_tokens), - output_cost: price(pricing.output_per_token, billable_output_tokens), - cache_cost: price(pricing.cache_read_input_per_token, usage.cache_read_tokens) - + price( - pricing.cache_creation_input_per_token, - usage.cache_creation_tokens, - ), + input_cost: price(input_rate, billable_input_tokens), + output_cost: price(output_rate, billable_output_tokens), + cache_cost: price(cache_read_rate, usage.cache_read_tokens) + + price(cache_write_rate, usage.cache_creation_tokens), reasoning_cost: price(pricing.output_reasoning_per_token, usage.reasoning_tokens), total_cost: 0.0, }; diff --git a/crates/tinyagents-harness/src/cost/test.rs b/crates/tinyagents-harness/src/cost/test.rs index 7790bf99..e9b6f8ed 100644 --- a/crates/tinyagents-harness/src/cost/test.rs +++ b/crates/tinyagents-harness/src/cost/test.rs @@ -82,3 +82,89 @@ fn cost_totals_accumulate() { assert!((totals.input_cost - 2.0).abs() < 1e-9); assert!((totals.total_cost - 2.0).abs() < 1e-9); } + +// --------------------------------------------------------------------------- +// Tiered pricing +// --------------------------------------------------------------------------- + +fn tiered_pricing() -> ModelPricing { + ModelPricing { + // Flat rate is the fallback the highest tier partially relies on. + input_per_token: Some(0.001), + output_per_token: Some(0.002), + tiers: vec![ + PriceTier { + up_to_tokens: Some(200_000), + input: Some(0.001), + output: Some(0.002), + cache_read: Some(0.0001), + cache_write: None, + }, + PriceTier { + up_to_tokens: None, + input: Some(0.002), + output: Some(0.004), + cache_read: None, + cache_write: None, + }, + ], + ..ModelPricing::default() + } +} + +#[test] +fn tiered_pricing_uses_the_matching_capped_tier() { + let usage = Usage::new(100_000, 1000); + let cost = estimate_cost(&tiered_pricing(), &usage); + assert!((cost.input_cost - 100.0).abs() < 1e-9); + assert!((cost.output_cost - 2.0).abs() < 1e-9); +} + +#[test] +fn tiered_pricing_falls_through_to_the_unlimited_tier_above_the_cap() { + let usage = Usage::new(300_000, 1000); + let cost = estimate_cost(&tiered_pricing(), &usage); + // The unlimited tier's higher input rate (0.002) applies once the call's + // input tokens exceed the capped tier's 200K threshold. + assert!((cost.input_cost - 600.0).abs() < 1e-9); + assert!((cost.output_cost - 4.0).abs() < 1e-9); +} + +#[test] +fn tiered_pricing_falls_back_to_flat_rate_for_unset_tier_fields() { + // The unlimited tier leaves `cache_read` unset; it must fall back to the + // flat `cache_read_input_per_token` rather than pricing at zero. + let mut pricing = tiered_pricing(); + pricing.cache_read_input_per_token = Some(0.0002); + let usage = Usage { + input_tokens: 300_000, + output_tokens: 0, + total_tokens: 300_000, + cache_read_tokens: 50_000, + cache_creation_tokens: 0, + reasoning_tokens: 0, + charged_amount: None, + context_window_tokens: None, + }; + let cost = estimate_cost(&pricing, &usage); + assert!((cost.cache_cost - 10.0).abs() < 1e-9); +} + +#[test] +fn no_tiers_declared_uses_flat_pricing() { + assert!(select_tier(&pricing(), 1_000_000).is_none()); +} + +#[test] +fn select_tier_picks_the_smallest_sufficient_cap() { + let pricing = tiered_pricing(); + assert_eq!( + select_tier(&pricing, 50_000).unwrap().up_to_tokens, + Some(200_000) + ); + assert_eq!( + select_tier(&pricing, 200_000).unwrap().up_to_tokens, + Some(200_000) + ); + assert_eq!(select_tier(&pricing, 200_001).unwrap().up_to_tokens, None); +} diff --git a/crates/tinyagents-harness/src/cost/types.rs b/crates/tinyagents-harness/src/cost/types.rs index 130d1fbf..d5a79367 100644 --- a/crates/tinyagents-harness/src/cost/types.rs +++ b/crates/tinyagents-harness/src/cost/types.rs @@ -29,6 +29,44 @@ pub struct ModelPricing { /// Price per reasoning output token. #[serde(default)] pub output_reasoning_per_token: Option, + /// Context-size-tiered pricing (some providers charge more once a call's + /// context crosses a threshold, e.g. Gemini's price step above 200K + /// input tokens). Empty by default: when set, [`crate::cost::estimate_cost`] + /// selects the tier matching the call's input-token count and uses its + /// rates in place of the flat fields above (falling back to the flat + /// fields for any rate a matched tier leaves `None`). Tiers do not need + /// to be pre-sorted; the matching tier is the one with the smallest + /// [`PriceTier::up_to_tokens`] that is still `>=` the call's input-token + /// count, or the tier with no `up_to_tokens` (unlimited) when the count + /// exceeds every capped tier. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub tiers: Vec, +} + +/// One context-size pricing tier. See [`ModelPricing::tiers`]. +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub struct PriceTier { + /// Upper bound (inclusive) on input tokens this tier applies to. `None` + /// means "no upper bound" — the tier for any call whose input-token + /// count exceeds every other tier's `up_to_tokens`. + #[serde(default)] + pub up_to_tokens: Option, + /// Price per input token in this tier. Falls back to + /// [`ModelPricing::input_per_token`] when `None`. + #[serde(default)] + pub input: Option, + /// Price per output token in this tier. Falls back to + /// [`ModelPricing::output_per_token`] when `None`. + #[serde(default)] + pub output: Option, + /// Price per cached input token in this tier. Falls back to + /// [`ModelPricing::cache_read_input_per_token`] when `None`. + #[serde(default)] + pub cache_read: Option, + /// Price per input token written to a prompt cache in this tier. Falls + /// back to [`ModelPricing::cache_creation_input_per_token`] when `None`. + #[serde(default)] + pub cache_write: Option, } /// A breakdown of estimated cost for one or more model calls, in the pricing diff --git a/crates/tinyagents-harness/src/error.rs b/crates/tinyagents-harness/src/error.rs index d8288464..72ac88c5 100644 --- a/crates/tinyagents-harness/src/error.rs +++ b/crates/tinyagents-harness/src/error.rs @@ -18,6 +18,7 @@ pub type Result = std::result::Result; /// execution, model/tool invocation, run limits and policy, graph durability, /// and graph execution. #[derive(Debug, Error)] +#[non_exhaustive] pub enum TinyAgentsError { /// A graph was compiled or run without a configured `START` edge, so there /// is no entry node to begin execution from. @@ -129,6 +130,69 @@ pub enum TinyAgentsError { #[error("tool error: {0}")] Tool(String), + /// A tool or an output validator ([`crate::structured::OutputValidator`]) + /// reported a *recoverable* failure the model should be asked to fix and + /// retry, rather than one that ends the run. + /// + /// Mirrors Pydantic AI's `ModelRetry`. A tool returning this from + /// [`tinytools::Tool::execute`] is folded into a recoverable + /// [`tinytools::ToolResult::retry`] result instead of aborting the run — + /// see `agent_loop/tools.rs`'s `map_tool_dispatch_error`. An + /// [`crate::structured::OutputValidator`] returning it on the agent + /// loop's final turn drives the output-validation retry loop (A3): the + /// message is pushed back to the model as a repair prompt and the turn + /// continues, bounded by + /// [`crate::runtime::RunPolicy::output_retry`]'s `max_attempts`. + /// Contrast with [`Self::ToolFailed`], which is permanent. + #[error("retryable failure: {0}")] + ModelRetry(String), + + /// A tool reported a **permanent** failure that must not be retried — + /// the counterpart to [`Self::ModelRetry`]. Folded into a + /// [`tinytools::ToolResult::failed`] result (still recoverable at the + /// transcript level — the model sees the message — but + /// [`crate::retry::RetryMiddleware`] and any other retry policy treat it + /// as non-retryable rather than re-attempting the call). + #[error("permanent tool failure: {0}")] + ToolFailed(String), + + /// A tool (from [`tinytools::Tool::execute`]) or a `before_tool` + /// middleware asked for **human approval** before this call runs (A2). + /// + /// The agent loop does not treat this as a failure: it finishes the rest + /// of the batch, lists the call under + /// [`crate::tool::DeferredToolRequests::approvals`] with `metadata` + /// attached, and exits with `AgentRun::deferred` set (or resolves it + /// inline through a registered + /// [`crate::tool::DeferredToolHandler`]). Mirrors Pydantic AI's + /// `ApprovalRequired`. Never retried by [`crate::retry::is_retryable`]. + #[error("tool call requires approval")] + ApprovalRequired { + /// Host-only context for the approver (never shown to the model). + metadata: serde_json::Value, + }, + + /// A tool asked the **host** to execute this call out of band (A2): + /// the loop lists it under [`crate::tool::DeferredToolRequests::calls`] + /// and expects a [`crate::tool::DeferredCallResult`] on resume. Raised + /// automatically for a tool registered through + /// [`crate::tool::ToolRegistry::register_external`], and equally by + /// [`crate::tool::toolset::ExternalToolSet`]'s + /// [`crate::tool::toolset::ToolSet::call`] (gap B3, mirroring Pydantic + /// AI's `defer_loading`/deferred-tools model) — both call sites raise + /// this same variant so a host sees one deferred-call signal regardless + /// of which registration path advertised the tool. `metadata` is + /// host-only context describing how to execute the call; the call's own + /// name and arguments are already carried on the + /// [`crate::tool::DeferredToolRequests`] entry, so most callers leave it + /// `Value::Null`. Mirrors Pydantic AI's `CallDeferred`. Never retried by + /// [`crate::retry::is_retryable`]. + #[error("tool call deferred to the host")] + CallDeferred { + /// Host-only context describing how to execute the call. + metadata: serde_json::Value, + }, + /// A run referenced a tool name that is not present in the /// [`crate::tool::ToolRegistry`]. The payload is the tool name. #[error("tool `{0}` is not registered")] @@ -164,9 +228,26 @@ pub enum TinyAgentsError { EmptyResponse, /// The run exceeded its wall-clock deadline. + /// + /// Terminal: the run itself is out of time, so retrying or falling back + /// to another model would just spin until the next deadline check fails + /// identically. See [`TinyAgentsError::CallTimeout`] for the per-call + /// counterpart, which *is* retryable. #[error("run timed out: {0}")] Timeout(String), + /// A single call (currently: a model call bounded by + /// [`crate::limits::RunLimits::max_model_call_ms`]) ran past its own + /// ceiling while the run still has wall-clock budget left. + /// + /// Unlike [`TinyAgentsError::Timeout`], this does not mean the run is out + /// of time — it means *this one call* wedged. [`crate::retry::is_retryable`] + /// treats it as transient, and the model-resolution retry/fallback loop + /// (`invoke_model_resolving`) does not treat it as a reason to skip the + /// fallback chain the way it does a run-deadline `Timeout`. + #[error("call timed out: {0}")] + CallTimeout(String), + /// The run was cancelled before completion. #[error("run cancelled")] Cancelled, @@ -247,6 +328,57 @@ pub enum TinyAgentsError { /// underlying driver message. #[error("storage error: {0}")] Storage(String), + + /// One or more `.rag` language diagnostics, collected together instead of + /// stopping at the first offending reference or construct. + /// + /// The payload is [`RenderedDiagnostic`], not + /// `tinyagents_language::Diagnostic`, because `tinyagents-language` + /// depends on this crate for [`Result`]/`TinyAgentsError`; holding the + /// language crate's structured type here would create an import cycle. + /// `tinyagents_language::diagnostic::into_diagnostics_error` builds this + /// variant from a `Vec` by rendering each + /// one down to its message, code, and resolved position. Never + /// constructed with an empty vector. + #[error("{}", render_diagnostics_summary(.0))] + Diagnostics(Vec), +} + +/// One `.rag` language diagnostic, rendered to a crate-boundary-safe, +/// serializable payload for [`TinyAgentsError::Diagnostics`]. +/// +/// See that variant's docs for why this mirrors (rather than reuses) +/// `tinyagents_language::Diagnostic`. +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct RenderedDiagnostic { + /// The diagnostic's stable code (e.g. `E-rag-unknown-model`), if any. + pub code: Option, + /// The headline message, without source context. + pub message: String, + /// The 1-based line the diagnostic's primary span begins at. + pub line: usize, + /// The 1-based column the diagnostic's primary span begins at. + pub column: usize, + /// The full presentation: the caret-underline rendering against source + /// when it was available at construction time, otherwise the + /// source-free `message` plus a `-->` position anchor. + pub rendered: String, +} + +/// Renders the [`TinyAgentsError::Diagnostics`] `Display` text: the first +/// diagnostic's full rendering, plus a `(and N more)` suffix when there is +/// more than one. +fn render_diagnostics_summary(diagnostics: &[RenderedDiagnostic]) -> String { + match diagnostics.split_first() { + Some((first, [])) => first.rendered.clone(), + Some((first, rest)) => format!( + "{} (and {} more diagnostic{})", + first.rendered, + rest.len(), + if rest.len() == 1 { "" } else { "s" } + ), + None => "no diagnostics".to_string(), + } } impl From for TinyAgentsError { @@ -257,6 +389,7 @@ impl From for TinyAgentsError { tinyinference_llm::Error::Validation(message) => Self::Validation(message), tinyinference_llm::Error::Serialization(error) => Self::Serialization(error), tinyinference_llm::Error::Catalog(message) => Self::Model(message), + tinyinference_llm::Error::Unsupported(message) => Self::Validation(message), } } } @@ -290,7 +423,7 @@ impl TinyAgentsError { if error.code.as_deref() == Some(tinyinference_llm::providers::openai::CONTEXT_OVERFLOW_CODE) { - tinyagents_tracing::debug!( + tracing::debug!( "[error] promoting provider `{}` context-overflow code to a typed error", error.provider ); diff --git a/crates/tinyagents-harness/src/events/mod.rs b/crates/tinyagents-harness/src/events/mod.rs index 35ca1534..3d2fd5ea 100644 --- a/crates/tinyagents-harness/src/events/mod.rs +++ b/crates/tinyagents-harness/src/events/mod.rs @@ -165,14 +165,16 @@ impl EventSink { inner.next_offset += 1; let id = crate::ids::EventId::new(format!("{}-evt-{offset}", inner.stream_id)); let record = EventRecord { id, offset, event }; - // Most production invocations do not attach an observer. Avoid a - // record clone, queue allocation, and serialized drain cycle in - // that common path; offsets still advance so a later subscriber - // starts at the correct position and never sees earlier events. + // The id/offset must still be minted with no listeners — callers + // (e.g. `HarnessRunStatus::set_last_event`) rely on the returned + // record regardless of whether anyone is watching — but with + // nothing registered there is nothing to fan out to, so the + // `Arc` clone, enqueue, and drain loop below are pure overhead on + // every emit of a run nobody is observing. Skip them. if inner.listeners.is_empty() { return record; } - let listeners = Arc::clone(&inner.listeners); + let listeners = inner.listeners.clone(); inner.pending.push_back((record.clone(), listeners)); let should_drain = !inner.dispatching; if should_drain { diff --git a/crates/tinyagents-harness/src/events/test.rs b/crates/tinyagents-harness/src/events/test.rs index bdf4fe09..925e5b38 100644 --- a/crates/tinyagents-harness/src/events/test.rs +++ b/crates/tinyagents-harness/src/events/test.rs @@ -150,6 +150,7 @@ fn completed_events_deserialize_without_started_at_ms() { duration_ms: Some(12), output_bytes: Some(5), error: None, + metadata: None, }; let json = serde_json::to_value(&event).unwrap(); assert_eq!(json["started_at_ms"], 1_704_067_199_000u64); @@ -499,3 +500,87 @@ fn every_started_variant_pairs_with_both_a_completed_and_a_failed_variant() { ); } } + +#[test] +fn custom_event_round_trips_with_its_call_id_and_payload() { + let event = AgentEvent::Custom { + call_id: Some(crate::ids::CallId::new("call-9")), + payload: serde_json::json!({ "progress": 0.5 }), + }; + assert_eq!(event.kind(), "custom"); + let json = serde_json::to_value(&event).unwrap(); + assert_eq!(json["kind"], "custom"); + assert_eq!(json["call_id"], "call-9"); + assert_eq!(json["payload"]["progress"], 0.5); + let back: AgentEvent = serde_json::from_value(json).unwrap(); + assert_eq!(back, event); + + // Outside a tool call the correlation is absent and omitted on the wire. + let event = AgentEvent::Custom { + call_id: None, + payload: serde_json::json!("ping"), + }; + let json = serde_json::to_value(&event).unwrap(); + assert!(json.get("call_id").is_none()); + let back: AgentEvent = serde_json::from_value(json).unwrap(); + assert_eq!(back, event); +} + +#[test] +fn tool_completed_metadata_is_optional_and_round_trips() { + let event: AgentEvent = + serde_json::from_str(r#"{"kind":"tool_completed","call_id":"t1","tool_name":"lookup"}"#) + .expect("pre-metadata tool_completed still deserializes"); + assert!(matches!( + event, + AgentEvent::ToolCompleted { metadata: None, .. } + )); + + let event = AgentEvent::ToolCompleted { + call_id: crate::ids::CallId::new("t2"), + tool_name: "lookup".to_string(), + started_at_ms: None, + input: None, + output: None, + duration_ms: None, + output_bytes: None, + error: None, + metadata: Some(serde_json::json!({ "coords": [1, 2] })), + }; + let json = serde_json::to_value(&event).unwrap(); + assert_eq!(json["metadata"]["coords"][0], 1); + let back: AgentEvent = serde_json::from_value(json).unwrap(); + assert_eq!(back, event); +} + +// ── EventSink::emit zero-listener fast path ───────────────────────────────── + +#[test] +fn emit_with_no_listeners_still_mints_ids_in_offset_order() { + let sink = EventSink::new(); + assert_eq!(sink.len(), 0); + + let first = sink.emit(AgentEvent::StateUpdate); + let second = sink.emit(AgentEvent::StateUpdate); + + assert_eq!(first.offset, 0); + assert_eq!(second.offset, 1); + assert_ne!(first.id, second.id); +} + +#[test] +fn emit_delivers_normally_once_a_listener_subscribes_after_a_quiet_run() { + let sink = EventSink::new(); + // Emitted while nobody is listening: takes the fast path. + sink.emit(AgentEvent::StateUpdate); + + let recorder = Arc::new(RecordingListener::new()); + sink.subscribe(recorder.clone()); + + // Emitted once a listener exists: should be delivered and continue the + // same offset sequence rather than resetting or skipping an id. + let delivered = sink.emit(AgentEvent::StateUpdate); + assert_eq!(delivered.offset, 1); + assert_eq!(recorder.events().len(), 1); + assert_eq!(recorder.events()[0].offset, 1); +} diff --git a/crates/tinyagents-harness/src/events/types.rs b/crates/tinyagents-harness/src/events/types.rs index 98444a17..292be873 100644 --- a/crates/tinyagents-harness/src/events/types.rs +++ b/crates/tinyagents-harness/src/events/types.rs @@ -36,6 +36,7 @@ use tinyinference_llm::usage::{Usage, UsageTotals}; /// the event type without inspecting nested fields. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case", tag = "kind")] +#[non_exhaustive] pub enum AgentEvent { /// A new harness run has been initiated. RunStarted { @@ -142,6 +143,36 @@ pub enum AgentEvent { tool_name: String, }, + /// A tool call was deferred out of the loop (A2): it needs a human + /// approval or host-side execution before it can be answered. The loop + /// finishes the batch's other calls and exits with + /// `AgentRun::deferred`, or resolves it inline through a registered + /// `DeferredToolHandler`. Terminal partner of a `ToolStarted` when the + /// tool itself raised the deferral mid-execution. + ToolDeferred { + /// Identifier of the deferred call. + call_id: CallId, + /// Why it was deferred (`approval_required`, `call_deferred`, + /// `external`, or a middleware-supplied reason). + reason: String, + }, + + /// A previously deferred call was approved on resume and is about to + /// execute (with the model's or the approver's edited arguments). + ToolApproved { + /// Identifier of the approved call. + call_id: CallId, + }, + + /// A previously deferred call was denied on resume; no tool runs and the + /// model sees `message` as a tool-error result. + ToolDenied { + /// Identifier of the denied call. + call_id: CallId, + /// The denial message handed to the model. + message: String, + }, + /// A tool-selection middleware filtered the model-visible tool set before a /// model call. Makes exposure decisions auditable: a UI or log can see /// which tools were withheld from the model and by which policy. @@ -152,6 +183,17 @@ pub enum AgentEvent { excluded: Vec, /// Number of tools left exposed to the model. remaining: usize, + /// Per-tool reason a [`crate::tool::toolset::ToolSet`] adaptor + /// changed or withheld a tool this turn, keyed by the tool's + /// original name. + /// + /// Additive (`docs/sdk-gaps/tools.md` §9's "explainable exposure + /// decisions"): `#[serde(default)]` keeps events recorded before + /// this field existed deserializable, and a middleware that only + /// reports `excluded` (no explanations) leaves this empty rather + /// than failing to construct the event. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + explanations: Vec<(String, crate::tool::ToolExposureExplanation)>, }, /// A tool invocation has been dispatched. @@ -204,6 +246,14 @@ pub enum AgentEvent { /// event itself rather than a live outcome side-channel. #[serde(default, skip_serializing_if = "Option::is_none")] error: Option, + /// Host-only metadata the tool attached to its result + /// (`tinytools::ToolResult::metadata`, B2). Carried here and on + /// [`crate::middleware::AgentRun::tool_metadata`] for events, + /// persistence, and telemetry; **never** rendered into the transcript + /// the model sees. Present regardless of payload capture: it is the + /// tool's deliberate host-facing channel, not captured I/O. + #[serde(default, skip_serializing_if = "Option::is_none")] + metadata: Option, }, /// A tool invocation failed and the run is propagating the error rather @@ -238,6 +288,24 @@ pub enum AgentEvent { error: String, }, + /// A resumed run reconciled an unresolved tool-effect-ledger row left + /// behind by an interrupted prior attempt (B5). + /// + /// Emitted by + /// [`crate::runtime::AgentHarness::reconcile_tool_effects`] for each + /// `started`-but-never-settled effect belonging to the last assistant + /// tool-call turn, once it has decided what to do per the tool's + /// [`tinytools::ToolReplay`] declaration. + ToolEffectReconciled { + /// Identifier of the reconciled tool call. + call_id: CallId, + /// What the reconciliation did: `"re_execute"` when the call was left + /// pending for the loop to run again (`ToolReplay::Safe`), or + /// `"interrupted"` when a synthesized tool-error result was appended + /// instead (`ToolReplay::Never`). + action: String, + }, + /// A model call failed and the run is propagating the error. /// /// The terminal partner of [`AgentEvent::ModelStarted`] on the error path; @@ -507,6 +575,58 @@ pub enum AgentEvent { to_tokens: u64, }, + /// A durable, rule-driven compaction ran and produced a + /// [`crate::summarization::CompactionRecord`]. + /// + /// Distinguished from [`Self::Compressed`] (the older, simpler + /// event `ContextCompressionMiddleware`'s original `before_model` path + /// emits) by carrying [`crate::summarization::CompactionReason`] and by + /// always being emitted for a compaction produced through + /// `crate::summarization::compaction` — including the + /// overflow → compact → retry recovery path, which has no other event of + /// its own. Both events fire for the same compaction on the `before_model` + /// path; a listener that only cares about *whether* the transcript shrank + /// can ignore `reason` and treat this exactly like `Compressed`. + Compacted { + /// Why this compaction ran. + reason: crate::summarization::CompactionReason, + /// Estimated total tokens of the transcript before compaction. + tokens_before: u64, + /// Estimated total tokens of the transcript after compaction. + tokens_after: u64, + }, + + /// The final turn's structured-output extraction failed schema + /// validation, or a registered + /// [`crate::structured::OutputValidator`] rejected the value with + /// [`crate::error::TinyAgentsError::ModelRetry`], and the loop is + /// re-asking the model instead of failing the run (A3's + /// output-validation retry loop; see + /// [`crate::runtime::RunPolicy::output_retry`]). + OutputRetry { + /// The 1-based retry attempt this event reports (1 is the first + /// re-ask after the original extraction failed). + attempt: u8, + /// The extraction/validation error handed back to the model as the + /// repair prompt. + error: String, + }, + + /// The agent loop appended one or more messages from a + /// [`crate::run_queue::RunQueue`] lane to the working transcript at a + /// safe turn boundary (A4): `Steer` after a tool batch or at a natural + /// finish, `Followup` at a natural finish. Emitted once per boundary + /// with the number of messages applied; `Collect` items never produce + /// this event because they are not applied to the transcript. Payload + /// text is deliberately not carried (events are payload-free by default). + QueuedMessageApplied { + /// Which lane the messages came from. + lane: crate::run_queue::QueueLane, + /// How many messages were appended at this boundary (`1` under + /// [`QueueMode::OneAtATime`][crate::run_queue::QueueMode::OneAtATime]). + count: usize, + }, + /// A graph routing decision produced a named route. RouteSelected { /// The route name chosen by the router. @@ -613,10 +733,30 @@ pub enum AgentEvent { message: String, }, + /// An application-defined event a tool (or any holder of the run's + /// [`EventSink`][crate::events::EventSink]) emitted through + /// [`ToolExecutionContext::custom`][crate::tool::ToolExecutionContext::custom] + /// (B1). The harness attaches no meaning to `payload`; it exists so a + /// tool can report structured progress — a download percentage, an + /// intermediate finding, a UI hint — on the same ordered stream as the + /// loop's own events, without the harness growing a variant per use. + Custom { + /// The tool call the event was emitted from, when it came from a + /// tool; `None` when emitted outside a call. + #[serde(default, skip_serializing_if = "Option::is_none")] + call_id: Option, + /// Application-defined payload, passed through verbatim. + payload: serde_json::Value, + }, + /// A middleware hook reported a failure. /// - /// Defined for future emit alongside [`AgentEvent::MiddlewareStarted`] / - /// [`AgentEvent::MiddlewareCompleted`] so a failing hook is observable. + /// Emitted by the lifecycle-hook driver ([`crate::middleware`]'s + /// `run_stack_hook!` macro) immediately after + /// [`AgentEvent::MiddlewareCompleted`] when a hook returns `Err`, so a + /// failing middleware is observable alongside + /// [`AgentEvent::MiddlewareStarted`] / [`AgentEvent::MiddlewareCompleted`] + /// instead of only surfacing as the run's terminal error. MiddlewareFailed { /// Registered name of the middleware that failed. name: String, @@ -624,6 +764,22 @@ pub enum AgentEvent { error: String, }, + /// A cross-provider handoff transform rewrote part of the outgoing + /// transcript immediately before a model call, because it carried + /// assistant content from a different provider/api/model than the one + /// about to receive it (a mid-session model switch, an explicit + /// per-request override, or a fallback to a different provider). Emitted + /// only when at least one message changed — same-origin runs (the + /// common case) never emit this. + /// + /// See the harness's cross-provider handoff transform for the exact + /// rules (redacted/signed thinking, tool-call id normalization, image + /// downgrade). + HandoffTransformApplied { + /// Number of messages rewritten by the transform for this call. + changes: usize, + }, + /// A streaming model call's chunk stream was closed (gracefully or by /// cancellation). /// @@ -685,10 +841,14 @@ impl AgentEvent { AgentEvent::ToolsAdvertised { .. } => "tool.advertised", AgentEvent::ToolSearched { .. } => "tool.searched", AgentEvent::DeferredToolCall { .. } => "tool.deferred_call", + AgentEvent::ToolDeferred { .. } => "tool.deferred", + AgentEvent::ToolApproved { .. } => "tool.approved", + AgentEvent::ToolDenied { .. } => "tool.denied", AgentEvent::ToolsFiltered { .. } => "tool.filtered", AgentEvent::ToolStarted { .. } => "tool.started", AgentEvent::ToolCompleted { .. } => "tool.completed", AgentEvent::ToolFailed { .. } => "tool.failed", + AgentEvent::ToolEffectReconciled { .. } => "tool.effect_reconciled", AgentEvent::ModelFailed { .. } => "model.failed", AgentEvent::SubAgentFailed { .. } => "subagent.failed", AgentEvent::UnknownToolCall { .. } => "tool.unknown", @@ -715,6 +875,9 @@ impl AgentEvent { AgentEvent::SubAgentReused { .. } => "subagent.reused", AgentEvent::Steered { .. } => "agent.steered", AgentEvent::Compressed { .. } => "context.compressed", + AgentEvent::Compacted { .. } => "context.compacted", + AgentEvent::OutputRetry { .. } => "output.retry", + AgentEvent::QueuedMessageApplied { .. } => "queue.applied", AgentEvent::RouteSelected { .. } => "route.selected", AgentEvent::UsageRecorded { .. } => "usage.recorded", AgentEvent::CostRecorded { .. } => "cost.recorded", @@ -722,7 +885,9 @@ impl AgentEvent { AgentEvent::MemoryLoaded => "memory.loaded", AgentEvent::MemorySaved => "memory.saved", AgentEvent::ToolProgress { .. } => "tool.progress", + AgentEvent::Custom { .. } => "custom", AgentEvent::MiddlewareFailed { .. } => "middleware.failed", + AgentEvent::HandoffTransformApplied { .. } => "handoff.transform_applied", AgentEvent::StreamClosed => "stream.closed", AgentEvent::RunCompleted { .. } => "run.completed", AgentEvent::RunFailed { .. } => "run.failed", diff --git a/crates/tinyagents-harness/src/handoff.rs b/crates/tinyagents-harness/src/handoff.rs index e646ef41..b2b44e89 100644 --- a/crates/tinyagents-harness/src/handoff.rs +++ b/crates/tinyagents-harness/src/handoff.rs @@ -19,6 +19,15 @@ //! * the [`ResultHandoffCache`] store itself (FIFO-evicting, `Arc`-shared); //! * the [`build_handoff_placeholder`] renderer used when rewriting tool //! results into history. +//! +//! # Not on the agent loop path (M-10) +//! +//! This is a host utility, not something [`crate::agent_loop`] calls on its +//! own: nothing in the built-in loop invokes [`apply_handoff`] or registers +//! the extraction tool it references. A host wires this in itself — calling +//! `apply_handoff` on each tool result before it is appended to history, and +//! registering an extraction tool (named per [`HandoffConfig::extractor_tool_name`]) +//! that reads from the same [`ResultHandoffCache`]. use std::collections::HashMap; use std::sync::Mutex as StdMutex; @@ -48,6 +57,49 @@ pub const HANDOFF_PREVIEW_CHARS: usize = 1500; /// user/orchestrator to narrow the request. pub const HANDOFF_MAX_ENTRIES: usize = 8; +// ── Host configuration ─────────────────────────────────────────────────────── + +/// Host-specific naming this module needs but does not own: which tool name +/// is the extractor (so its own output passes through uncleaned/unstashed), +/// and which literal prefixes mark a result as already an error. +/// +/// Both were previously hardcoded to one particular host's conventions +/// (`extract_from_result`, a bare `result_text.starts_with("Error")`) even +/// though the rest of this module is host-agnostic (M-9). A different host +/// — a different extractor tool name, or error-carrying results that do not +/// start with the literal word "Error" — could not use this cache without +/// forking the module. [`HandoffConfig::default`] reproduces the historical +/// behaviour exactly, so existing callers are unaffected. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct HandoffConfig { + /// The tool name whose own output skips cleaning/stashing (it is already + /// a narrowed, host-curated response to a targeted query). + pub extractor_tool_name: String, + /// Literal prefixes that mark `result_text` as already an error, which + /// also skips cleaning/stashing (an error message should reach the model + /// verbatim, not truncated or placeholder-replaced). + pub error_prefixes: Vec, +} + +impl Default for HandoffConfig { + fn default() -> Self { + Self { + extractor_tool_name: "extract_from_result".to_string(), + error_prefixes: vec!["Error".to_string()], + } + } +} + +impl HandoffConfig { + /// `true` when `result_text` starts with any of + /// [`HandoffConfig::error_prefixes`]. + fn is_error_result(&self, result_text: &str) -> bool { + self.error_prefixes + .iter() + .any(|prefix| result_text.starts_with(prefix.as_str())) + } +} + // ── Store ────────────────────────────────────────────────────────────────── /// Per-spawn cache of oversized tool payloads. One instance is built at @@ -127,22 +179,28 @@ impl ResultHandoffCache { /// the path in tests — and because the alternative this replaced was an /// environment-variable backdoor named after one particular host. Pass /// [`HANDOFF_OVERSIZE_THRESHOLD_TOKENS`] for the default. +/// +/// `config` supplies the host's extractor tool name and error-prefix +/// heuristics (M-9); pass [`HandoffConfig::default`] to reproduce the +/// historical hardcoded behaviour. pub fn apply_handoff( cache: &ResultHandoffCache, + config: &HandoffConfig, tool_name: &str, task_id: &str, agent_id: &str, result_text: String, threshold_tokens: usize, ) -> String { - let skip_cleaning = tool_name == "extract_from_result" || result_text.starts_with("Error"); + let skip_cleaning = + tool_name == config.extractor_tool_name || config.is_error_result(&result_text); let cleaned = if skip_cleaning { result_text } else { let pre_len = result_text.len(); let cleaned = clean_tool_output(&result_text); if cleaned.len() < pre_len { - tinyagents_tracing::debug!( + tracing::debug!( tool = %tool_name, before_bytes = pre_len, after_bytes = cleaned.len(), @@ -155,8 +213,8 @@ pub fn apply_handoff( let tokens = cleaned.len().div_ceil(4); if !skip_cleaning && tokens > threshold_tokens { let id = cache.store(tool_name.to_string(), cleaned.clone()); - let placeholder = build_handoff_placeholder(tool_name, &id, &cleaned); - tinyagents_tracing::info!( + let placeholder = build_handoff_placeholder(config, tool_name, &id, &cleaned); + tracing::info!( task_id = %task_id, agent_id = %agent_id, tool = %tool_name, @@ -176,22 +234,28 @@ pub fn apply_handoff( /// Build the placeholder text that replaces an oversized tool result in /// the sub-agent's history. Shows the payload size (estimated tokens and -/// raw bytes), a preview, and a call shape for the `extract_from_result` -/// tool. The sub-agent decides whether to answer from the preview or -/// dispatch the extractor. +/// raw bytes), a preview, and a call shape for the configured extractor +/// tool ([`HandoffConfig::extractor_tool_name`]). The sub-agent decides +/// whether to answer from the preview or dispatch the extractor. /// /// Token count is estimated at ~4 chars/token (same heuristic as the /// trigger threshold in [`HANDOFF_OVERSIZE_THRESHOLD_TOKENS`]), so the /// unit the sub-agent sees matches the unit the runtime used to decide /// to hand off in the first place. -pub fn build_handoff_placeholder(tool_name: &str, result_id: &str, raw: &str) -> String { +pub fn build_handoff_placeholder( + config: &HandoffConfig, + tool_name: &str, + result_id: &str, + raw: &str, +) -> String { let preview: String = raw.chars().take(HANDOFF_PREVIEW_CHARS).collect(); let raw_tokens = raw.len().div_ceil(4); + let extractor = &config.extractor_tool_name; format!( "[oversized tool output: {raw_tokens} tokens ({raw_bytes} bytes) — stashed as result_id=\"{result_id}\"]\n\ Preview (first {preview_chars} chars):\n{preview}\n\n\ If the preview does not answer your task, call:\n\ - extract_from_result(result_id=\"{result_id}\", query=\"\")\n\ + {extractor}(result_id=\"{result_id}\", query=\"\")\n\ Good queries name the exact fields/identifiers you need \ (e.g. \"subject and sender of the 5 most recent messages\"). \ Tool: {tool_name}", diff --git a/crates/tinyagents-harness/src/handoff_test.rs b/crates/tinyagents-harness/src/handoff_test.rs index 12c606ff..ac6b3a4a 100644 --- a/crates/tinyagents-harness/src/handoff_test.rs +++ b/crates/tinyagents-harness/src/handoff_test.rs @@ -68,7 +68,15 @@ fn eviction_is_fifo_and_bounded() { #[test] fn a_small_result_passes_through_untouched_and_is_not_cached() { let c = cache(); - let out = apply_handoff(&c, "search", "task-1", "agent-1", "small".to_string(), 10); + let out = apply_handoff( + &c, + &HandoffConfig::default(), + "search", + "task-1", + "agent-1", + "small".to_string(), + 10, + ); assert_eq!(out, "small"); assert!( c.get("res_1").is_none(), @@ -80,7 +88,15 @@ fn a_small_result_passes_through_untouched_and_is_not_cached() { fn an_oversized_result_is_stashed_and_replaced_by_a_placeholder() { let c = cache(); let raw = big(4_000); // ~1000 tokens at the 4-chars/token heuristic - let out = apply_handoff(&c, "gmail_list", "task-1", "agent-1", raw.clone(), 10); + let out = apply_handoff( + &c, + &HandoffConfig::default(), + "gmail_list", + "task-1", + "agent-1", + raw.clone(), + 10, + ); assert_ne!(out, raw, "the raw payload must not reach history"); assert!(out.contains("oversized tool output")); @@ -105,8 +121,24 @@ fn the_threshold_is_honoured_in_both_directions() { // Same payload, two thresholds: this is the parameter that replaced the // env-var backdoor, so it has to actually decide the outcome. let raw = big(400); // ~100 tokens - let below = apply_handoff(&cache(), "t", "task", "agent", raw.clone(), 10); - let above = apply_handoff(&cache(), "t", "task", "agent", raw.clone(), 10_000); + let below = apply_handoff( + &cache(), + &HandoffConfig::default(), + "t", + "task", + "agent", + raw.clone(), + 10, + ); + let above = apply_handoff( + &cache(), + &HandoffConfig::default(), + "t", + "task", + "agent", + raw.clone(), + 10_000, + ); assert!(below.contains("oversized tool output")); assert_eq!(above, raw); } @@ -117,8 +149,71 @@ fn an_error_result_passes_through_however_large() { // behind an extraction call would hide the failure it needs to react to. let c = cache(); let err = format!("Error: {}", big(8_000)); - let out = apply_handoff(&c, "gmail_list", "task", "agent", err.clone(), 1); + let out = apply_handoff( + &c, + &HandoffConfig::default(), + "gmail_list", + "task", + "agent", + err.clone(), + 1, + ); + assert_eq!(out, err); +} + +/// M-9 regression: the extractor tool name and the "already an error" +/// prefix used to be hardcoded (`extract_from_result`, a bare +/// `starts_with("Error")`). A host with different conventions must be able +/// to configure both instead of forking the module. +#[test] +fn a_host_can_configure_its_own_extractor_name_and_error_prefix() { + let config = HandoffConfig { + extractor_tool_name: "fetch_full_result".to_string(), + error_prefixes: vec!["FAILED:".to_string()], + }; + + // The custom extractor's own output is never re-stashed. + let c = cache(); + let raw = big(8_000); + let out = apply_handoff( + &c, + &config, + "fetch_full_result", + "task", + "agent", + raw.clone(), + 1, + ); + assert_eq!(out, raw); + + // A result whose custom error prefix matches passes through unstashed, + // even though it does not start with the historical "Error". + let c = cache(); + let err = format!("FAILED: {}", big(8_000)); + let out = apply_handoff(&c, &config, "gmail_list", "task", "agent", err.clone(), 1); assert_eq!(out, err); + + // A result starting with the *historical* "Error" prefix is NOT treated + // as an error under this custom config, and is stashed like any other + // oversized payload — the heuristic is fully replaced, not merged. + let c = cache(); + let historical_error = format!("Error: {}", big(8_000)); + let out = apply_handoff( + &c, + &config, + "gmail_list", + "task", + "agent", + historical_error.clone(), + 1, + ); + assert_ne!(out, historical_error); + assert!(out.contains("oversized tool output")); + + // The placeholder advertises the configured extractor tool name. + let placeholder = build_handoff_placeholder(&config, "gmail_list", "res_1", "payload"); + assert!(placeholder.contains("fetch_full_result")); + assert!(!placeholder.contains("extract_from_result")); } #[test] @@ -128,7 +223,15 @@ fn an_extraction_result_is_never_re_stashed() { // model another placeholder — a loop that never converges. let c = cache(); let raw = big(8_000); - let out = apply_handoff(&c, "extract_from_result", "task", "agent", raw.clone(), 1); + let out = apply_handoff( + &c, + &HandoffConfig::default(), + "extract_from_result", + "task", + "agent", + raw.clone(), + 1, + ); assert_eq!(out, raw); } @@ -137,7 +240,7 @@ fn an_extraction_result_is_never_re_stashed() { #[test] fn the_placeholder_reports_size_and_previews_the_head() { let raw = format!("HEAD-MARKER{}", big(5_000)); - let text = build_handoff_placeholder("gmail_list", "res_1", &raw); + let text = build_handoff_placeholder(&HandoffConfig::default(), "gmail_list", "res_1", &raw); assert!(text.contains("res_1")); assert!(text.contains("gmail_list")); @@ -157,7 +260,7 @@ fn the_placeholder_reports_size_and_previews_the_head() { #[test] fn a_short_payload_previews_whole_without_padding() { - let text = build_handoff_placeholder("t", "res_1", "tiny"); + let text = build_handoff_placeholder(&HandoffConfig::default(), "t", "res_1", "tiny"); assert!(text.contains("tiny")); // The preview length is reported to the model; claiming the full budget // for a 4-char payload would misdescribe what it is looking at. @@ -167,7 +270,7 @@ fn a_short_payload_previews_whole_without_padding() { #[test] fn the_preview_is_capped_for_a_large_payload() { let raw = big(HANDOFF_PREVIEW_CHARS * 4); - let text = build_handoff_placeholder("t", "res_1", &raw); + let text = build_handoff_placeholder(&HandoffConfig::default(), "t", "res_1", &raw); assert!(text.contains(&format!("first {HANDOFF_PREVIEW_CHARS} chars"))); } @@ -176,6 +279,6 @@ fn the_preview_never_splits_a_multibyte_character() { // Taken by chars, not bytes — a byte-indexed cut here would panic rather // than merely misformat. let raw = "é".repeat(HANDOFF_PREVIEW_CHARS * 2); - let text = build_handoff_placeholder("t", "res_1", &raw); + let text = build_handoff_placeholder(&HandoffConfig::default(), "t", "res_1", &raw); assert!(text.is_char_boundary(text.len())); } diff --git a/crates/tinyagents-harness/src/host/mod.rs b/crates/tinyagents-harness/src/host/mod.rs index c180d57f..6a13fa25 100644 --- a/crates/tinyagents-harness/src/host/mod.rs +++ b/crates/tinyagents-harness/src/host/mod.rs @@ -103,6 +103,16 @@ pub struct HostCapabilities { /// Procedural memory of how this agent has performed before. `None` means /// no experience is recorded or recalled. pub experience: Option>, + /// Whether a resolved [`AgentDefinition`] that declares no tools (an + /// empty or absent `tools` list) denies every tool, instead of granting + /// the whole registered catalogue. + /// + /// Defaults to `true` (fail-closed): policy metadata that is missing is + /// treated as "nothing authorized", not as "unrestricted" (I-9). Set to + /// `false` only to restore the legacy behavior for a host that relied on + /// an empty list meaning unrestricted — new hosts should leave this on + /// and declare tools explicitly. + pub fail_closed_tool_allowlist: bool, } impl HostCapabilities { @@ -129,9 +139,21 @@ impl HostCapabilities { learning: None, tool_outcomes: None, experience: None, + fail_closed_tool_allowlist: true, } } + /// Opts this host out of the default fail-closed tool allow-list, + /// restoring the legacy behavior where a definition that declares no + /// tools is granted the entire registered catalogue. + /// + /// Prefer declaring tools explicitly per definition instead of calling + /// this; it exists for hosts migrating from the pre-I-9 behavior. + pub fn with_legacy_unrestricted_tool_allowlist(mut self) -> Self { + self.fail_closed_tool_allowlist = false; + self + } + /// Supplies durable user memory. pub fn with_memory(mut self, memory: Arc) -> Self { self.memory = Some(memory); @@ -185,6 +207,7 @@ impl Clone for HostCapabilities { learning: self.learning.clone(), tool_outcomes: self.tool_outcomes.clone(), experience: self.experience.clone(), + fail_closed_tool_allowlist: self.fail_closed_tool_allowlist, } } } diff --git a/crates/tinyagents-harness/src/ids/test.rs b/crates/tinyagents-harness/src/ids/test.rs index 2ecd603f..c2dfad6e 100644 --- a/crates/tinyagents-harness/src/ids/test.rs +++ b/crates/tinyagents-harness/src/ids/test.rs @@ -61,6 +61,14 @@ fn status_and_phase_use_snake_case() { serde_json::to_string(&ExecutionStatus::Interrupted).unwrap(), "\"interrupted\"" ); + assert_eq!( + serde_json::to_string(&ExecutionStatus::Drained).unwrap(), + "\"drained\"" + ); + assert_eq!( + serde_json::from_str::("\"drained\"").unwrap(), + ExecutionStatus::Drained + ); assert_eq!( serde_json::to_string(&HarnessPhase::BuildingRequest).unwrap(), "\"building_request\"" diff --git a/crates/tinyagents-harness/src/ids/types.rs b/crates/tinyagents-harness/src/ids/types.rs index b6f2f4d5..00391524 100644 --- a/crates/tinyagents-harness/src/ids/types.rs +++ b/crates/tinyagents-harness/src/ids/types.rs @@ -23,7 +23,11 @@ pub struct RunId(pub(crate) String); pub struct ThreadId(pub(crate) String); /// Identifies an individual model or tool call inside a run. -#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +/// +/// Also `Ord`, so it can key the `BTreeMap`s in +/// [`crate::tool::DeferredToolRequests`]/[`crate::tool::DeferredToolResults`] +/// deterministically. +#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)] pub struct CallId(pub(crate) String); /// Identifies a single emitted harness event. @@ -95,6 +99,11 @@ pub enum ExecutionStatus { Failed, /// Cancelled before completion. Cancelled, + /// Stopped gracefully at a superstep boundary by a drain request (see + /// `tinyagents_graph::DrainSignal`): the in-flight superstep was allowed + /// to finish and the still-pending work was checkpointed, so the run is + /// resumable exactly like an interrupted one. + Drained, } /// The active operation within a harness run, used for compact status reads. diff --git a/crates/tinyagents-harness/src/lib.rs b/crates/tinyagents-harness/src/lib.rs index acbbb21c..447f0eef 100644 --- a/crates/tinyagents-harness/src/lib.rs +++ b/crates/tinyagents-harness/src/lib.rs @@ -9,16 +9,61 @@ //! The harness is intentionally split by feature. Each submodule owns one //! substantial part of model/tool orchestration so the implementation can grow //! without creating one large runtime file. - -#![cfg_attr( - not(feature = "tracing"), - allow(dead_code, unused_imports, unused_variables) -)] +//! +//! # Cargo features +//! +//! - `sqlite` — durable, file-backed stores (`rusqlite`). +//! - `builtin-tools` — the bundled [`tools`] (currently the time tool). The +//! old name `tools` is kept as a deprecated alias. +//! - `multimodal` — image/audio/binary content resolution ([`multimodal`]), +//! pulling in `reqwest` and `flate2`. +//! - `claude-code` — the Claude Code CLI and Claude Agent SDK provider +//! adapters under [`providers`]. +//! - `langfuse` — the Langfuse observability exporter under [`observability`], +//! pulling in `reqwest`. +//! - `tracing` — a no-op compatibility alias; tracing instrumentation is +//! always compiled in. +//! +//! `claude-code` and `langfuse` are part of `default` so existing consumers +//! see no change; disable default features to opt out of either. +//! +//! # Host utilities not on the agent loop path (M-10) +//! +//! [`handoff`] and the [`memory`] module's +//! [`memory::ChatHistory`]/[`memory::ShortTermMemory`] are exported for a +//! host to build on, but [`agent_loop`] does not call into any of them on its +//! own — they are opt-in plumbing, not implicit loop behavior. (A +//! [`run_queue::RunQueue`] *is* drained by the loop once attached via +//! [`context::RunContext::with_run_queue`]; see that module's docs.) +//! +//! - [`handoff`] is a progressive-disclosure cache for oversized tool +//! results; a host calls [`handoff::apply_handoff`] itself before +//! appending a tool result to history, and registers an extraction tool +//! that reads the same [`handoff::ResultHandoffCache`]. +//! - [`memory::ChatHistory`]/[`memory::ShortTermMemory`] persist a thread's +//! transcript across runs; a host reads history into a run's `input` and +//! appends the run's messages back afterward. +//! +//! Wiring any of these directly into the loop is deliberately future work +//! rather than default behavior, so a host that does not need one pays +//! nothing for it. +//! +//! # Vendor re-exports +//! +//! The harness pins exact versions of the `tinyinference-llm`, `tinytools`, +//! and `tinytools-agent` vendor crates and exposes their public types +//! (e.g. `ChatMessage`, tool schemas) across its own API. Downstream crates +//! must reach those types through [`tinyinference_llm`], [`tinytools`], and +//! [`tinytools_agent`] re-exported here rather than depending on the vendor +//! crates directly, or the compiler will see two distinct copies of the same +//! type. pub mod agent_loop; pub mod artifacts; +mod blocking; pub mod cache; pub mod cancel; +pub mod capability; pub mod config; pub mod context; pub mod cost; @@ -36,6 +81,7 @@ pub mod no_progress; pub mod observability; pub mod prompt; pub mod providers; +pub(crate) mod relaxed_json; pub mod retriever; pub mod retry; pub mod run_queue; @@ -48,10 +94,25 @@ pub mod summarization; pub mod testkit; pub mod token_estimation; pub mod tool; -#[cfg(feature = "tools")] +#[cfg(feature = "builtin-tools")] pub mod tools; +/// Re-exported vendor crates. Downstream consumers should reach these +/// dependencies' types through these re-exports (e.g. +/// `tinyagents_harness::tinyinference_llm::ChatMessage`) rather than adding +/// their own `tinyinference-llm` / `tinytools` / `tinytools-agent` +/// dependency, since the harness pins exact vendor versions and a second, +/// independent dependency would produce a duplicate, incompatible copy of +/// the same types. +pub use tinyinference_llm; +pub use tinytools; +pub use tinytools_agent; + pub use cancel::CancellationToken; +pub use capability::{ + Capability, CapabilityToolSet, LOAD_CAPABILITY_TOOL_NAME, LoadCapabilityTool, + ModelRequestDefaults, +}; pub use cost::CostTotals; pub use error::{Result, TinyAgentsError}; pub use ids::*; @@ -64,10 +125,13 @@ pub use no_progress::{ pub use observability::{ AgentCallLatency, AgentLatencyMetrics, AgentObservation, FanOutSink, HarnessEventJournal, HarnessStatusStore, InMemoryEventJournal, InMemoryStatusStore, JournalSink, JsonlSink, + RedactingSink, StoreEventJournal, +}; +#[cfg(feature = "langfuse")] +pub use observability::{ LangfuseAuth, LangfuseClient, LangfuseScore, LangfuseScoreValue, LangfuseTraceConfig, - ProcessProfile, ProcessProfiler, ProcessSnapshot, RedactingSink, SinkHealth, StoreEventJournal, }; -pub use run_queue::{QueueLane, QueueStatus, RunQueue}; +pub use run_queue::{QueueLane, QueueMode, QueueStatus, RunQueue, RunQueueHandle}; pub use steering::{ SteeringCommand, SteeringCommandKind, SteeringHandle, SteeringOutcome, SteeringPolicy, }; diff --git a/crates/tinyagents-harness/src/limits/mod.rs b/crates/tinyagents-harness/src/limits/mod.rs index 66b1a234..9698d1de 100644 --- a/crates/tinyagents-harness/src/limits/mod.rs +++ b/crates/tinyagents-harness/src/limits/mod.rs @@ -75,6 +75,14 @@ impl RunLimits { self.behavior = behavior; self } + + /// Caps how many tool calls a concurrently-executed batch may run at + /// once. `None` removes the cap. See + /// [`RunLimits::max_tool_concurrency`]. + pub fn with_max_tool_concurrency(mut self, n: Option) -> Self { + self.max_tool_concurrency = n; + self + } } /// Tracks live counters for a single harness run and enforces [`RunLimits`]. @@ -102,6 +110,22 @@ impl LimitTracker { } } + /// Resets the wall-clock start to now, leaving the call counters and + /// limits untouched. + /// + /// [`RunContext::new`][crate::context::RunContext::new] constructs the + /// tracker (and therefore stamps `started_at`) at context-construction + /// time, which is not always the same moment the run actually starts + /// doing work — a context built ahead of time and queued, or reused + /// across a retry of the *surrounding* host operation, would otherwise + /// have its wall-clock deadline silently burn down before the agent loop + /// issues its first model call (M-8). The agent loop calls this at the + /// top of the run so the deadline is always measured from when the run + /// actually began. + pub fn restart(&mut self) { + self.started_at = Instant::now(); + } + /// Records one model call and returns an error if the cap is exceeded. /// /// The counter is incremented **before** the check so the limit is @@ -148,7 +172,7 @@ impl LimitTracker { fn exhausted(&self, kind: LimitKind, cap: usize) -> Result { match self.limits.behavior { LimitBehavior::Error => { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::limits", limit_kind = kind.as_str(), cap, @@ -163,7 +187,7 @@ impl LimitTracker { ))) } LimitBehavior::StopWithPartial => { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::limits", limit_kind = kind.as_str(), cap, @@ -286,7 +310,7 @@ impl LimitTracker { /// defaulted. See the note on [`LimitTracker::tighten_call_limits`] for what /// the agent loop has to do about it. pub fn sync_call_limits(&mut self, max_model_calls: usize, max_tool_calls: usize) { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::limits", from_model_calls = self.limits.max_model_calls, from_tool_calls = self.limits.max_tool_calls, @@ -323,7 +347,7 @@ impl LimitTracker { pub fn tighten_call_limits(&mut self, max_model_calls: usize, max_tool_calls: usize) { let model = self.limits.max_model_calls.min(max_model_calls); let tool = self.limits.max_tool_calls.min(max_tool_calls); - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::limits", from_model_calls = self.limits.max_model_calls, from_tool_calls = self.limits.max_tool_calls, diff --git a/crates/tinyagents-harness/src/limits/test.rs b/crates/tinyagents-harness/src/limits/test.rs index 52dcc2c0..c3f0a2a1 100644 --- a/crates/tinyagents-harness/src/limits/test.rs +++ b/crates/tinyagents-harness/src/limits/test.rs @@ -154,6 +154,31 @@ fn rollback_tool_calls_uncounts_calls_that_never_ran() { assert_eq!(tracker.tool_calls(), 0); } +#[test] +fn restart_resets_the_wall_clock_start_without_touching_counters() { + // M-8 regression: `started_at` is stamped when the tracker (via + // `RunContext::new`) is constructed, not when the run actually starts + // doing work. A context built ahead of time and left to sit burns down + // its deadline before the first model call. `restart` must reset the + // clock while leaving the call counters alone. + let mut tracker = + LimitTracker::new(RunLimits::default().with_max_wall_clock_ms(Some(1_000_000))); + tracker.record_model_call().unwrap(); + tracker.record_tool_call().unwrap(); + std::thread::sleep(std::time::Duration::from_millis(20)); + let elapsed_before_restart = tracker.elapsed(); + assert!(elapsed_before_restart >= std::time::Duration::from_millis(20)); + + tracker.restart(); + + assert!( + tracker.elapsed() < elapsed_before_restart, + "restart should reset the wall-clock start, not extend it" + ); + assert_eq!(tracker.model_calls(), 1, "counters must survive a restart"); + assert_eq!(tracker.tool_calls(), 1, "counters must survive a restart"); +} + #[test] fn limit_kind_labels_match_the_event_layer() { // The limits module keeps its own `LimitKind` so it need not depend on the diff --git a/crates/tinyagents-harness/src/limits/types.rs b/crates/tinyagents-harness/src/limits/types.rs index fe3a2f35..6b0d68b7 100644 --- a/crates/tinyagents-harness/src/limits/types.rs +++ b/crates/tinyagents-harness/src/limits/types.rs @@ -69,6 +69,17 @@ pub struct RunLimits { /// What the run should do when a call cap is reached. Defaults to /// [`LimitBehavior::Error`], which is the historical behaviour. pub behavior: LimitBehavior, + /// Caps how many tool calls in one concurrently-executed batch (see + /// [`should_execute_tools_concurrently`][crate::agent_loop] and its + /// module docs) may be in flight at once. `None` (the default) leaves the + /// batch unbounded — every eligible call in the turn starts together, as + /// before this field existed. + /// + /// Only applies to the concurrent tool path; the serial path always runs + /// one call at a time regardless of this setting. A `Some(0)` behaves the + /// same as `Some(1)`: at least one call must be in flight to make + /// progress. + pub max_tool_concurrency: Option, } /// What a run does when it reaches a configured call cap. @@ -180,6 +191,7 @@ impl Default for RunLimits { max_retries_per_call: 3, max_depth: Self::DEFAULT_MAX_DEPTH, behavior: LimitBehavior::Error, + max_tool_concurrency: None, } } } diff --git a/crates/tinyagents-harness/src/middleware/library/README.md b/crates/tinyagents-harness/src/middleware/library/README.md index 5e0ffe67..eadfc0c4 100644 --- a/crates/tinyagents-harness/src/middleware/library/README.md +++ b/crates/tinyagents-harness/src/middleware/library/README.md @@ -103,6 +103,24 @@ Grouped by extension shape: `tokio::time` paused-time tests; `RateLimitMiddleware` takes an injectable clock and configurable poll interval. Preserve this when adding new middleware: prefer computing a delay over unconditionally awaiting one. +- **Retry engines are coordinated, not yet unified (R-3, partial).** + `invoke_model_resolving` (the loop's own base-call attempt engine), + `RetryMiddleware`, and `ModelFallbackMiddleware` each still implement their + own attempt loop. Double-retrying is prevented by + `ModelMiddleware::overrides_retry` — a registered `RetryMiddleware` (or any + middleware overriding it) makes the base call skip its own retry loop + entirely, so attempts are bounded by whichever engine actually runs, never + by their product (`middleware::library::test:: + retry_middleware_correlates_retry_scheduled_with_the_loops_call_id` and + `agent_loop::test:: + retry_middleware_and_run_policy_retry_do_not_multiply_attempts` cover this). + `RetryMiddleware` also mirrors the loop's own call id onto + `RunContext::active_model_call` so its `RetryScheduled` events correlate + with the `ModelStarted`/`ModelCompleted` pair for the same attempt. Making + the loop's engine the *only* one and turning these middlewares into pure + policy overrides (`code-review-harness.md` R-3) is still open — it changes + their semantics from "execute" to "configure" and touches every test that + exercises them directly via `MiddlewareStack::run_wrapped_model`. - **Budget reservations are keyed by `RunContext::instance_id`, not `run_id`.** Concurrent runs sharing a `BudgetTracker` may share a caller-supplied `run_id`; only the process-unique instance id keeps each run releasing diff --git a/crates/tinyagents-harness/src/middleware/library/budget.rs b/crates/tinyagents-harness/src/middleware/library/budget.rs index 89d4fe9a..9b23a953 100644 --- a/crates/tinyagents-harness/src/middleware/library/budget.rs +++ b/crates/tinyagents-harness/src/middleware/library/budget.rs @@ -213,6 +213,35 @@ impl Middleware for BudgetMidd self.label } + /// Control-outcome override (A1): a budget already exhausted *before* + /// this call is not a run failure — it is exactly the "stop cleanly with + /// whatever the run produced so far" case `MiddlewareControl::JumpTo` + /// `(LoopTarget::End)` exists for, so this stops the loop gracefully + /// instead of erroring the whole run out from under a partial transcript. + /// The preflight *reservation* check (a single oversized call, handled in + /// [`Self::before_model`] below) stays a hard `Err`: it is an admission + /// refusal for one call, not "the run is over". + async fn before_model_control( + &self, + ctx: &mut RunContext, + state: &State, + request: &mut ModelRequest, + ) -> Result { + { + let guard = self.tracker.lock_recovering(); + if let Some(reason) = self.limits.exceeded_reason(&guard) { + drop(guard); + ctx.emit(AgentEvent::BudgetExceeded { + reason, + blocked: true, + }); + return Ok(MiddlewareControl::JumpTo(crate::context::LoopTarget::End)); + } + } + self.before_model(ctx, state, request).await?; + Ok(MiddlewareControl::Continue) + } + async fn before_model( &self, ctx: &mut RunContext, @@ -227,7 +256,11 @@ impl Middleware for BudgetMidd let estimated = estimated_input_tokens(request); { let mut guard = self.tracker.lock_recovering(); - // (1) Already exhausted before this call. + // (1) Already exhausted before this call. Reachable when this + // hook is invoked directly (bypassing `before_model_control`, + // which the agent loop actually drives) — kept as a hard `Err` + // here for that direct-call case; see `before_model_control` for + // the loop's actual (graceful) behavior. if let Some(reason) = self.limits.exceeded_reason(&guard) { drop(guard); ctx.emit(AgentEvent::BudgetExceeded { @@ -296,7 +329,7 @@ impl Middleware for BudgetMidd // on a budget it never actually touched. The reservation is still // released above — that part is real bookkeeping. if response.served_from_cache { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::middleware", label = self.label, "[budget] skipping accounting for a cache-served response" diff --git a/crates/tinyagents-harness/src/middleware/library/context.rs b/crates/tinyagents-harness/src/middleware/library/context.rs index fa360653..134601b7 100644 --- a/crates/tinyagents-harness/src/middleware/library/context.rs +++ b/crates/tinyagents-harness/src/middleware/library/context.rs @@ -12,7 +12,9 @@ use crate::middleware::{ PromptCacheGuardMiddleware, }; use crate::summarization::{ - ConcatSummarizer, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy, trim_messages, + CompactionContext, CompactionDecision, CompactionReason, CompactionRecord, ConcatSummarizer, + OverflowClassifier, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy, + find_cut_point, summarize_with_split, trim_messages, }; // ── MessageTrimMiddleware ───────────────────────────────────────────────────── @@ -77,9 +79,41 @@ impl ContextCompressionMiddleware { records: std::sync::Mutex::new(std::collections::VecDeque::new()), max_records: DEFAULT_COMPRESSION_RECORD_CAP, on_failure: CompressionFailurePolicy::default(), + last_summary: std::sync::Mutex::new(None), + max_turn_tokens: None, + overflow_classifier: OverflowClassifier::default(), + before_compaction: None, } } + /// Sets the token budget above which a single turn handed to the + /// summarizer is split into two halves and merged (see + /// [`summarize_with_split`]). Unset (the default) never splits. + pub fn with_max_turn_tokens(mut self, max_turn_tokens: u64) -> Self { + self.max_turn_tokens = Some(max_turn_tokens); + self + } + + /// Replaces the [`OverflowClassifier`] consulted by the + /// overflow → compact → retry recovery path (this middleware's + /// [`ModelMiddleware::wrap_model`] implementation). Defaults to + /// [`OverflowClassifier::default`]'s built-in provider patterns. + pub fn with_overflow_classifier(mut self, classifier: OverflowClassifier) -> Self { + self.overflow_classifier = classifier; + self + } + + /// Installs a `before_compaction` hook consulted before every compaction + /// this middleware runs (proactive threshold or reactive overflow + /// recovery). See [`CompactionDecision`]. + pub fn with_before_compaction( + mut self, + hook: impl Fn(&CompactionContext) -> CompactionDecision + Send + Sync + 'static, + ) -> Self { + self.before_compaction = Some(std::sync::Arc::new(hook)); + self + } + /// Sets the [`CompressionFailurePolicy`] applied when the [`Summarizer`] /// returns an `Err`. Defaults to /// [`CompressionFailurePolicy::FallbackTrim`]. @@ -145,7 +179,7 @@ impl Middleware for ContextCom return Ok(()); } - let (to_summarize, mut to_keep) = self.policy.plan(&request.messages); + let (to_summarize, to_keep) = self.policy.plan(&request.messages); // Nothing old enough to compress (e.g. keep_last covers everything): // leave the transcript untouched rather than summarizing an empty set. if to_summarize.is_empty() { @@ -153,7 +187,62 @@ impl Middleware for ContextCom } let from_tokens = total_message_tokens(&request.messages); - let record = match self.summarizer.summarize(&to_summarize).await { + // `plan` splits by count (`non_system[..first_kept_index]` is exactly + // `to_summarize`); the record's `first_kept_index` is therefore just + // its length — see `compaction::CompactionRecord::first_kept_index`. + let first_kept_index = to_summarize.len(); + + match self.hook_decision( + CompactionReason::Threshold, + from_tokens, + &to_summarize, + &to_keep, + ) { + CompactionDecision::Decline => return Ok(()), + CompactionDecision::UseSummary(text) => { + let record = SummaryRecord { + summary: Message::system(text), + provenance: crate::summarization::CompressionProvenance { + source_ids: Vec::new(), + original_token_estimate: 0, + summary_token_estimate: 0, + reason: "before_compaction hook supplied the summary".to_string(), + }, + }; + let new_messages = splice_summary(to_keep, record.summary.clone()); + let to_tokens = total_message_tokens(&new_messages); + self.finish_compaction( + ctx, + record, + first_kept_index, + from_tokens, + to_tokens, + CompactionReason::Threshold, + ); + request.messages = new_messages; + ctx.emit(AgentEvent::Compressed { + from_tokens, + to_tokens, + }); + return Ok(()); + } + CompactionDecision::Proceed => {} + } + + let previous_summary = self + .last_summary + .lock() + .expect("last_summary mutex poisoned") + .clone(); + let record = match summarize_with_split( + self.summarizer.as_ref(), + &to_summarize, + self.max_turn_tokens.unwrap_or(u64::MAX), + previous_summary, + crate::token_estimation::estimate_message_tokens, + ) + .await + { Ok(record) => record, Err(err) => { // A summarizer failure hits precisely the longest, most valuable @@ -203,21 +292,250 @@ impl Middleware for ContextCom }; // `plan` returns `to_keep` as `[system prompts..., recent turns...]`. - // Insert the summary *after* the leading system prompts, not at index 0: - // a system prompt must stay first so its persistent instructions keep - // priority and the cacheable prefix is not churned. The summary of the - // elided older turns then sits between the system prompt and the kept - // recent turns, in chronological position. - let system_prefix = to_keep - .iter() - .take_while(|m| matches!(m, tinyinference_llm::message::Message::System(_))) - .count(); - let recent = to_keep.split_off(system_prefix); - let mut new_messages = Vec::with_capacity(to_keep.len() + recent.len() + 1); - new_messages.append(&mut to_keep); - new_messages.push(record.summary.clone()); - new_messages.extend(recent); + // `splice_summary` inserts the summary *after* the leading system + // prompts, not at index 0: a system prompt must stay first so its + // persistent instructions keep priority and the cacheable prefix is + // not churned. The summary of the elided older turns then sits + // between the system prompt and the kept recent turns, in + // chronological position. + let new_messages = splice_summary(to_keep, record.summary.clone()); + let to_tokens = total_message_tokens(&new_messages); + + self.finish_compaction( + ctx, + record, + first_kept_index, + from_tokens, + to_tokens, + CompactionReason::Threshold, + ); + request.messages = new_messages; + + ctx.emit(AgentEvent::Compressed { + from_tokens, + to_tokens, + }); + Ok(()) + } +} + +#[async_trait] +impl ModelMiddleware + for ContextCompressionMiddleware +{ + fn name(&self) -> &str { + self.label + } + + /// Implements pi's overflow → compact → retry recovery + /// (`docs/runtime-comparison/pi.md` §4.5): the wrapped model call runs + /// once; if it fails with an error + /// [`Self::overflow_classifier`][ContextCompressionMiddleware] classifies + /// as a provider context-window overflow, this compacts the transcript + /// once (recorded with [`CompactionReason::Overflow`]) and retries the + /// *same* turn exactly once more. A second overflow (or a decline from + /// the `before_compaction` hook) propagates the error instead of retrying + /// again, so a pathological transcript that cannot be shrunk under the + /// window cannot loop forever. + async fn wrap_model( + &self, + ctx: &mut RunContext, + state: &State, + request: ModelRequest, + next: ModelHandler<'_, State, Ctx>, + ) -> Result { + let first_error = match next.run(ctx, state, request.clone()).await { + Ok(outcome) => return Ok(outcome), + Err(error) => error, + }; + + let Some(_overflow) = self.overflow_classifier.classify(&first_error) else { + return Err(first_error); + }; + + let keep_recent_tokens = self.policy.trigger_budget(); + let Some(cut) = find_cut_point( + &request.messages, + keep_recent_tokens, + crate::token_estimation::estimate_message_tokens, + ) else { + // Nothing safe to cut (e.g. the whole transcript is already + // within budget, or is a single indivisible tool-call pair) — + // there is no compaction that could help, so surface the + // original provider error. + return Err(first_error); + }; + + let (system, non_system) = partition_messages_system(&request.messages); + let to_summarize = non_system[..cut.index].to_vec(); + let mut to_keep = system; + to_keep.extend(non_system[cut.index..].iter().cloned()); + let from_tokens = cut.tokens_before + cut.tokens_after; + + match self.hook_decision( + CompactionReason::Overflow, + from_tokens, + &to_summarize, + &to_keep, + ) { + CompactionDecision::Decline => return Err(first_error), + CompactionDecision::Proceed => {} + CompactionDecision::UseSummary(text) => { + let record = SummaryRecord { + summary: Message::system(text), + provenance: crate::summarization::CompressionProvenance { + source_ids: Vec::new(), + original_token_estimate: 0, + summary_token_estimate: 0, + reason: "before_compaction hook supplied the summary".to_string(), + }, + }; + let mut retried = request.clone(); + let new_messages = splice_summary(to_keep, record.summary.clone()); + let to_tokens = total_message_tokens(&new_messages); + self.finish_compaction( + ctx, + record, + cut.index, + from_tokens, + to_tokens, + CompactionReason::Overflow, + ); + retried.messages = new_messages; + return next.run(ctx, state, retried).await; + } + } + + let previous_summary = self + .last_summary + .lock() + .expect("last_summary mutex poisoned") + .clone(); + let record = match summarize_with_split( + self.summarizer.as_ref(), + &to_summarize, + self.max_turn_tokens.unwrap_or(u64::MAX), + previous_summary, + crate::token_estimation::estimate_message_tokens, + ) + .await + { + Ok(record) => record, + // Compaction itself failed: nothing changed, so surface the + // original overflow rather than a confusing summarizer error. + Err(_) => return Err(first_error), + }; + + let new_messages = splice_summary(to_keep, record.summary.clone()); let to_tokens = total_message_tokens(&new_messages); + self.finish_compaction( + ctx, + record, + cut.index, + from_tokens, + to_tokens, + CompactionReason::Overflow, + ); + + let mut retried = request; + retried.messages = new_messages; + // The retry is the *last* attempt: a second overflow propagates + // rather than looping — see this method's docs. + next.run(ctx, state, retried).await + } +} + +/// Inserts `summary` into `to_keep` right after any leading system messages, +/// so a system prompt stays first (preserving both its instruction priority +/// and the cacheable prefix) and the summary sits chronologically between it +/// and the kept recent turns. +fn splice_summary(mut to_keep: Vec, summary: Message) -> Vec { + let system_prefix = to_keep + .iter() + .take_while(|m| matches!(m, Message::System(_))) + .count(); + let recent = to_keep.split_off(system_prefix); + let mut new_messages = Vec::with_capacity(to_keep.len() + recent.len() + 1); + new_messages.append(&mut to_keep); + new_messages.push(summary); + new_messages.extend(recent); + new_messages +} + +/// [`crate::summarization::pairing`] partitions operate on non-system +/// slices; this mirrors that split for callers outside the `summarization` +/// module (`compaction::find_cut_point` already partitions internally, but +/// its caller here also needs the same partition to rebuild `to_keep`). +fn partition_messages_system(messages: &[Message]) -> (Vec, Vec) { + let system = messages + .iter() + .filter(|m| matches!(m, Message::System(_))) + .cloned() + .collect(); + let non_system = messages + .iter() + .filter(|m| !matches!(m, Message::System(_))) + .cloned() + .collect(); + (system, non_system) +} + +impl ContextCompressionMiddleware { + /// Consults the `before_compaction` hook, when one is installed; + /// defaults to [`CompactionDecision::Proceed`] otherwise. + fn hook_decision( + &self, + reason: CompactionReason, + tokens_before: u64, + to_summarize: &[Message], + to_keep: &[Message], + ) -> CompactionDecision { + match &self.before_compaction { + Some(hook) => hook(&CompactionContext { + reason, + tokens_before, + to_summarize_count: to_summarize.len(), + to_keep_count: to_keep.len(), + }), + None => CompactionDecision::Proceed, + } + } + + /// Finalizes a successful compaction: records `record` in the in-process + /// history, updates [`Self::last_summary`] for the next iterative + /// compaction, builds a [`CompactionRecord`], persists it through + /// [`RunContext::compaction_sink`] when attached, and emits + /// [`AgentEvent::Compacted`]. Does **not** emit `Compressed` — callers + /// that also want the legacy event emit it themselves. + fn finish_compaction( + &self, + ctx: &mut RunContext, + record: SummaryRecord, + first_kept_index: usize, + tokens_before: u64, + tokens_after: u64, + reason: CompactionReason, + ) { + *self + .last_summary + .lock() + .expect("last_summary mutex poisoned") = Some(record.summary.text()); + + let compaction_record = CompactionRecord { + summary: record.summary.text(), + first_kept_index, + tokens_before, + tokens_after, + usage: None, + details: serde_json::json!({ "source_ids": record.provenance.source_ids }), + reason, + }; + + if let Some(sink) = &ctx.compaction_sink + && let Err(err) = sink.persist(&compaction_record) + { + tracing::debug!("[context_compression] compaction sink persist failed: {err}"); + } { let mut records = self.records.lock().expect("records mutex poisoned"); @@ -228,13 +546,12 @@ impl Middleware for ContextCom records.push_back(record); } } - request.messages = new_messages; - ctx.emit(AgentEvent::Compressed { - from_tokens, - to_tokens, + ctx.emit(AgentEvent::Compacted { + reason, + tokens_before, + tokens_after, }); - Ok(()) } } @@ -369,7 +686,7 @@ impl Middleware for Microcompa // of, a signature, a diff — so leave it intact and reclaim // tokens elsewhere. if t.trusted_verbatim { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::middleware", tool_call_id = %t.tool_call_id, "[microcompact] skipping a trusted_verbatim tool result" @@ -462,7 +779,7 @@ impl Middleware for PromptCach && prev_run == &run_id && !prev.is_prefix_stable_against(&layout) { - tinyagents_tracing::debug!( + tracing::debug!( "[cache] prompt_cache_guard: prefix invalidated run={run_id} \ before={} after={}", prev.fingerprint(), diff --git a/crates/tinyagents-harness/src/middleware/library/mod.rs b/crates/tinyagents-harness/src/middleware/library/mod.rs index cf6a5d5f..bc8524b2 100644 --- a/crates/tinyagents-harness/src/middleware/library/mod.rs +++ b/crates/tinyagents-harness/src/middleware/library/mod.rs @@ -36,7 +36,7 @@ use std::time::{Duration, Instant}; use async_trait::async_trait; -use crate::context::{RunConfig, RunContext}; +use crate::context::{MiddlewareControl, RunConfig, RunContext}; use crate::error::{Result, TinyAgentsError}; use crate::events::AgentEvent; use crate::ids::CallId; diff --git a/crates/tinyagents-harness/src/middleware/library/observe.rs b/crates/tinyagents-harness/src/middleware/library/observe.rs index 52ed83ab..55f8c49b 100644 --- a/crates/tinyagents-harness/src/middleware/library/observe.rs +++ b/crates/tinyagents-harness/src/middleware/library/observe.rs @@ -265,6 +265,9 @@ impl Middleware for RedactionM } } ToolContent::Json { data } => hits += self.redact_value(data), + // Image/File blocks carry no free text to redact; the media + // type/name fields are structural, not user data. + ToolContent::Image { .. } | ToolContent::File { .. } => {} } } if let Some(markdown) = &mut result.markdown_formatted { diff --git a/crates/tinyagents-harness/src/middleware/library/resilience.rs b/crates/tinyagents-harness/src/middleware/library/resilience.rs index 11675646..a292c52d 100644 --- a/crates/tinyagents-harness/src/middleware/library/resilience.rs +++ b/crates/tinyagents-harness/src/middleware/library/resilience.rs @@ -37,6 +37,10 @@ impl ModelMiddleware for Retry self.label } + fn overrides_retry(&self) -> bool { + true + } + async fn wrap_model( &self, ctx: &mut RunContext, @@ -59,7 +63,16 @@ impl ModelMiddleware for Retry // one step too high. let backoff_attempt = attempt; attempt += 1; - let call_id = CallId::new(format!("{}-model", ctx.run_id())); + // Prefer the loop's own call id (mirrored onto the + // context for exactly this purpose, see I-7) so + // `RetryScheduled` events correlate with the same + // call id `ModelStarted`/`ModelCompleted` use. Falls + // back to a run-scoped id for a caller that invokes + // this middleware outside the agent loop. + let call_id = ctx + .active_model_call + .clone() + .unwrap_or_else(|| CallId::new(format!("{}-model", ctx.run_id()))); ctx.emit(AgentEvent::RetryScheduled { call_id, attempt }); // Sleep for the backoff only when the policy opts in // (`with_backoff_sleep`); a no-op otherwise. diff --git a/crates/tinyagents-harness/src/middleware/library/test.rs b/crates/tinyagents-harness/src/middleware/library/test.rs index c3bef4d5..d1023dd6 100644 --- a/crates/tinyagents-harness/src/middleware/library/test.rs +++ b/crates/tinyagents-harness/src/middleware/library/test.rs @@ -7,7 +7,7 @@ use std::time::{Duration, Instant}; use serde_json::json; use super::*; -use crate::context::{RunConfig, RunContext}; +use crate::context::{MiddlewareControl, RunConfig, RunContext}; use crate::error::{Result, TinyAgentsError}; use crate::events::{AgentEvent, EventRecord, RecordingListener}; use crate::middleware::{BoxModelFuture, MiddlewareStack, ModelBaseCall, ToolInvocationIdentity}; @@ -116,6 +116,51 @@ async fn retry_middleware_retries_then_succeeds() { assert_eq!(scheduled, 2); } +#[tokio::test] +async fn retry_middleware_correlates_retry_scheduled_with_the_loops_call_id() { + // R-3: the loop's `invoke_model_resolving` mirrors its own call id onto + // `ctx.active_model_call` specifically so a retrying middleware's + // `RetryScheduled` events carry the same id as that attempt's + // `ModelStarted`/`ModelCompleted` pair, letting a consumer join retries to + // the call they belong to instead of only to the run. + let (mut ctx, recorder) = ctx_with_recorder(); + let call_id = crate::ids::CallId::new("run-model-3"); + ctx.active_model_call = Some(call_id.clone()); + + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push_model_middleware(Arc::new(RetryMiddleware::new( + RetryPolicy::default().with_max_attempts(3), + ))); + + let base = FakeModelBase::new(|n, _req| { + if n < 2 { + Err(TinyAgentsError::Model("transient".to_string())) + } else { + Ok(ok_response()) + } + }); + + stack + .run_wrapped_model(&mut ctx, &(), ModelRequest::default(), &base) + .await + .expect("retry should eventually succeed"); + + let scheduled: Vec<_> = events(&recorder) + .into_iter() + .filter_map(|e| match e { + AgentEvent::RetryScheduled { call_id, attempt } => Some((call_id, attempt)), + _ => None, + }) + .collect(); + assert_eq!(scheduled.len(), 2); + for (event_call_id, _attempt) in &scheduled { + assert_eq!( + *event_call_id, call_id, + "RetryScheduled must carry the same call id as the attempt it retries" + ); + } +} + #[tokio::test(start_paused = true)] async fn retry_middleware_sleeps_the_documented_backoff_schedule() { // Regression test: the middleware used to compute the backoff from the @@ -603,12 +648,17 @@ async fn budget_warns_then_blocks_on_token_exhaustion() { .any(|e| matches!(e, AgentEvent::BudgetExceeded { blocked: false, .. })) ); - // Now preflight fails closed. - let err = stack + // Now preflight fails closed — gracefully (A1): the control-outcome hook + // the stack actually drives (`before_model_control`) requests + // `JumpTo(End)` instead of erroring the whole run out. + stack .run_before_model(&mut ctx, &(), &mut req) .await - .expect_err("budget exhausted should block"); - assert!(matches!(err, TinyAgentsError::LimitExceeded(_))); + .expect("an exhausted budget stops the run gracefully, not with an error"); + assert!(matches!( + ctx.take_control(), + Some(MiddlewareControl::JumpTo(crate::context::LoopTarget::End)) + )); } #[tokio::test] @@ -647,11 +697,14 @@ async fn budget_prices_usage_and_enforces_cost() { ); let mut req = ModelRequest::new(vec![Message::user("go")]); - let err = stack + stack .run_before_model(&mut ctx, &(), &mut req) .await - .expect_err("cost budget exhausted should block"); - assert!(matches!(err, TinyAgentsError::LimitExceeded(_))); + .expect("a cost budget exhausted stops the run gracefully, not with an error"); + assert!(matches!( + ctx.take_control(), + Some(MiddlewareControl::JumpTo(crate::context::LoopTarget::End)) + )); } #[tokio::test] @@ -679,11 +732,14 @@ async fn budget_enforces_cached_input_token_limit() { .unwrap(); let mut req = ModelRequest::new(vec![Message::user("next")]); - let err = stack + stack .run_before_model(&mut ctx, &(), &mut req) .await - .expect_err("cached input budget exhausted should block"); - assert!(matches!(err, TinyAgentsError::LimitExceeded(_))); + .expect("a cached-input budget exhausted stops the run gracefully, not with an error"); + assert!(matches!( + ctx.take_control(), + Some(MiddlewareControl::JumpTo(crate::context::LoopTarget::End)) + )); } #[tokio::test] @@ -1236,11 +1292,19 @@ async fn human_approval_interrupts_without_callback() { stack.push(Arc::new(HumanApprovalMiddleware::new(["wire_transfer"]))); let mut call = tool_call("wire_transfer"); - let err = stack + // A1: the flagged call now requests `MiddlewareControl::Interrupt` + // through the control-outcome hook the stack actually drives + // (`before_tool_control`), rather than erroring `run_before_tool` out + // directly — the agent loop honors the queued control at its next safe + // checkpoint with the same `TinyAgentsError::Interrupted`. + stack .run_before_tool(&mut ctx, &(), &mut call) .await - .expect_err("flagged tool requires approval"); - assert!(matches!(err, TinyAgentsError::Interrupted { .. })); + .expect("the hook itself succeeds; the interrupt is queued as control"); + assert!(matches!( + ctx.take_control(), + Some(MiddlewareControl::Interrupt { .. }) + )); } #[tokio::test] @@ -1259,11 +1323,14 @@ async fn human_approval_consults_callback() { .expect("callback approves wire_transfer"); let mut rejected = tool_call("delete"); - let err = stack + stack .run_before_tool(&mut ctx, &(), &mut rejected) .await - .expect_err("callback rejects delete"); - assert!(matches!(err, TinyAgentsError::Interrupted { .. })); + .expect("the hook itself succeeds; the rejection is queued as control"); + assert!(matches!( + ctx.take_control(), + Some(MiddlewareControl::Interrupt { .. }) + )); } // ── StructuredOutputValidatorMiddleware ───────────────────────────────────── diff --git a/crates/tinyagents-harness/src/middleware/library/tool_policy.rs b/crates/tinyagents-harness/src/middleware/library/tool_policy.rs index 572ae94d..f8efce5b 100644 --- a/crates/tinyagents-harness/src/middleware/library/tool_policy.rs +++ b/crates/tinyagents-harness/src/middleware/library/tool_policy.rs @@ -20,8 +20,13 @@ impl ToolAllowlistMiddleware { } /// Returns `true` if `name` is on the allowlist. + /// + /// Delegates to [`crate::tool::toolset::tool_name_allowed`] — the exact + /// membership test [`crate::tool::toolset::FilteredToolSet::allowing`] + /// uses — so this middleware and its `ToolSet` counterpart cannot drift + /// (`docs/runtime-comparison/pydantic-ai.md` §4). pub fn allows(&self, name: &str) -> bool { - self.allowed.contains(name) + crate::tool::toolset::tool_name_allowed(&self.allowed, name) } } @@ -37,7 +42,7 @@ impl Middleware for ToolAllowl _state: &State, call: &mut ToolCall, ) -> Result<()> { - if !self.allowed.contains(&call.name) { + if !self.allows(&call.name) { return Err(TinyAgentsError::Validation(format!( "tool `{}` is not on the allowlist", call.name @@ -349,7 +354,10 @@ impl Middleware _state: &State, request: &mut ModelRequest, ) -> Result<()> { - request.tools.retain(|schema| (self.predicate)(schema)); + // Delegates to `PreparedToolSet`'s retain helper — see + // `crate::tool::toolset::retain_matching_schemas`'s doc comment for + // why this is shared rather than a second `retain` implementation. + crate::tool::toolset::retain_matching_schemas(&mut request.tools, self.predicate.as_ref()); Ok(()) } } @@ -461,6 +469,15 @@ impl Middleware if !excluded.is_empty() { ctx.emit(AgentEvent::ToolsFiltered { by: self.label.to_string(), + explanations: excluded + .iter() + .map(|name| { + ( + name.clone(), + crate::tool::ToolExposureExplanation::FilteredOut, + ) + }) + .collect(), excluded, remaining: request.tools.len(), }); @@ -479,6 +496,7 @@ impl HumanApprovalMiddleware { label: "human_approval", flagged: flagged.into_iter().map(Into::into).collect(), approve: None, + outcome: None, } } @@ -488,6 +506,50 @@ impl HumanApprovalMiddleware { self.approve = Some(approve); self } + + /// Attaches a callback that decides [`ApprovalOutcome::Allow`], + /// [`ApprovalOutcome::Deny`], or [`ApprovalOutcome::Defer`] for each + /// flagged call (A2). Takes precedence over [`Self::with_approval`]. + pub fn with_approval_outcome(mut self, outcome: ApprovalOutcomeFn) -> Self { + self.outcome = Some(outcome); + self + } + + /// Resolves one flagged call to a control outcome, or a signal error the + /// tool-admission path turns into a deferral / a denial answer. + /// + /// A call the resume path already approved is always allowed, so the + /// same gate cannot defer it a second time. + fn decide(&self, ctx: &RunContext, call: &ToolCall) -> Result { + if !self.flagged.contains(&call.name) || ctx.is_call_approved(&call.id) { + return Ok(MiddlewareControl::Continue); + } + if let Some(outcome) = &self.outcome { + return match outcome(call) { + ApprovalOutcome::Allow => Ok(MiddlewareControl::Continue), + ApprovalOutcome::Deny(message) => Err(TinyAgentsError::ToolFailed(message)), + ApprovalOutcome::Defer => Err(TinyAgentsError::ApprovalRequired { + metadata: serde_json::json!({ + "gate": self.label, + "tool": call.name, + }), + }), + }; + } + let approved = self + .approve + .as_ref() + .map(|approve| approve(call)) + .unwrap_or(false); + if approved { + Ok(MiddlewareControl::Continue) + } else { + Ok(MiddlewareControl::Interrupt { + node: "tool".to_string(), + message: format!("tool `{}` requires human approval", call.name), + }) + } + } } #[async_trait] @@ -498,23 +560,38 @@ impl Middleware for HumanAppro async fn before_tool( &self, - _ctx: &mut RunContext, + ctx: &mut RunContext, _state: &State, call: &mut ToolCall, ) -> Result<()> { - if self.flagged.contains(&call.name) { - let approved = self - .approve - .as_ref() - .map(|approve| approve(call)) - .unwrap_or(false); - if !approved { - return Err(TinyAgentsError::Interrupted { - node: "tool".to_string(), - message: format!("tool `{}` requires human approval", call.name), - }); + match self.decide(ctx, call)? { + MiddlewareControl::Interrupt { node, message } => { + Err(TinyAgentsError::Interrupted { node, message }) } + _ => Ok(()), } - Ok(()) + } + + /// Control-outcome override (A1): a flagged, unapproved call requests + /// [`MiddlewareControl::Interrupt`] instead of erroring the run out + /// directly. The agent loop drains the request at its next safe + /// checkpoint — the same place any other interrupt is honored — and + /// surfaces the identical [`TinyAgentsError::Interrupted`], so callers + /// driving the harness through the ordinary loop see no behavior change; + /// what changes is that the interrupt is now expressed in the shared + /// control vocabulary a durable HITL host can also inspect via + /// [`RunContext::take_control`][crate::context::RunContext::take_control] + /// before it is drained, rather than only as a thrown error. + /// + /// With an [`ApprovalOutcomeFn`] installed (A2), `Deny` and `Defer` are + /// returned as the `ToolFailed` / `ApprovalRequired` signals that tool + /// admission answers or defers without failing the run. + async fn before_tool_control( + &self, + ctx: &mut RunContext, + _state: &State, + call: &mut ToolCall, + ) -> Result { + self.decide(ctx, call) } } diff --git a/crates/tinyagents-harness/src/middleware/library/types.rs b/crates/tinyagents-harness/src/middleware/library/types.rs index eba03668..7ea03204 100644 --- a/crates/tinyagents-harness/src/middleware/library/types.rs +++ b/crates/tinyagents-harness/src/middleware/library/types.rs @@ -41,14 +41,36 @@ use tinytools::{ToolPolicy, ToolSideEffects}; /// the configured [`RetryPolicy`][crate::retry::RetryPolicy] still /// permits another attempt, retries. Each scheduled retry emits an /// [`AgentEvent::RetryScheduled`][crate::events::AgentEvent::RetryScheduled] -/// with a [`CallId`][crate::ids::CallId] derived from the run id. +/// with the same [`CallId`][crate::ids::CallId] the agent loop is using for +/// the in-flight call (mirrored onto +/// [`RunContext::active_model_call`][crate::context::RunContext::active_model_call]), +/// falling back to a run-scoped id only when this middleware runs outside the +/// agent loop. +/// +/// # An alternative to `RunPolicy::retry`, not a companion +/// +/// This middleware and the loop's own [`RunPolicy::retry`][crate::runtime::RunPolicy::retry] +/// are two implementations of the same idea. Registering both does not +/// compose them: [`crate::middleware::MiddlewareStack::has_retry_override`] +/// tells the loop's base call to skip its own retry loop whenever any +/// `ModelMiddleware` reports [`ModelMiddleware::overrides_retry`][crate::middleware::ModelMiddleware::overrides_retry] +/// (this middleware always does), so only this middleware's `RetryPolicy` +/// governs the attempt count — `RunPolicy::retry` is ignored for the base +/// call while it is registered. Without that guard the two layers would +/// multiply attempts (`mw.max_attempts x policy.retry.max_attempts x +/// |fallback|` provider calls for one logical failure); see I-7. Prefer +/// `RunPolicy::retry` for the common case (it also drives the fallback +/// chain) and reach for this middleware only when retry needs to run at a +/// specific point in the wrap onion (e.g. after a guardrail middleware has +/// already inspected the request). Full unification into one retry engine is +/// tracked as a later phase. /// /// # Sleeping /// -/// Like the agent loop's own retry path, this middleware *computes* the backoff -/// from the policy but does **not** sleep, keeping the loop fast and tests -/// deterministic. A production integration may sleep for -/// [`RetryMiddleware::backoff_for_attempt`] before each retry. +/// This middleware sleeps for the policy's computed backoff between attempts +/// only when the policy opts in via +/// [`RetryPolicy::with_backoff_sleep`][crate::retry::RetryPolicy::with_backoff_sleep]; +/// otherwise it retries back-to-back, keeping tests fast and deterministic. /// /// # Failure mode /// @@ -285,6 +307,17 @@ pub struct BudgetMiddleware { /// /// Rejections at `before_tool` surface as /// [`TinyAgentsError::Validation`][crate::error::TinyAgentsError::Validation]. +/// +/// # Relationship to `crate::tool::toolset` +/// +/// This middleware's policy classification (side-effect/background-safe/ +/// approval enforcement from the vendor `tinytools` declaration) is a +/// different axis from [`crate::tool::toolset::FilteredToolSet`] (an +/// arbitrary per-tool predicate) and +/// [`crate::tool::toolset::ApprovalRequiredToolSet`] (which only *sets* the +/// approval flag this middleware enforces) — kept as its own implementation +/// rather than rebased onto either, since neither adaptor reads +/// [`ToolPolicy`] as a whole. pub struct ToolPolicyMiddleware { pub(crate) label: &'static str, pub(crate) policies: std::collections::HashMap, @@ -361,6 +394,17 @@ pub type ContextualToolPredicate = /// or from explicit allow/deny lists with /// [`from_lists`](Self::from_lists) (deny wins; when an allow-list is present a /// tool must appear in it — fail-closed for unknown tools). +/// +/// # Relationship to `crate::tool::toolset` +/// +/// [`DynamicToolSelectionMiddleware`] and +/// [`crate::tool::toolset::PreparedToolSet`] both operate on a bare +/// [`ToolSchema`] predicate; this middleware additionally reads +/// [`ToolSelectionContext`] (depth, tags, the requested model), which a +/// `ToolSet::tools`'s own `ctx: &RunContext` argument can already carry +/// through `Ctx` — kept as its own predicate type rather than folded into +/// [`crate::tool::toolset::PreparedToolSet::filtering`] to avoid coupling +/// every `Ctx` to this specific context shape. pub struct ContextualToolSelectionMiddleware { pub(crate) label: &'static str, pub(crate) predicate: ContextualToolPredicate, @@ -374,6 +418,28 @@ pub struct ContextualToolSelectionMiddleware { /// middleware then raises an interrupt). pub type ApprovalFn = Arc bool + Send + Sync>; +/// What an approval callback decided about one flagged [`ToolCall`] (A2). +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ApprovalOutcome { + /// Run the call now. + Allow, + /// Do not run it; the model sees the message as a tool-error result and + /// the run continues (no interrupt, no deferral). + Deny(String), + /// Hand the call back to the host: the loop finishes the batch's other + /// calls and exits with `AgentRun::deferred` listing this call under + /// `approvals` (or resolves it through a registered + /// [`DeferredToolHandler`][crate::tool::DeferredToolHandler]). On resume + /// the middleware sees the approval through + /// [`RunContext::is_call_approved`][crate::context::RunContext::is_call_approved] + /// and lets the call through. + Defer, +} + +/// A richer approval callback returning an [`ApprovalOutcome`] instead of a +/// bare `bool`; see [`HumanApprovalMiddleware::with_approval_outcome`]. +pub type ApprovalOutcomeFn = Arc ApprovalOutcome + Send + Sync>; + /// Lifecycle middleware implementing a simple human-in-the-loop gate for /// sensitive tools. /// @@ -384,6 +450,11 @@ pub type ApprovalFn = Arc bool + Send + Sync>; /// [`TinyAgentsError::Interrupted`][crate::error::TinyAgentsError::Interrupted] /// (node `"tool"`) so the run pauses for human input. /// +/// An [`ApprovalOutcomeFn`] (see [`Self::with_approval_outcome`]) replaces +/// the bare `bool` with [`ApprovalOutcome::{Allow, Deny, Defer}`]: `Deny` +/// answers the model with a tool-error result instead of interrupting, and +/// `Defer` turns the call into a resumable deferred request (A2). +/// /// # HITL hookup /// /// This is the harness-native signal; the full graph interrupt/resume path is a @@ -394,6 +465,8 @@ pub struct HumanApprovalMiddleware { pub(crate) label: &'static str, pub(crate) flagged: std::collections::HashSet, pub(crate) approve: Option, + /// Takes precedence over `approve` when set (A2). + pub(crate) outcome: Option, } // ── StructuredOutputValidatorMiddleware ─────────────────────────────────────── diff --git a/crates/tinyagents-harness/src/middleware/mod.rs b/crates/tinyagents-harness/src/middleware/mod.rs index 1a30d6f8..348a0428 100644 --- a/crates/tinyagents-harness/src/middleware/mod.rs +++ b/crates/tinyagents-harness/src/middleware/mod.rs @@ -36,39 +36,63 @@ pub use library::*; use std::sync::Arc; -use crate::context::RunContext; +use crate::context::{MiddlewareControl, RunContext}; use crate::error::{Result, TinyAgentsError}; use crate::events::AgentEvent; use tinyinference_llm::model::{ModelDelta, ModelRequest, ModelResponse}; use tinyinference_llm::tool::{ToolCall, ToolDelta}; use tinytools::ToolResult; -/// Runs one per-middleware lifecycle hook across the whole stack, bracketing -/// each call with `MiddlewareStarted`/`MiddlewareCompleted` events and fanning -/// `on_error` out to every middleware on the first failure (so the originating -/// error is never masked). +/// Runs one per-middleware **control-outcome** hook across the whole stack, +/// bracketing each *actually invoked* call with +/// `MiddlewareStarted`/`MiddlewareCompleted` events, fanning `on_error` out to +/// every middleware on the first failure, and resolving the phase's +/// [`MiddlewareControl`] per the precedence rule documented on +/// [`Middleware::is_observer`]: the first non-[`MiddlewareControl::Continue`] +/// outcome wins; every hook after it is skipped unless +/// [`Middleware::is_observer`] returns `true` for it, in which case it still +/// runs (for observation) but its own control outcome is discarded. The +/// winning control (if any) is installed via +/// [`RunContext::request_control`], exactly as if a hook had called it +/// directly — this macro is the single place that bridges "hook returned a +/// control" and "hook called `request_control`" into one mechanism. /// -/// This is factored as a macro rather than an async helper because each hook -/// takes different arguments and borrows `ctx` mutably across its `await`, which -/// a closure-based helper cannot express without heap-boxing every call. -/// -/// Crucially, `MiddlewareCompleted` is emitted on *both* the success and error -/// paths: a hook that returns `Err` can no longer leave a dangling -/// `MiddlewareStarted` with no matching `Completed` in the event stream. `$iter` -/// selects registration order (`.iter()`) or reverse order (`.iter().rev()`); -/// `$call` is the (un-awaited) hook invocation on `$mw`. +/// Factored as a macro (not an async helper) for the same reason as before +/// control outcomes existed: each hook takes different arguments and borrows +/// `ctx` mutably across its `await`, which a closure-based helper cannot +/// express without heap-boxing every call. `$iter` selects registration order +/// (`.iter()`) or reverse order (`.iter().rev()`); `$call` is the (un-awaited) +/// `_control` hook invocation on `$mw`. macro_rules! run_stack_hook { ($self:ident, $ctx:ident, $iter:expr, |$mw:ident| $call:expr) => {{ + let mut winning: Option = None; for $mw in $iter { + if winning.is_some() && !$mw.is_observer() { + continue; + } let name = $mw.name().to_string(); $ctx.emit(AgentEvent::MiddlewareStarted { name: name.clone() }); let result = $call.await; - $ctx.emit(AgentEvent::MiddlewareCompleted { name }); - if let Err(e) = result { - $self.fan_out_on_error($ctx, &e).await; - return Err(e); + $ctx.emit(AgentEvent::MiddlewareCompleted { name: name.clone() }); + match result { + Ok(control) => { + if winning.is_none() && !matches!(control, MiddlewareControl::Continue) { + winning = Some(control); + } + } + Err(e) => { + $ctx.emit(AgentEvent::MiddlewareFailed { + name, + error: e.to_string(), + }); + $self.fan_out_on_error($ctx, &e).await; + return Err(e); + } } } + if let Some(control) = winning { + $ctx.request_control(control); + } Ok(()) }}; } @@ -85,6 +109,23 @@ impl AgentRun { pub fn text(&self) -> Option { self.final_response.as_ref().map(|r| r.text()) } + + /// Deserializes [`Self::structured`] into `T`, when the run produced a + /// structured output. + /// + /// A typed convenience over `run.structured`, mirroring Pydantic AI's + /// `result.output` (A3). Returns + /// [`TinyAgentsError::StructuredOutput`][crate::error::TinyAgentsError::StructuredOutput] + /// when the run produced no structured value, or when the value does not + /// deserialize into `T`. + pub fn structured_as(&self) -> Result { + let value = self.structured.clone().ok_or_else(|| { + TinyAgentsError::StructuredOutput("run produced no structured output".to_string()) + })?; + serde_json::from_value(value).map_err(|error| { + TinyAgentsError::StructuredOutput(format!("deserialization failed: {error}")) + }) + } } // ── MiddlewareStack ─────────────────────────────────────────────────────────── @@ -137,6 +178,31 @@ impl MiddlewareStack { self.model_middlewares.len() } + /// Returns `true` when a registered [`ModelMiddleware`] already retries + /// the model call itself (see [`ModelMiddleware::overrides_retry`]). + /// + /// The agent loop's base call uses this to skip its own + /// [`crate::runtime::RunPolicy::retry`] loop, so `RetryMiddleware` and the + /// loop's built-in retry do not multiply attempts together (I-7). + pub fn has_retry_override(&self) -> bool { + self.model_middlewares.iter().any(|mw| mw.overrides_retry()) + } + + /// Returns `true` when any registered lifecycle [`Middleware`] asks to + /// stop after the turn currently completing (see + /// [`Middleware::should_stop_after_turn`]). + /// + /// Called by the agent loop at the turn boundary — after tool execution, + /// before the loop would otherwise continue — so an aggregate stop + /// condition (a tally across the whole turn's tool results, not any + /// single call) can end the run as cleanly as + /// [`crate::context::MiddlewareControl::JumpTo`]`(`[`crate::context::LoopTarget::End`]`)`. + pub fn any_should_stop_after_turn(&self, ctx: &RunContext, run: &AgentRun) -> bool { + self.middlewares + .iter() + .any(|mw| mw.should_stop_after_turn(ctx, run)) + } + /// Returns the number of registered around-agent middleware layers. pub fn agent_middleware_len(&self) -> usize { self.agent_middlewares.len() @@ -175,7 +241,7 @@ impl MiddlewareStack { /// order. pub async fn run_before_agent(&self, ctx: &mut RunContext, state: &State) -> Result<()> { run_stack_hook!(self, ctx, self.middlewares.iter(), |mw| mw - .before_agent(ctx, state)) + .before_agent_control(ctx, state)) } /// Runs every middleware's [`Middleware::after_agent`] in reverse @@ -187,7 +253,7 @@ impl MiddlewareStack { run: &mut AgentRun, ) -> Result<()> { run_stack_hook!(self, ctx, self.middlewares.iter().rev(), |mw| mw - .after_agent(ctx, state, run)) + .after_agent_control(ctx, state, run)) } /// Runs every middleware's [`Middleware::before_model`] in registration @@ -199,7 +265,7 @@ impl MiddlewareStack { request: &mut ModelRequest, ) -> Result<()> { run_stack_hook!(self, ctx, self.middlewares.iter(), |mw| mw - .before_model(ctx, state, request)) + .before_model_control(ctx, state, request)) } /// Runs every middleware's [`Middleware::on_model_delta`] in registration @@ -237,31 +303,91 @@ impl MiddlewareStack { response: &mut ModelResponse, ) -> Result<()> { run_stack_hook!(self, ctx, self.middlewares.iter().rev(), |mw| mw - .after_model(ctx, state, response)) + .after_model_control(ctx, state, response)) } /// Runs every middleware's [`Middleware::before_tool`] in registration /// order, threading the mutable tool call through each. + /// + /// Unlike the other stack runners this one recognises the per-call + /// signals of A2/A3 — `ApprovalRequired`, `CallDeferred`, `ToolFailed`, + /// `ModelRetry` — as *decisions about the call* rather than hook + /// failures: they propagate to admission (which defers or answers the + /// call) without a `MiddlewareFailed` event or an `on_error` fan-out. + /// An `ApprovalRequired` for a call the resume path already approved + /// ([`RunContext::is_call_approved`]) is treated as `Continue`, so a + /// gate that cannot see the approval does not re-defer the call. pub async fn run_before_tool( &self, ctx: &mut RunContext, state: &State, call: &mut ToolCall, ) -> Result<()> { - run_stack_hook!(self, ctx, self.middlewares.iter(), |mw| mw - .before_tool(ctx, state, call)) + let mut winning: Option = None; + for mw in self.middlewares.iter() { + if winning.is_some() && !mw.is_observer() { + continue; + } + let name = mw.name().to_string(); + ctx.emit(AgentEvent::MiddlewareStarted { name: name.clone() }); + let result = mw.before_tool_control(ctx, state, call).await; + ctx.emit(AgentEvent::MiddlewareCompleted { name: name.clone() }); + match result { + Ok(control) => { + if winning.is_none() && !matches!(control, MiddlewareControl::Continue) { + winning = Some(control); + } + } + Err(TinyAgentsError::ApprovalRequired { .. }) if ctx.is_call_approved(&call.id) => { + } + Err( + signal @ (TinyAgentsError::ApprovalRequired { .. } + | TinyAgentsError::CallDeferred { .. } + | TinyAgentsError::ToolFailed(_) + | TinyAgentsError::ModelRetry(_)), + ) => return Err(signal), + Err(e) => { + ctx.emit(AgentEvent::MiddlewareFailed { + name, + error: e.to_string(), + }); + self.fan_out_on_error(ctx, &e).await; + return Err(e); + } + } + } + if let Some(control) = winning { + ctx.request_control(control); + } + Ok(()) } /// Runs every middleware's [`Middleware::on_tool_delta`] in registration /// order for one streamed tool-progress delta. + /// + /// Like [`Self::run_on_model_delta`], and for the same reason (M-12): + /// this is **not** bracketed by `MiddlewareStarted`/`MiddlewareCompleted` + /// events. It used to be the one delta hook still routed through + /// `run_stack_hook!`, so a stack of `N` middlewares produced `2*N` + /// bookkeeping events per streamed tool-progress delta — noise a + /// `ModelCompleted`-based exporter had to filter, for a hook that (unlike + /// `before_tool`/`after_tool`) can fire many times per call. Both delta + /// hooks now agree: bracket every non-delta hook, skip both delta hooks. + /// A caller that needs to observe delta-level middleware activity should + /// instrument the hook implementation itself. pub async fn run_on_tool_delta( &self, ctx: &mut RunContext, state: &State, delta: &mut ToolDelta, ) -> Result<()> { - run_stack_hook!(self, ctx, self.middlewares.iter(), |mw| mw - .on_tool_delta(ctx, state, delta)) + for mw in self.middlewares.iter() { + if let Err(e) = mw.on_tool_delta(ctx, state, delta).await { + self.fan_out_on_error(ctx, &e).await; + return Err(e); + } + } + Ok(()) } /// Runs every middleware's [`Middleware::after_tool`] in reverse @@ -274,7 +400,7 @@ impl MiddlewareStack { result: &mut ToolResult, ) -> Result<()> { run_stack_hook!(self, ctx, self.middlewares.iter().rev(), |mw| mw - .after_tool(ctx, state, invocation, result)) + .after_tool_control(ctx, state, invocation, result)) } /// Runs every middleware's [`Middleware::on_error`] in registration order, diff --git a/crates/tinyagents-harness/src/middleware/test.rs b/crates/tinyagents-harness/src/middleware/test.rs index ad48575c..d8471efc 100644 --- a/crates/tinyagents-harness/src/middleware/test.rs +++ b/crates/tinyagents-harness/src/middleware/test.rs @@ -8,7 +8,10 @@ use super::*; use crate::context::{RunConfig, RunContext}; use crate::error::{Result, TinyAgentsError}; use crate::events::{AgentEvent, RecordingListener}; -use crate::summarization::{SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy}; +use crate::summarization::{ + CompactionContext, CompactionDecision, CompactionRecord, CompactionSink, SummarizationPolicy, + Summarizer, SummaryRecord, TrimStrategy, +}; use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message, UserMessage}; use tinyinference_llm::model::{ModelRequest, ModelResponse, PromptSegment, SegmentRole}; use tinyinference_llm::tool::ToolCall; @@ -34,6 +37,7 @@ fn response_with_usage(usage: Usage) -> ModelResponse { content: vec![ContentBlock::Text("ok".to_string())], tool_calls: Vec::new(), usage: None, + origin: None, }, usage: Some(usage), finish_reason: None, @@ -240,6 +244,36 @@ async fn failing_hook_still_emits_balanced_completed_event() { ); } +/// I-3 regression: `run_stack_hook!` must emit `AgentEvent::MiddlewareFailed` +/// for a hook that returns `Err`, not just fan `on_error` out privately. The +/// variant existed but nothing in the stack emitted it before this fix. +#[tokio::test] +async fn failing_hook_emits_middleware_failed() { + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(FailingMiddleware)); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + + let mut request = ModelRequest::default(); + let _ = stack.run_before_model(&mut c, &(), &mut request).await; + + let failed: Vec = recorder + .events() + .into_iter() + .map(|r| r.event) + .filter(|e| matches!(e, AgentEvent::MiddlewareFailed { .. })) + .collect(); + assert_eq!( + failed, + vec![AgentEvent::MiddlewareFailed { + name: "failing".to_string(), + error: TinyAgentsError::Middleware("boom".to_string()).to_string(), + }], + ); +} + #[tokio::test] async fn on_model_delta_hook_emits_no_bracketing_events() { // The per-delta hook runs on the streaming hot path, so it must NOT emit @@ -280,6 +314,47 @@ async fn on_model_delta_hook_emits_no_bracketing_events() { ); } +#[tokio::test] +async fn on_tool_delta_hook_emits_no_bracketing_events() { + // M-12 regression: `run_on_tool_delta` was the one delta hook still + // routed through `run_stack_hook!`, so it emitted + // `MiddlewareStarted`/`MiddlewareCompleted` on every streamed + // tool-progress delta while `run_on_model_delta` (the sibling hook, same + // hot-path shape) did not. The two delta hooks must agree. + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(LoggingMiddleware::new())); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + + let mut delta = tinyinference_llm::tool::ToolDelta { + call_id: "call-1".to_string(), + content: "partial args".to_string(), + tool_name: Some("search".to_string()), + ..Default::default() + }; + stack + .run_on_tool_delta(&mut c, &(), &mut delta) + .await + .unwrap(); + + let bracketing = recorder + .events() + .into_iter() + .filter(|r| { + matches!( + r.event, + AgentEvent::MiddlewareStarted { .. } | AgentEvent::MiddlewareCompleted { .. } + ) + }) + .count(); + assert_eq!( + bracketing, 0, + "the tool-delta hook must not bracket middleware with events" + ); +} + #[tokio::test] async fn message_trim_middleware_shrinks_request() { let mw = MessageTrimMiddleware::new(TrimStrategy::KeepLast(1)); @@ -1073,6 +1148,7 @@ fn response_text(text: &str) -> ModelResponse { content: vec![ContentBlock::Text(text.to_string())], tool_calls: Vec::new(), usage: None, + origin: None, }, usage: None, finish_reason: None, @@ -1480,3 +1556,402 @@ async fn agent_run_text_reflects_final_response() { run.final_response = Some(response_with_usage(Usage::new(1, 1))); assert_eq!(run.text(), Some("ok".to_string())); } + +// ── ContextCompressionMiddleware: overflow → compact → retry ────────────────── + +/// A model base that fails its first `fail_times` calls with a classified +/// context-overflow error, then succeeds. +struct OverflowThenSucceedBase { + calls: Arc>, + fail_times: usize, +} + +impl ModelBaseCall<(), ()> for OverflowThenSucceedBase { + fn call<'a>( + &'a self, + _ctx: &'a mut RunContext, + _state: &'a (), + _request: ModelRequest, + ) -> BoxModelFuture<'a> { + Box::pin(async move { + let attempt = { + let mut n = self.calls.lock().unwrap(); + *n += 1; + *n + }; + if attempt <= self.fail_times { + Err(TinyAgentsError::Model( + "This model's maximum context length is 100 tokens. However, your \ + messages resulted in 900 tokens." + .to_string(), + )) + } else { + Ok(response_text("recovered")) + } + }) + } +} + +/// A large-enough transcript that `find_cut_point` finds a real cut under a +/// small `keep_recent_tokens` budget: several long user/assistant turns, no +/// tool calls (pairing is exercised separately by `summarization::compaction` +/// tests). +fn overflow_prone_messages() -> Vec { + let big = "word ".repeat(60); + vec![ + user(&format!("first {big}")), + Message::assistant(format!("second {big}")), + user(&format!("third {big}")), + Message::assistant(format!("fourth {big}")), + user(&format!("fifth {big}")), + ] +} + +fn small_window_policy() -> SummarizationPolicy { + SummarizationPolicy::default() + .with_context_window(100) + .with_threshold_fraction(0.5) +} + +#[tokio::test] +async fn context_compression_overflow_retries_once_and_compacts() { + let calls = Arc::new(Mutex::new(0)); + let base = OverflowThenSucceedBase { + calls: calls.clone(), + fail_times: 1, + }; + let mw = Arc::new(ContextCompressionMiddleware::new(small_window_policy())); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push_model_middleware(mw.clone()); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + + let request = ModelRequest { + messages: overflow_prone_messages(), + ..Default::default() + }; + let response = stack + .run_wrapped_model(&mut c, &(), request, &base) + .await + .unwrap() + .into_response(); + + assert_eq!(response.text(), "recovered"); + // One failing call + one successful retry = exactly two base invocations. + assert_eq!(*calls.lock().unwrap(), 2); + + let compacted: Vec = recorder + .events() + .into_iter() + .map(|r| r.event) + .filter(|e| matches!(e, AgentEvent::Compacted { .. })) + .collect(); + assert_eq!(compacted.len(), 1); + assert!(matches!( + compacted[0], + AgentEvent::Compacted { + reason: crate::summarization::CompactionReason::Overflow, + .. + } + )); +} + +#[tokio::test] +async fn context_compression_overflow_propagates_after_second_failure() { + // Fails every call: the retry itself also overflows, so the middleware + // must give up after exactly one retry rather than looping. + let calls = Arc::new(Mutex::new(0)); + let base = OverflowThenSucceedBase { + calls: calls.clone(), + fail_times: usize::MAX, + }; + let mw = Arc::new(ContextCompressionMiddleware::new(small_window_policy())); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push_model_middleware(mw.clone()); + + let mut c = ctx(); + let request = ModelRequest { + messages: overflow_prone_messages(), + ..Default::default() + }; + let result = stack.run_wrapped_model(&mut c, &(), request, &base).await; + + assert!(result.is_err()); + // Original call + exactly one retry = two base invocations, not more. + assert_eq!(*calls.lock().unwrap(), 2); +} + +#[tokio::test] +async fn context_compression_wrap_model_ignores_unrelated_errors() { + // A non-overflow failure must propagate untouched, with no compaction + // attempted and no retry. + let calls = Arc::new(Mutex::new(0)); + struct AlwaysFailsBase { + calls: Arc>, + } + impl ModelBaseCall<(), ()> for AlwaysFailsBase { + fn call<'a>( + &'a self, + _ctx: &'a mut RunContext, + _state: &'a (), + _request: ModelRequest, + ) -> BoxModelFuture<'a> { + Box::pin(async move { + *self.calls.lock().unwrap() += 1; + Err(TinyAgentsError::Tool("boom".to_string())) + }) + } + } + let base = AlwaysFailsBase { + calls: calls.clone(), + }; + let mw = Arc::new(ContextCompressionMiddleware::new(small_window_policy())); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push_model_middleware(mw); + + let mut c = ctx(); + let request = ModelRequest { + messages: overflow_prone_messages(), + ..Default::default() + }; + let result = stack.run_wrapped_model(&mut c, &(), request, &base).await; + + assert!(matches!(result, Err(TinyAgentsError::Tool(_)))); + assert_eq!(*calls.lock().unwrap(), 1); +} + +#[tokio::test] +async fn context_compression_before_compaction_decline_leaves_transcript_untouched() { + let calls = Arc::new(Mutex::new(0)); + let base = OverflowThenSucceedBase { + calls: calls.clone(), + fail_times: 1, + }; + let mw = Arc::new( + ContextCompressionMiddleware::new(small_window_policy()) + .with_before_compaction(|_ctx: &CompactionContext| CompactionDecision::Decline), + ); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push_model_middleware(mw); + + let mut c = ctx(); + let request = ModelRequest { + messages: overflow_prone_messages(), + ..Default::default() + }; + let result = stack.run_wrapped_model(&mut c, &(), request, &base).await; + + // Declined: no retry happens, the original overflow error propagates. + assert!(result.is_err()); + assert_eq!(*calls.lock().unwrap(), 1); +} + +#[tokio::test] +async fn context_compression_threshold_decline_leaves_transcript_untouched() { + let policy = SummarizationPolicy { + keep_last: 1, + ..SummarizationPolicy::default() + } + .with_context_window(100) + .with_threshold_fraction(0.5); + let mw = Arc::new( + ContextCompressionMiddleware::new(policy) + .with_before_compaction(|_ctx: &CompactionContext| CompactionDecision::Decline), + ); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let mut c = ctx(); + let big = "a".repeat(200); + let before = vec![ + user(&format!("{big}-1")), + user(&format!("{big}-2")), + user(&format!("{big}-3")), + ]; + let mut request = ModelRequest { + messages: before.clone(), + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + + assert_eq!(request.messages, before); + assert!(mw.records().is_empty()); +} + +/// An in-memory [`CompactionSink`] that records every persisted +/// [`CompactionRecord`], for asserting the durable-persistence contract +/// without depending on `tinyagents-session`. +#[derive(Default)] +struct RecordingCompactionSink { + records: Mutex>, +} + +impl CompactionSink for RecordingCompactionSink { + fn persist(&self, record: &CompactionRecord) -> Result<()> { + self.records.lock().unwrap().push(record.clone()); + Ok(()) + } +} + +#[tokio::test] +async fn context_compression_persists_compaction_when_sink_is_attached() { + let policy = SummarizationPolicy { + keep_last: 1, + ..SummarizationPolicy::default() + } + .with_context_window(100) + .with_threshold_fraction(0.5); + let mw = Arc::new(ContextCompressionMiddleware::new(policy)); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let sink = Arc::new(RecordingCompactionSink::default()); + let mut c = ctx().with_compaction_sink(sink.clone()); + + let big = "a".repeat(200); + let mut request = ModelRequest { + messages: vec![ + user(&format!("{big}-1")), + user(&format!("{big}-2")), + user(&format!("{big}-3")), + ], + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + + let persisted = sink.records.lock().unwrap(); + assert_eq!(persisted.len(), 1); + assert_eq!( + persisted[0].reason, + crate::summarization::CompactionReason::Threshold + ); + assert_eq!(persisted[0].first_kept_index, 2); + assert!(persisted[0].tokens_before > 0); +} + +#[tokio::test] +async fn context_compression_no_sink_means_no_persistence_attempt() { + // No sink attached: compaction still runs and emits `Compacted`, it just + // has nowhere to persist to. This is mostly a "doesn't panic" check. + let policy = SummarizationPolicy { + keep_last: 1, + ..SummarizationPolicy::default() + } + .with_context_window(100) + .with_threshold_fraction(0.5); + let mw = Arc::new(ContextCompressionMiddleware::new(policy)); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw); + + let recorder = Arc::new(RecordingListener::new()); + let mut c = ctx(); + c.events.subscribe(recorder.clone()); + assert!(c.compaction_sink.is_none()); + + let big = "a".repeat(200); + let mut request = ModelRequest { + messages: vec![ + user(&format!("{big}-1")), + user(&format!("{big}-2")), + user(&format!("{big}-3")), + ], + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + + let compacted = recorder + .events() + .into_iter() + .filter(|r| matches!(r.event, AgentEvent::Compacted { .. })) + .count(); + assert_eq!(compacted, 1); +} + +#[tokio::test] +async fn context_compression_iterative_summary_threads_previous_summary() { + // Two successive threshold-triggered compactions on the same middleware + // instance: the second must see the first's summary as + // `SummaryRequest::previous_summary`. + let requests: Arc>> = + Arc::new(Mutex::new(Vec::new())); + + struct RecordingSummarizer { + requests: Arc>>, + } + + #[async_trait] + impl Summarizer for RecordingSummarizer { + async fn summarize(&self, messages: &[Message]) -> Result { + self.summarize_request(&crate::summarization::SummaryRequest::new( + messages.to_vec(), + )) + .await + } + + async fn summarize_request( + &self, + request: &crate::summarization::SummaryRequest, + ) -> Result { + self.requests.lock().unwrap().push(request.clone()); + crate::summarization::ConcatSummarizer + .summarize(&request.messages) + .await + } + } + + let policy = SummarizationPolicy { + keep_last: 1, + ..SummarizationPolicy::default() + } + .with_context_window(100) + .with_threshold_fraction(0.5); + let mw = Arc::new(ContextCompressionMiddleware::with_summarizer( + policy, + Box::new(RecordingSummarizer { + requests: requests.clone(), + }), + )); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(mw.clone()); + + let mut c = ctx(); + let big = "a".repeat(200); + + let mut request = ModelRequest { + messages: vec![ + user(&format!("{big}-1")), + user(&format!("{big}-2")), + user(&format!("{big}-3")), + ], + ..Default::default() + }; + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + + // Grow the (already-compacted) transcript back over threshold and compact + // again. + request.messages.push(user(&format!("{big}-4"))); + request.messages.push(user(&format!("{big}-5"))); + stack + .run_before_model(&mut c, &(), &mut request) + .await + .unwrap(); + + let seen = requests.lock().unwrap(); + assert_eq!(seen.len(), 2); + assert_eq!(seen[0].previous_summary, None); + assert!(seen[1].previous_summary.is_some()); +} diff --git a/crates/tinyagents-harness/src/middleware/types.rs b/crates/tinyagents-harness/src/middleware/types.rs index 32888f2f..a6a99481 100644 --- a/crates/tinyagents-harness/src/middleware/types.rs +++ b/crates/tinyagents-harness/src/middleware/types.rs @@ -23,7 +23,7 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; use crate::cache::CacheLayoutEvent; -use crate::context::RunContext; +use crate::context::{MiddlewareControl, RunContext}; use crate::error::{Result, TinyAgentsError}; use crate::events::HarnessRunStatus; use crate::ids::{CallId, RunId}; @@ -92,6 +92,10 @@ pub struct AgentRun { pub final_response: Option, /// Parsed structured output, when the run requested a structured format. pub structured: Option, + /// Which schema variant matched, when [`Self::structured`] was extracted + /// under [`crate::structured::StructuredStrategy::ToolCallUnion`] (A6). + /// `None` for every other strategy, and whenever `structured` is `None`. + pub structured_variant: Option, /// Cumulative token usage across every model call in the run. pub usage: UsageTotals, /// Number of model calls dispatched during the run. @@ -119,6 +123,42 @@ pub struct AgentRun { /// lifts it and a fresh invocation continues from /// [`AgentRun::messages`]. pub paused: Option, + /// Set when the run stopped because one or more tool calls were + /// **deferred** (A2): they need a human approval or host-side execution + /// before the loop can continue. Like [`Self::paused`], this is not a + /// completion — there is no `final_response`, and + /// [`HarnessRunStatus`][crate::events::HarnessRunStatus] reports the run + /// `Interrupted`. Persist [`Self::messages`] together with this value, + /// resolve it into a [`crate::tool::DeferredToolResults`], and resume + /// with [`crate::runtime::AgentHarness::resume_deferred`]. + pub deferred: Option, + /// Messages the host pushed onto the run queue's `Collect` lane (A4), + /// drained once when the run ends — on every exit path, including + /// errors. They are delivered here for the host to act on and are + /// **never** appended to the transcript or sent to the model. Empty when + /// the run had no queue. + pub collected: Vec, + /// Host-only metadata tools attached to their results + /// (`tinytools::ToolResult::metadata`, B2), one entry per answered call + /// that carried any, in fold order. Kept beside [`Self::executed_tools`] + /// rather than inside it so the name list stays a plain `Vec`. + /// The same value rides the call's + /// [`AgentEvent::ToolCompleted`][crate::events::AgentEvent::ToolCompleted]; + /// neither copy is ever rendered into [`Self::messages`]. + pub tool_metadata: Vec, +} + +/// Host-only metadata one tool call returned, as recorded on +/// [`AgentRun::tool_metadata`] (B2). +#[derive(Clone, Debug, PartialEq)] +pub struct ToolResultMetadata { + /// The call that produced it — matches the transcript row and the + /// `ToolCompleted` event. + pub call_id: CallId, + /// The tool the call named (after any unknown-tool rewrite). + pub tool_name: String, + /// The metadata verbatim; never shown to the model. + pub metadata: serde_json::Value, } // ── Middleware trait ────────────────────────────────────────────────────────── @@ -245,6 +285,115 @@ pub trait Middleware: Send + Sync { async fn on_error(&self, _ctx: &mut RunContext, _error: &TinyAgentsError) -> Result<()> { Ok(()) } + + // ── Control-outcome hooks ──────────────────────────────────────────── + // + // Each hook above has a `_control`-suffixed counterpart the + // [`MiddlewareStack`] actually drives. The default implementation below + // calls the plain hook and returns [`MiddlewareControl::Continue`], so + // every existing `Middleware` impl that only overrides the plain hooks + // keeps compiling and behaving exactly as before (A1's source-compat + // shim). Override a `_control` hook directly (instead of, not in + // addition to, the plain one) when the outcome needs to steer the loop — + // stop, jump, interrupt, or queue a state update. See + // `docs/modules/harness/middleware.md` for the precedence rule the stack + // applies across a phase's hooks and the checkpoints the loop honors a + // returned control at. + + /// Control-outcome counterpart of [`Self::before_agent`]. + async fn before_agent_control( + &self, + ctx: &mut RunContext, + state: &State, + ) -> Result { + self.before_agent(ctx, state).await?; + Ok(MiddlewareControl::Continue) + } + + /// Control-outcome counterpart of [`Self::after_agent`]. + async fn after_agent_control( + &self, + ctx: &mut RunContext, + state: &State, + run: &mut AgentRun, + ) -> Result { + self.after_agent(ctx, state, run).await?; + Ok(MiddlewareControl::Continue) + } + + /// Control-outcome counterpart of [`Self::before_model`]. + async fn before_model_control( + &self, + ctx: &mut RunContext, + state: &State, + request: &mut ModelRequest, + ) -> Result { + self.before_model(ctx, state, request).await?; + Ok(MiddlewareControl::Continue) + } + + /// Control-outcome counterpart of [`Self::after_model`]. + async fn after_model_control( + &self, + ctx: &mut RunContext, + state: &State, + response: &mut ModelResponse, + ) -> Result { + self.after_model(ctx, state, response).await?; + Ok(MiddlewareControl::Continue) + } + + /// Control-outcome counterpart of [`Self::before_tool`]. + async fn before_tool_control( + &self, + ctx: &mut RunContext, + state: &State, + call: &mut ToolCall, + ) -> Result { + self.before_tool(ctx, state, call).await?; + Ok(MiddlewareControl::Continue) + } + + /// Control-outcome counterpart of [`Self::after_tool`]. + async fn after_tool_control( + &self, + ctx: &mut RunContext, + state: &State, + invocation: &ToolInvocationIdentity, + result: &mut ToolResult, + ) -> Result { + self.after_tool(ctx, state, invocation, result).await?; + Ok(MiddlewareControl::Continue) + } + + /// Whether this middleware still runs (for observation) in a phase where + /// an earlier middleware already produced a winning control outcome. + /// + /// The stack applies the *first* non-[`MiddlewareControl::Continue`] + /// outcome in a phase and, by default, skips every hook after it — an + /// early-exit tool guard or a budget stop should not pay for hooks whose + /// work is now moot. A middleware that must still observe every call + /// regardless (a usage accountant, an audit log) overrides this to + /// `true`; its own control outcome is then ignored; only the first + /// winning one is ever applied. See `docs/modules/harness/middleware.md`. + fn is_observer(&self) -> bool { + false + } + + /// Whether the loop should stop after the turn currently completing, + /// evaluated once at the turn boundary (after tool execution, before the + /// loop would otherwise continue to the next model call). + /// + /// Defaults to `false`. A middleware that returns `true` here has the + /// same effect as requesting + /// [`MiddlewareControl::JumpTo`]`(`[`crate::context::LoopTarget::End`]`)` + /// from `after_tool_control`, but expresses "stop once this turn settles" + /// without needing to compute that decision inside `after_tool_control` + /// itself (useful when the decision depends on the whole turn's tool + /// results, not just one call). + fn should_stop_after_turn(&self, _ctx: &RunContext, _run: &AgentRun) -> bool { + false + } } // ── Wrap (around-call) middleware ───────────────────────────────────────────── @@ -349,18 +498,47 @@ pub trait ToolBaseCall: Send + Sync { /// `next` rather than by distinct enum variants; the enum only needs to carry /// the resolved response. It is `#[non_exhaustive]` so future control variants /// can be added without breaking callers. +// `Response(ModelResponse)` is large relative to `Command`'s payload; boxing +// it would ripple through every construction/destructure site across the +// crate (including the `From` impl below and every wrap +// middleware) for a value that lives only as long as one model call, so the +// size skew is accepted here rather than threaded through as indirection. #[derive(Clone, Debug)] #[non_exhaustive] +#[allow(clippy::large_enum_variant)] pub enum MiddlewareModelOutcome { /// The response to use as the result of the wrapped model call. Response(ModelResponse), + /// Short-circuit with a [`MiddlewareControl`] instead of a response — for + /// example a wrap middleware that decides, before ever calling `next`, + /// that the run should stop or jump. There is no response to hand back in + /// this case, so callers that need one (see [`Self::into_response`]) get + /// an empty placeholder; the control itself is recovered separately, via + /// [`Self::into_response_with_control`], and applied through the same + /// [`RunContext::request_control`][crate::context::RunContext::request_control] + /// path a lifecycle hook's control-outcome return uses. + Command { + /// The control outcome to apply. + control: MiddlewareControl, + }, } impl MiddlewareModelOutcome { - /// Unwraps the contained [`ModelResponse`]. + /// Unwraps the contained [`ModelResponse`], or an empty placeholder for + /// [`Self::Command`] (see that variant's docs — prefer + /// [`Self::into_response_with_control`] when a `Command` must not be + /// silently discarded). pub fn into_response(self) -> ModelResponse { + self.into_response_with_control().0 + } + + /// Splits this outcome into a [`ModelResponse`] (a placeholder for + /// [`Self::Command`]) and the [`MiddlewareControl`] to apply, when this + /// was a `Command` outcome. + pub fn into_response_with_control(self) -> (ModelResponse, Option) { match self { - Self::Response(response) => response, + Self::Response(response) => (response, None), + Self::Command { control } => (ModelResponse::assistant(String::new()), Some(control)), } } } @@ -381,13 +559,30 @@ impl From for MiddlewareModelOutcome { pub enum MiddlewareToolOutcome { /// The result to use as the result of the wrapped tool call. Result(ToolResult), + /// Short-circuit with a [`MiddlewareControl`] instead of a result. The + /// tool-wrap counterpart of [`MiddlewareModelOutcome::Command`]; see its + /// docs for the placeholder-result and control-recovery contract. + Command { + /// The control outcome to apply. + control: MiddlewareControl, + }, } impl MiddlewareToolOutcome { - /// Unwraps the contained [`ToolResult`]. + /// Unwraps the contained [`ToolResult`], or an empty error placeholder for + /// [`Self::Command`] (prefer [`Self::into_result_with_control`] when a + /// `Command` must not be silently discarded). pub fn into_result(self) -> ToolResult { + self.into_result_with_control().0 + } + + /// Splits this outcome into a [`ToolResult`] (a placeholder for + /// [`Self::Command`]) and the [`MiddlewareControl`] to apply, when this + /// was a `Command` outcome. + pub fn into_result_with_control(self) -> (ToolResult, Option) { match self { - Self::Result(result) => result, + Self::Result(result) => (result, None), + Self::Command { control } => (ToolResult::success(String::new()), Some(control)), } } } @@ -447,6 +642,21 @@ pub trait ModelMiddleware: Send + Syn /// `MiddlewareStarted`/`MiddlewareCompleted` events. fn name(&self) -> &str; + /// Whether this middleware already retries the model call itself (as + /// [`crate::middleware::library::RetryMiddleware`] does). + /// + /// [`MiddlewareStack::has_retry_override`] uses this to tell the loop's + /// base call to skip its own [`crate::runtime::RunPolicy::retry`] loop + /// when one is registered — otherwise the two retry layers compose + /// multiplicatively (`mw.max_attempts × policy.retry.max_attempts × + /// |fallback|` provider calls for one logical failure) instead of + /// replacing each other. See I-7; full unification into one engine is a + /// later phase. Defaults to `false` so an ordinary middleware is + /// unaffected. + fn overrides_retry(&self) -> bool { + false + } + /// Wraps the inner model pipeline. Call `next.run(ctx, state, request)` to /// proceed (zero or more times), or return a [`MiddlewareModelOutcome`] /// without calling it to short-circuit. @@ -620,6 +830,15 @@ pub enum CompressionFailurePolicy { PassThrough, } +/// The type of a `before_compaction` hook, consulted before every compaction +/// [`ContextCompressionMiddleware`] runs. Named to keep the struct field's +/// type simple (`clippy::type_complexity`). +pub type BeforeCompactionHook = std::sync::Arc< + dyn Fn(&crate::summarization::CompactionContext) -> crate::summarization::CompactionDecision + + Send + + Sync, +>; + /// Middleware that summarizes/compresses the request transcript, but **only** /// when it nears the model's context window. /// @@ -656,6 +875,25 @@ pub struct ContextCompressionMiddleware { pub(crate) max_records: usize, /// Recovery behaviour when [`Summarizer::summarize`] returns `Err`. pub(crate) on_failure: CompressionFailurePolicy, + /// The most recently produced compaction summary text, threaded into the + /// next compaction's [`crate::summarization::SummaryRequest::previous_summary`] + /// so an iterative [`Summarizer`] refines rather than restarts. `None` + /// until the first compaction on this middleware instance. + pub(crate) last_summary: Mutex>, + /// Token budget above which a single "turn" of messages handed to the + /// summarizer is itself split into two halves and merged (see + /// [`crate::summarization::summarize_with_split`]). `None` disables + /// splitting — the whole `to_summarize` slice is always summarized in one + /// call, matching the middleware's original behaviour. + pub(crate) max_turn_tokens: Option, + /// Classifies a model-call error as a provider context-window overflow, + /// consulted by [`ModelMiddleware::wrap_model`] for the + /// overflow → compact → retry recovery path. + pub(crate) overflow_classifier: crate::summarization::OverflowClassifier, + /// Optional hook consulted before every compaction (proactive or + /// overflow-triggered) that can decline it or substitute a summary. See + /// [`crate::summarization::CompactionDecision`]. + pub(crate) before_compaction: Option, } // ── MicrocompactMiddleware ──────────────────────────────────────────────────── diff --git a/crates/tinyagents-harness/src/multimodal/markers.rs b/crates/tinyagents-harness/src/multimodal/markers.rs index 41c42735..7d78a79f 100644 --- a/crates/tinyagents-harness/src/multimodal/markers.rs +++ b/crates/tinyagents-harness/src/multimodal/markers.rs @@ -150,7 +150,7 @@ pub fn extract_ollama_image_payload(image_ref: &str) -> Option { return None; } if !is_data_uri && looks_like_absolute_path(payload) { - tinyagents_tracing::debug!( + tracing::debug!( "[multimodal] image reference is shaped like a filesystem path, not image bytes" ); return None; @@ -172,7 +172,7 @@ pub fn extract_ollama_image_payload(image_ref: &str) -> Option { STANDARD_NO_PAD.decode(payload).is_ok() }; if !is_base64 { - tinyagents_tracing::debug!( + tracing::debug!( "[multimodal] image reference is not base64 (a filesystem path is not accepted here)" ); return None; diff --git a/crates/tinyagents-harness/src/multimodal/resolve.rs b/crates/tinyagents-harness/src/multimodal/resolve.rs index 06424190..7b4ae3c3 100644 --- a/crates/tinyagents-harness/src/multimodal/resolve.rs +++ b/crates/tinyagents-harness/src/multimodal/resolve.rs @@ -388,7 +388,7 @@ async fn build_file_payload( }); } - tinyagents_tracing::debug!( + tracing::debug!( target: "multimodal", file = %name, mime = %mime, @@ -407,7 +407,7 @@ async fn build_file_payload( match extractor.extract(&mime, &bytes).await { Ok(text) => Some(text), Err(reason) => { - tinyagents_tracing::warn!( + tracing::warn!( target: "multimodal", file = %name, mime = %mime, @@ -429,7 +429,7 @@ async fn build_file_payload( } = &payload && *truncated_chars > 0 { - tinyagents_tracing::info!( + tracing::info!( target: "multimodal", file = %name, truncated_chars, diff --git a/crates/tinyagents-harness/src/observability/langfuse/mod.rs b/crates/tinyagents-harness/src/observability/langfuse/mod.rs index 1132ea95..8207fd71 100644 --- a/crates/tinyagents-harness/src/observability/langfuse/mod.rs +++ b/crates/tinyagents-harness/src/observability/langfuse/mod.rs @@ -543,6 +543,11 @@ fn observation_event( duration_ms, output_bytes, error, + // The tool's host-only result metadata (B2) is not exported: the + // observation's `metadata` is the exporter's own correlation + // record, and a tool payload of arbitrary size/shape belongs in a + // deliberate mapping, not merged in by default. + metadata: _, } => { // Prefer the loop-captured start + real duration for the end time; // fall back to the journal timestamp. A failed call is marked ERROR @@ -652,42 +657,19 @@ pub fn clean_nulls(mut value: Value) -> Value { value } -/// Formats a Unix-epoch millisecond timestamp as the UTC ISO-8601 string -/// Langfuse's ingestion API expects (`YYYY-MM-DDTHH:MM:SS.sssZ`). +/// Renders a Unix epoch millisecond timestamp as the `YYYY-MM-DDTHH:MM:SS.sssZ` +/// form Langfuse's ingestion API expects. +/// +/// Delegates to `chrono` (M-9): `chrono` is already a non-optional workspace +/// dependency of this crate (`tools/time.rs` uses it under the `tools` +/// feature), so the hand-rolled Howard Hinnant civil-date conversion this +/// module carried was duplicating logic the dependency graph already pays +/// for, unconditionally, elsewhere. pub fn iso_ms(ms: u64) -> String { - use std::time::{Duration, UNIX_EPOCH}; - let system_time = UNIX_EPOCH + Duration::from_millis(ms); - let duration = system_time - .duration_since(UNIX_EPOCH) - .unwrap_or(Duration::from_secs(0)); - let secs = duration.as_secs(); - let millis = duration.subsec_millis(); - format_unix_iso(secs, millis) -} - -fn format_unix_iso(secs: u64, millis: u32) -> String { - // Howard Hinnant civil-date conversion for Unix days, dependency-free. - let days = (secs / 86_400) as i64; - let day_secs = secs % 86_400; - let (year, month, day) = civil_from_days(days); - let hour = day_secs / 3_600; - let minute = (day_secs % 3_600) / 60; - let second = day_secs % 60; - format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}.{millis:03}Z") -} - -fn civil_from_days(days: i64) -> (i32, u32, u32) { - let z = days + 719_468; - let era = if z >= 0 { z } else { z - 146_096 } / 146_097; - let doe = z - era * 146_097; - let yoe = (doe - doe / 1_460 + doe / 36_524 - doe / 146_096) / 365; - let y = yoe + era * 400; - let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); - let mp = (5 * doy + 2) / 153; - let d = doy - (153 * mp + 2) / 5 + 1; - let m = mp + if mp < 10 { 3 } else { -9 }; - let year = y + if m <= 2 { 1 } else { 0 }; - (year as i32, m as u32, d as u32) + chrono::DateTime::::from_timestamp_millis(i64::try_from(ms).unwrap_or(i64::MAX)) + .unwrap_or_default() + .format("%Y-%m-%dT%H:%M:%S%.3fZ") + .to_string() } #[cfg(test)] diff --git a/crates/tinyagents-harness/src/observability/langfuse/test.rs b/crates/tinyagents-harness/src/observability/langfuse/test.rs index ec1b9d4c..3289ee90 100644 --- a/crates/tinyagents-harness/src/observability/langfuse/test.rs +++ b/crates/tinyagents-harness/src/observability/langfuse/test.rs @@ -170,6 +170,7 @@ fn populates_generation_and_tool_io_when_captured() { duration_ms: Some(250), output_bytes: Some(5), error: None, + metadata: None, }, ), ], @@ -272,6 +273,7 @@ fn call_scoped_observation_ids_are_unique_per_trace() { duration_ms: None, output_bytes: None, error: None, + metadata: None, }, ), ], diff --git a/crates/tinyagents-harness/src/observability/mod.rs b/crates/tinyagents-harness/src/observability/mod.rs index 2d2934f5..adc30b38 100644 --- a/crates/tinyagents-harness/src/observability/mod.rs +++ b/crates/tinyagents-harness/src/observability/mod.rs @@ -24,6 +24,7 @@ //! bounded queue drops rather than stalls; backend errors are reported, not //! propagated), and `flush` blocks until the durable log has caught up. +#[cfg(feature = "langfuse")] mod langfuse; mod profile; mod types; @@ -32,12 +33,14 @@ mod worker; #[doc(hidden)] pub use worker::{AppendWorker, DEFAULT_DRAIN_CAPACITY}; +#[cfg(feature = "langfuse")] pub use langfuse::{ LangfuseAuth, LangfuseClient, LangfuseScore, LangfuseScoreValue, LangfuseTraceConfig, }; pub use profile::{ProcessProfile, ProcessProfiler, ProcessSnapshot}; // Shared Langfuse payload helpers reused by the graph observability exporter so // ISO-8601 timestamp formatting and null-field pruning live in one place. +#[cfg(feature = "langfuse")] #[doc(hidden)] pub use langfuse::{clean_nulls, iso_ms}; pub use types::*; diff --git a/crates/tinyagents-harness/src/observability/profile.rs b/crates/tinyagents-harness/src/observability/profile.rs index 72165d56..d5cf718d 100644 --- a/crates/tinyagents-harness/src/observability/profile.rs +++ b/crates/tinyagents-harness/src/observability/profile.rs @@ -176,27 +176,28 @@ fn current_rss_bytes() -> Option { None } -#[cfg(unix)] +#[cfg(target_os = "linux")] fn cpu_time_us() -> Option<(u64, u64)> { - let mut usage = std::mem::MaybeUninit::::zeroed(); - // SAFETY: `getrusage` initializes the provided `rusage` on success and we - // only call `assume_init` after checking its zero return code. - if unsafe { libc::getrusage(libc::RUSAGE_SELF, usage.as_mut_ptr()) } != 0 { - return None; - } - // SAFETY: established by the successful `getrusage` call above. - let usage = unsafe { usage.assume_init() }; - Some((timeval_us(usage.ru_utime), timeval_us(usage.ru_stime))) + static CLOCK_TICKS_PER_SECOND: std::sync::OnceLock = std::sync::OnceLock::new(); + let stat = std::fs::read_to_string("/proc/self/stat").ok()?; + let fields: Vec<_> = stat.rsplit_once(')')?.1.split_whitespace().collect(); + let user_ticks = fields.get(11)?.parse::().ok()?; + let system_ticks = fields.get(12)?.parse::().ok()?; + let ticks_per_second = *CLOCK_TICKS_PER_SECOND.get_or_init(|| { + std::process::Command::new("getconf") + .arg("CLK_TCK") + .output() + .ok() + .and_then(|output| String::from_utf8(output.stdout).ok()) + .and_then(|value| value.trim().parse::().ok()) + .filter(|value| *value > 0) + .unwrap_or(100) + }); + let to_micros = |ticks: u64| ticks.saturating_mul(1_000_000) / ticks_per_second; + Some((to_micros(user_ticks), to_micros(system_ticks))) } -#[cfg(unix)] -fn timeval_us(value: libc::timeval) -> u64 { - let seconds = u64::try_from(value.tv_sec).unwrap_or(0); - let micros = u64::try_from(value.tv_usec).unwrap_or(0); - seconds.saturating_mul(1_000_000).saturating_add(micros) -} - -#[cfg(not(unix))] +#[cfg(not(target_os = "linux"))] fn cpu_time_us() -> Option<(u64, u64)> { None } diff --git a/crates/tinyagents-harness/src/observability/test.rs b/crates/tinyagents-harness/src/observability/test.rs index 941b4782..96dff4b4 100644 --- a/crates/tinyagents-harness/src/observability/test.rs +++ b/crates/tinyagents-harness/src/observability/test.rs @@ -146,6 +146,7 @@ fn agent_latency_metrics_include_model_tool_and_run_elapsed() { duration_ms: None, output_bytes: None, error: None, + metadata: None, }, ), obs( diff --git a/crates/tinyagents-harness/src/observability/worker.rs b/crates/tinyagents-harness/src/observability/worker.rs index a7988d9f..25ff9b64 100644 --- a/crates/tinyagents-harness/src/observability/worker.rs +++ b/crates/tinyagents-harness/src/observability/worker.rs @@ -178,7 +178,7 @@ impl AppendWorker { Msg::Item(item) => match append(item).await { Ok(()) => { if failure_run > 0 { - tinyagents_tracing::warn!( + tracing::warn!( target: "tinyagents::observability", sink = name, lost = failure_run, @@ -194,7 +194,7 @@ impl AppendWorker { let now = Instant::now(); if failure_run == 1 { last_report = Some(now); - tinyagents_tracing::error!( + tracing::error!( target: "tinyagents::observability", sink = name, error = %error, @@ -202,7 +202,7 @@ impl AppendWorker { ); } else if should_report(last_report, now, cooldown) { last_report = Some(now); - tinyagents_tracing::warn!( + tracing::warn!( target: "tinyagents::observability", sink = name, error = %error, @@ -220,7 +220,7 @@ impl AppendWorker { // The channel closed mid-failure: report once on the way out // so a run that never recovered is not silently quiet. if failure_run > 0 { - tinyagents_tracing::warn!( + tracing::warn!( target: "tinyagents::observability", sink = name, lost = failure_run, diff --git a/crates/tinyagents-harness/src/providers/claude_agent_sdk/mod.rs b/crates/tinyagents-harness/src/providers/claude_agent_sdk/mod.rs index 914eaf3e..7f271b22 100644 --- a/crates/tinyagents-harness/src/providers/claude_agent_sdk/mod.rs +++ b/crates/tinyagents-harness/src/providers/claude_agent_sdk/mod.rs @@ -105,9 +105,11 @@ fn build_invocation( /// transcript is present so the model can distinguish its own prior output /// from the next user turn. fn render_transcript(messages: &[Message]) -> String { + // `Message::Custom` is a host-side out-of-band record (e.g. a compaction + // marker); it never rides to a provider transcript. let non_system: Vec<&Message> = messages .iter() - .filter(|message| !matches!(message, Message::System(_))) + .filter(|message| !matches!(message, Message::System(_) | Message::Custom(_))) .collect(); if non_system.len() == 1 { return non_system[0].text(); @@ -121,6 +123,7 @@ fn render_transcript(messages: &[Message]) -> String { Message::Assistant(_) => "ASSISTANT", Message::Tool(_) => "TOOL", Message::System(_) => unreachable!("system messages were filtered"), + Message::Custom(_) => unreachable!("custom messages were filtered"), }; format!("[{role}]\n{}\n[/{role}]", message.text()) }) @@ -191,7 +194,7 @@ impl ClaudeAgentSdkProvider { .stdin(std::process::Stdio::piped()) .kill_on_drop(true); - tinyagents_tracing::debug!( + tracing::debug!( "[claude_agent_sdk] spawning claude binary={} model={} message_len={}", self.config.binary, model, @@ -199,7 +202,7 @@ impl ClaudeAgentSdkProvider { ); let mut child = cmd.spawn().map_err(|source| { - tinyagents_tracing::warn!( + tracing::warn!( error = %source, binary = %self.config.binary, "[claude_agent_sdk] failed to spawn claude binary" @@ -254,7 +257,7 @@ impl ClaudeAgentSdkProvider { if line.is_empty() { continue; } - tinyagents_tracing::trace!( + tracing::trace!( "[claude_agent_sdk] ndjson line received line_len={}", line.len() ); @@ -268,7 +271,7 @@ impl ClaudeAgentSdkProvider { total_cost_usd, }) => { if let Some(cost) = total_cost_usd { - tinyagents_tracing::debug!( + tracing::debug!( "[claude_agent_sdk] request completed total_cost_usd={:.6}", cost ); @@ -285,12 +288,10 @@ impl ClaudeAgentSdkProvider { error_message = Some(error.message); } Ok(SdkMessage::Unknown) => { - tinyagents_tracing::trace!( - "[claude_agent_sdk] unknown ndjson message type, skipping" - ); + tracing::trace!("[claude_agent_sdk] unknown ndjson message type, skipping"); } Err(e) => { - tinyagents_tracing::warn!( + tracing::warn!( error = %e, line_len = line.len(), "[claude_agent_sdk] failed to parse ndjson line" @@ -316,7 +317,7 @@ impl ClaudeAgentSdkProvider { anyhow::anyhow!("[claude_agent_sdk] subprocess timed out while waiting for exit") })??; let stderr_output = stderr_task.await.unwrap_or_default(); - tinyagents_tracing::debug!("[claude_agent_sdk] subprocess exited status={}", status); + tracing::debug!("[claude_agent_sdk] subprocess exited status={}", status); if !status.success() { anyhow::bail!( @@ -335,7 +336,7 @@ impl ClaudeAgentSdkProvider { .filter(|s| !s.is_empty()) .unwrap_or_else(|| text_parts.join("")); - tinyagents_tracing::debug!( + tracing::debug!( "[claude_agent_sdk] response collected output_len={}", output.len() ); diff --git a/crates/tinyagents-harness/src/providers/claude_code/auth_status.rs b/crates/tinyagents-harness/src/providers/claude_code/auth_status.rs index 16afdd24..56d1ddd3 100644 --- a/crates/tinyagents-harness/src/providers/claude_code/auth_status.rs +++ b/crates/tinyagents-harness/src/providers/claude_code/auth_status.rs @@ -159,7 +159,7 @@ pub fn parse_auth_status_json(raw: &str) -> AuthSource { /// `OPENHUMAN_CLAUDE_CLI` override via [`version_check::resolve_binary`]. fn probe_via_cli() -> AuthSource { let Some(bin) = version_check::resolve_binary() else { - log::debug!("[claude-code][auth] no `claude` binary on PATH; auth state unknown"); + tracing::debug!("[claude-code][auth] no `claude` binary on PATH; auth state unknown"); return AuthSource::Unknown { reason: Some("`claude` CLI not found on PATH".to_string()), }; @@ -176,7 +176,7 @@ fn probe_via_cli() -> AuthSource { { Ok(c) => c, Err(e) => { - log::warn!("[claude-code][auth] spawn failed bin={bin_str} err={e}"); + tracing::warn!("[claude-code][auth] spawn failed bin={bin_str} err={e}"); return AuthSource::Unknown { reason: Some(format!("spawn failed: {e}")), }; @@ -189,7 +189,7 @@ fn probe_via_cli() -> AuthSource { let status = match child.wait_timeout(AUTH_STATUS_TIMEOUT) { Ok(Some(s)) => s, Ok(None) => { - log::warn!( + tracing::warn!( "[claude-code][auth] `claude auth status` timed out after {}s; killing bin={bin_str}", AUTH_STATUS_TIMEOUT.as_secs() ); @@ -203,7 +203,7 @@ fn probe_via_cli() -> AuthSource { }; } Err(e) => { - log::warn!("[claude-code][auth] wait failed bin={bin_str} err={e}"); + tracing::warn!("[claude-code][auth] wait failed bin={bin_str} err={e}"); let _ = child.kill(); let _ = child.wait(); return AuthSource::Unknown { @@ -219,7 +219,7 @@ fn probe_via_cli() -> AuthSource { if let Some(mut s) = child.stderr.take() { let _ = s.read_to_string(&mut stderr); } - log::debug!( + tracing::debug!( "[claude-code][auth] `claude auth status` exit={} stderr={}", status, stderr.trim() @@ -234,7 +234,7 @@ fn probe_via_cli() -> AuthSource { let _ = s.read_to_string(&mut stdout); } let source = parse_auth_status_json(stdout.trim()); - log::debug!( + tracing::debug!( "[claude-code][auth] probe classified source={}", match &source { AuthSource::Subscription { .. } => "subscription", @@ -258,7 +258,7 @@ pub fn probe() -> AuthStatus { if let Ok(k) = std::env::var("ANTHROPIC_API_KEY") && !k.trim().is_empty() { - log::debug!("[claude-code][auth] ANTHROPIC_API_KEY present → api_key_env"); + tracing::debug!("[claude-code][auth] ANTHROPIC_API_KEY present → api_key_env"); return AuthStatus { source: AuthSource::ApiKeyEnv, last_checked, diff --git a/crates/tinyagents-harness/src/providers/claude_code/driver.rs b/crates/tinyagents-harness/src/providers/claude_code/driver.rs index 54815db8..7f720c76 100644 --- a/crates/tinyagents-harness/src/providers/claude_code/driver.rs +++ b/crates/tinyagents-harness/src/providers/claude_code/driver.rs @@ -323,20 +323,20 @@ fn append_system_prompt_args( }; let path = dir.join("append-system-prompt.txt"); - log::debug!( + tracing::debug!( "[claude-code][driver] append-system-prompt file write start path={} bytes={}", path.display(), prompt.len() ); if let Err(error) = std::fs::write(&path, prompt) { - log::warn!( + tracing::warn!( "[claude-code][driver] append-system-prompt file write failed path={} error={}", path.display(), error ); return Err(error); } - log::debug!( + tracing::debug!( "[claude-code][driver] append-system-prompt file write complete path={} bytes={}", path.display(), prompt.len() @@ -373,19 +373,19 @@ pub(crate) async fn run_turn(ctx: TurnContext<'_>) -> anyhow::Result { match write_mcp_http_config(scratch.path(), endpoint.addr, &endpoint.token) { Ok(p) => { - log::debug!( + tracing::debug!( "[claude-code][driver] wrote http mcp-config path={} url=http://{}/ (authenticated)", p.display(), endpoint.addr ); mcp_config_path = Some(p); } - Err(e) => log::warn!( + Err(e) => tracing::warn!( "[claude-code][driver] failed to write mcp-config: {e}; CC will run without OpenHuman MCP tools" ), } } - Err(e) => log::warn!( + Err(e) => tracing::warn!( "[claude-code][driver] in-process MCP HTTP server unavailable: {e}; CC running without OpenHuman MCP tools" ), } @@ -462,7 +462,7 @@ pub(crate) async fn run_turn(ctx: TurnContext<'_>) -> anyhow::Result) -> anyhow::Result) -> anyhow::Result) -> anyhow::Result) -> anyhow::Result inner?, Err(_elapsed) => { - log::error!("[claude-code][driver] turn timeout ({timeout:?}) exceeded; killing child"); + tracing::error!( + "[claude-code][driver] turn timeout ({timeout:?}) exceeded; killing child" + ); // kill_on_drop handles cleanup, but explicit kill gives us // a chance to collect stderr. let _ = child.kill().await; @@ -626,7 +628,7 @@ pub(crate) async fn run_turn(ctx: TurnContext<'_>) -> anyhow::Result Vec pub(crate) static ENV_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); #[cfg(test)] +#[allow(unsafe_code)] pub(crate) fn test_set_env(key: impl AsRef, value: impl AsRef) { // SAFETY: every moved environment-mutating test serializes access through // `ENV_TEST_LOCK`; no provider work runs concurrently in those tests. @@ -95,6 +96,7 @@ pub(crate) fn test_set_env(key: impl AsRef, value: impl AsRef) { // SAFETY: see `test_set_env`. unsafe { std::env::remove_var(key) } @@ -221,12 +223,13 @@ impl ClaudeCodeProvider { model_override: Option<&str>, thread_id: String, ) -> anyhow::Result { - let _permit = self - .semaphore - .clone() - .acquire_owned() - .await - .map_err(|error| anyhow::anyhow!("claude-code semaphore closed: {error}"))?; + // Acquire the per-thread mutex *before* the global concurrency + // semaphore (M-14). Reversed, N callers on one busy thread each hold + // a global permit while blocked on the same thread lock — that is + // head-of-line blocking for every *other* thread's turns, which the + // semaphore exists to admit. Waiting on the free, per-thread lock + // first means a caller only claims a global permit once it can + // actually make progress. let lock_key = thread_id.clone(); let thread_lock = { let mut locks = self @@ -239,6 +242,12 @@ impl ClaudeCodeProvider { .clone() }; let _thread_guard = thread_lock.lock().await; + let _permit = self + .semaphore + .clone() + .acquire_owned() + .await + .map_err(|error| anyhow::anyhow!("claude-code semaphore closed: {error}"))?; let append_system_prompt = coalesce_system_prompt(messages); let result = driver::run_turn(driver::TurnContext { bin_path: self.bin_path.clone(), @@ -332,18 +341,23 @@ fn request_messages(request: &ModelRequest) -> Vec { } messages .iter() + // `Message::Custom` is a host-side out-of-band record; never sent to + // the provider. + .filter(|message| !matches!(message, Message::Custom(_))) .map(|message| { let role = match message { Message::System(_) => "system", Message::User(_) => "user", Message::Assistant(_) => "assistant", Message::Tool(_) => "tool", + Message::Custom(_) => unreachable!("custom messages were filtered"), }; let content = match message { Message::System(value) => render_content(&value.content), Message::User(value) => render_content(&value.content), Message::Assistant(value) => render_content(&value.content), Message::Tool(value) => render_content(&value.content), + Message::Custom(_) => unreachable!("custom messages were filtered"), }; ChatMessage::new(role, content) }) @@ -387,6 +401,9 @@ fn render_content(content: &[ContentBlock]) -> String { } ContentBlock::Thinking { text, .. } => Some(text.clone()), ContentBlock::RedactedThinking { .. } => None, + ContentBlock::Audio(media) => Some(format!("[OH_AUDIO:{media:?}]")), + ContentBlock::Video(media) => Some(format!("[OH_VIDEO:{media:?}]")), + ContentBlock::Document(media) => Some(format!("[OH_DOCUMENT:{media:?}]")), }) .collect::>() // Content-block boundaries carry no implicit whitespace. Inserting a @@ -416,6 +433,7 @@ fn model_response(response: ChatResponse) -> ModelResponse { content: response.text.into_iter().map(ContentBlock::Text).collect(), tool_calls: Vec::new(), usage, + origin: None, }, usage, finish_reason: Some("stop".into()), diff --git a/crates/tinyagents-harness/src/providers/claude_code/settings.rs b/crates/tinyagents-harness/src/providers/claude_code/settings.rs index 14245379..4022aba3 100644 --- a/crates/tinyagents-harness/src/providers/claude_code/settings.rs +++ b/crates/tinyagents-harness/src/providers/claude_code/settings.rs @@ -40,14 +40,14 @@ pub fn load(workspace_dir: &Path) -> ClaudeCodeSettings { let path = settings_path(workspace_dir); match std::fs::read(&path) { Ok(bytes) => serde_json::from_slice(&bytes).unwrap_or_else(|e| { - log::warn!( + tracing::warn!( "[claude-code][settings] corrupt {} ({e}); using safe defaults", path.display() ); ClaudeCodeSettings::default() }), Err(e) => { - log::debug!( + tracing::debug!( "[claude-code][settings] no settings at {} ({e}); using defaults", path.display() ); @@ -64,7 +64,7 @@ pub fn save(workspace_dir: &Path, settings: &ClaudeCodeSettings) -> std::io::Res } let json = serde_json::to_vec_pretty(settings).map_err(std::io::Error::other)?; std::fs::write(&path, json)?; - log::debug!( + tracing::debug!( "[claude-code][settings] saved full_access={} → {}", settings.full_access, path.display() diff --git a/crates/tinyagents-harness/src/providers/claude_code/version_check.rs b/crates/tinyagents-harness/src/providers/claude_code/version_check.rs index a351c1fa..ed1e93b0 100644 --- a/crates/tinyagents-harness/src/providers/claude_code/version_check.rs +++ b/crates/tinyagents-harness/src/providers/claude_code/version_check.rs @@ -37,7 +37,7 @@ pub fn resolve_binary() -> Option { // Finder/Dock-launch case where `~/.local/bin` is absent from `PATH`. let found = first_existing(&well_known_candidates()); if let Some(p) = found.as_ref() { - log::debug!( + tracing::debug!( "[claude-code][version] `claude` not on PATH; resolved via well-known location {}", p.display() ); @@ -120,7 +120,7 @@ fn which_on_path(name: &str) -> Option { /// Probe the `claude` CLI and return its status. pub fn probe() -> CliStatus { let Some(path) = resolve_binary() else { - log::debug!("[claude-code][version] no `claude` binary on PATH"); + tracing::debug!("[claude-code][version] no `claude` binary on PATH"); return CliStatus::NotInstalled; }; let path_str = path.display().to_string(); @@ -132,7 +132,7 @@ pub fn probe() -> CliStatus { { Ok(o) => o, Err(e) => { - log::warn!("[claude-code][version] spawn failed path={path_str} err={e}"); + tracing::warn!("[claude-code][version] spawn failed path={path_str} err={e}"); return CliStatus::Unusable { path: path_str, reason: format!("spawn failed: {e}"), diff --git a/crates/tinyagents-harness/src/providers/mod.rs b/crates/tinyagents-harness/src/providers/mod.rs index 90863e30..1b27fde2 100644 --- a/crates/tinyagents-harness/src/providers/mod.rs +++ b/crates/tinyagents-harness/src/providers/mod.rs @@ -1,4 +1,6 @@ //! Model adapters whose behavior depends on TinyAgents prompt dialects. +#[cfg(feature = "claude-code")] pub mod claude_agent_sdk; +#[cfg(feature = "claude-code")] pub mod claude_code; diff --git a/crates/tinyagents-harness/src/relaxed_json.rs b/crates/tinyagents-harness/src/relaxed_json.rs new file mode 100644 index 00000000..205a3c93 --- /dev/null +++ b/crates/tinyagents-harness/src/relaxed_json.rs @@ -0,0 +1,538 @@ +//! Best-effort repair of the relaxed / malformed JSON small local models emit +//! for tool-call arguments, turning it back into strict JSON. +//! +//! ## Why this exists +//! +//! Some OpenAI-compatible gateways fail to detokenize a model's native +//! tool-call template cleanly, so the argument blob placed in +//! `function.arguments` is frequently *not strict JSON*: +//! +//! - **unquoted object keys** — `{tool:"X",arguments:{guild_id:"Y"}}` +//! - **redundant wrapping braces** — `{{tool:"X",arguments:{…}}}`, which the +//! model piles on (`{{{…}}}`, `{{{{…}}}}`) each time the previous attempt +//! bounced back as an error. +//! - **leaked chat-template quote tokens** — the gateway emits the model's +//! string-delimiter token as literal text instead of a `"`, so a value +//! arrives as `[<|">discord<|">]` rather than `["discord"]` (observed with +//! Kimi-family models served via GMI). +//! +//! Strict `serde_json::from_str` rejects all of these, so the call is marked +//! [`tinyinference_llm::tool::ToolCall::invalid`] and fed back to the model, which +//! "repairs" it by adding *another* brace — an infinite retry that burns the +//! step budget without ever executing the tool. A zero-argument call +//! (`NAME{}`) is the only shape that survives, because `{}` is valid strict +//! JSON. +//! +//! ## What it does +//! +//! Conservative, **meaning-preserving** repairs, composed and retried at each +//! brace depth: +//! +//! 0. substitute any leaked chat-template quote token (see +//! [`LEAKED_QUOTE_TOKENS`]) back to a literal `"`, once up front, +//! 1. peel a redundant outer brace layer that wraps exactly one object +//! (`{{…}}` → `{…}`), and +//! 2. quote bare identifier keys in object position (`{tool:…}` → +//! `{"tool":…}`), string- and array-aware so string contents and +//! array/value positions are never rewritten. +//! +//! The result is accepted **only** when it parses strictly *and* is a JSON +//! object, so a scalar scraped out of noise can never masquerade as arguments. +//! This is called only *after* strict parsing has already failed on the input +//! ([`super::convert::recover_tool_arguments`]), so a well-formed argument +//! object can never reach — or be rewritten by — this path. + +use serde_json::Value; + +/// Maximum redundant outer brace layers to peel. Bounds work on adversarial +/// `{{{{…}}}}` blobs while comfortably covering every depth seen in the wild +/// (≤5 layers before the model gives up). +const MAX_BRACE_PEEL: usize = 16; + +/// Chat-template string-delimiter tokens some gateways emit as literal text in +/// place of a `"` when they fail to detokenize a model's tool-call template +/// (seen with Kimi-family models via GMI: `[<|">discord<|">]`). Both the +/// asymmetric (`<|">`) and symmetric (`<|"|>`) renderings are covered; longer +/// forms are listed first so a substitution never leaves a partial token behind. +/// Substituted to `"`, not deleted — unlike the structural markers stripped in +/// `convert::TOOL_CALL_TEMPLATE_MARKERS`. +const LEAKED_QUOTE_TOKENS: &[&str] = &["<|\"|>", "<|\">"]; + +/// Attempts to recover a strict-JSON **object** from a relaxed/malformed +/// tool-call argument string, or `None` when no conservative repair yields a +/// strictly-parseable object. +/// +/// See the module docs for the repair strategy and the safety invariant (only +/// invoked after strict parsing has already failed). +pub fn recover_relaxed_object(raw: &str) -> Option { + let normalized = normalize_leaked_quote_tokens(raw); + let mut layer = normalized.trim().to_string(); + for _ in 0..=MAX_BRACE_PEEL { + // Try the current brace layer verbatim, then with bare keys quoted. + if let Some(obj) = parse_object(&layer) { + return Some(obj); + } + let quoted = quote_bare_keys(&layer); + if quoted != layer + && let Some(obj) = parse_object("ed) + { + return Some(obj); + } + + match peel_redundant_brace(&layer) { + Some(inner) => layer = inner, + None => break, + } + } + None +} + +/// Replaces any leaked chat-template quote token (see [`LEAKED_QUOTE_TOKENS`]) +/// with a literal `"`. Returns the input unchanged when no token is present, so +/// well-formed input is untouched. +fn normalize_leaked_quote_tokens(raw: &str) -> String { + let mut out = raw.to_string(); + for &token in LEAKED_QUOTE_TOKENS { + if out.contains(token) { + out = out.replace(token, "\""); + } + } + out +} + +/// Strictly parses `s`, returning it only when it is a JSON object. +fn parse_object(s: &str) -> Option { + match serde_json::from_str::(s) { + Ok(value @ Value::Object(_)) => Some(value), + _ => None, + } +} + +/// If `s` is `{ X }` where `X` is itself exactly one complete `{…}` object +/// (ignoring surrounding whitespace), returns `X` — removing one redundant +/// wrapping brace layer. +/// +/// Returns `None` when the outer braces are *not* redundant, so a legitimate +/// single-object argument is never unwrapped. This is safe because a bare +/// object nested directly inside another object with no key (`{{…}}`) is never +/// valid JSON, so peeling it can only ever move toward a valid parse. +fn peel_redundant_brace(s: &str) -> Option { + let trimmed = s.trim(); + let inner = trimmed.strip_prefix('{')?.strip_suffix('}')?.trim(); + // The inner content must itself be a single complete object; otherwise the + // outer braces are structural (real arguments), not redundant wrapping. + if inner.starts_with('{') && object_spans_all(inner) { + Some(inner.to_string()) + } else { + None + } +} + +/// True when `s` begins with `{` and the brace it opens closes exactly at the +/// end of `s` (string-aware) — i.e. `s` is a single `{…}` object with no +/// trailing content. Used to decide whether an outer brace layer is redundant. +fn object_spans_all(s: &str) -> bool { + if !s.starts_with('{') { + return false; + } + let mut depth = 0usize; + let mut in_string = false; + let mut escaped = false; + for (idx, ch) in s.char_indices() { + if in_string { + if escaped { + escaped = false; + } else if ch == '\\' { + escaped = true; + } else if ch == '"' { + in_string = false; + } + continue; + } + match ch { + '"' => in_string = true, + '{' => depth += 1, + '}' => { + // Guard against an unbalanced stray `}` underflowing. + depth = match depth.checked_sub(1) { + Some(d) => d, + None => return false, + }; + if depth == 0 { + // Matched the opening brace: redundant only if it is the last char. + return idx + ch.len_utf8() == s.len(); + } + } + _ => {} + } + } + false +} + +/// Whether `s` is inside a JSON object or array — governs when a `,` introduces +/// a new key (object) versus a new element (array). +#[derive(Clone, Copy, PartialEq, Eq)] +enum Container { + Object, + Array, +} + +/// Quotes bare identifier keys that appear in object-key position, e.g. +/// `{tool:1,a:{b:2}}` → `{"tool":1,"a":{"b":2}}`. +/// +/// String-literal and array aware: content inside `"…"` is never touched, and +/// identifiers in array or value position are left alone (so `["discord"]`, +/// `true`, numbers, and already-quoted keys pass through unchanged). Returns the +/// input verbatim when there is nothing to quote. +/// Reads a quote-delimited object key whose delimiters may be single quotes or +/// mismatched, returning the key text and the bytes consumed (including both +/// delimiters). +/// +/// `rest` begins at the opening quote. Models that lose track of their own +/// string delimiters produce `'city'`, `"city'`, and `'city"` interchangeably — +/// all three mean the same key, and strict JSON accepts none of them. +/// +/// Returns `None` for a well-formed `"key"` so the caller keeps using the +/// normal in-string path, and `None` for anything that does not look like a +/// key: the token must be terminated by `'` or `"` followed (after optional +/// whitespace) by a `:`, and must not span a line break or contain structural +/// JSON characters. That keeps a legitimate double-quoted key containing an +/// apostrophe (`{"it's fine": 1}`) from being truncated at the apostrophe, +/// because there the next character after `'` is not a colon. +fn take_quoted_key(rest: &str) -> Option<(String, usize)> { + let mut chars = rest.char_indices(); + let (_, open) = chars.next()?; + debug_assert!(open == '"' || open == '\''); + + let mut key = String::new(); + for (idx, ch) in chars { + match ch { + '"' | '\'' => { + let after = &rest[idx + ch.len_utf8()..]; + if after.trim_start().starts_with(':') { + // A perfectly well-formed key needs no rewriting; let the + // ordinary scanner handle it so behaviour is unchanged. + if open == '"' && ch == '"' { + return None; + } + return Some((key, idx + ch.len_utf8())); + } + // Not the end of a key — record it and keep looking. + key.push(ch); + } + // A key never spans a newline or contains structure; bail out and + // let the ordinary scanner deal with whatever this really is. + '\n' | '\r' | '{' | '}' | '[' | ']' | ':' => return None, + _ => key.push(ch), + } + } + None +} + +fn quote_bare_keys(s: &str) -> String { + let mut out = String::with_capacity(s.len() + 8); + let mut stack: Vec = Vec::new(); + let mut expect_key = false; + let mut in_string = false; + let mut escaped = false; + let mut chars = s.char_indices().peekable(); + + while let Some((idx, ch)) = chars.next() { + if in_string { + out.push(ch); + if escaped { + escaped = false; + } else if ch == '\\' { + escaped = true; + } else if ch == '"' { + in_string = false; + } + continue; + } + + match ch { + // A quote in key position may open a *mismatched* key delimiter + // (`"city'`) or a single-quoted one (`'city'`), neither of which the + // in-string scanner below can terminate correctly. Try that first; + // a well-formed `"key"` falls through to the normal path. + '"' | '\'' if expect_key && matches!(stack.last(), Some(Container::Object)) => { + match take_quoted_key(&s[idx..]) { + Some((key, consumed)) => { + out.push('"'); + out.push_str(&key.replace('\\', r"\\").replace('"', "\\\"")); + out.push('"'); + // Advance the iterator past the bytes just consumed. + while chars.peek().is_some_and(|&(next, _)| next < idx + consumed) { + chars.next(); + } + expect_key = false; + } + None => { + in_string = true; + expect_key = false; + out.push(ch); + } + } + } + '"' => { + in_string = true; + expect_key = false; + out.push(ch); + } + '{' => { + stack.push(Container::Object); + expect_key = true; + out.push(ch); + } + '}' => { + stack.pop(); + expect_key = false; + out.push(ch); + } + '[' => { + stack.push(Container::Array); + expect_key = false; + out.push(ch); + } + ']' => { + stack.pop(); + expect_key = false; + out.push(ch); + } + ',' => { + // A comma re-opens key position only inside an object. + expect_key = matches!(stack.last(), Some(Container::Object)); + out.push(ch); + } + ':' => { + expect_key = false; + out.push(ch); + } + c if c.is_whitespace() => out.push(ch), + c if expect_key + && matches!(stack.last(), Some(Container::Object)) + && (c.is_ascii_alphabetic() || c == '_') => + { + // Bare identifier key: consume it and wrap it in quotes. + let start = idx; + let mut end = idx + c.len_utf8(); + while let Some(&(next_idx, next_ch)) = chars.peek() { + if next_ch.is_ascii_alphanumeric() + || next_ch == '_' + || next_ch == '-' + || next_ch == '.' + { + end = next_idx + next_ch.len_utf8(); + chars.next(); + } else { + break; + } + } + out.push('"'); + out.push_str(&s[start..end]); + out.push('"'); + expect_key = false; + } + _ => { + expect_key = false; + out.push(ch); + } + } + } + out +} + +/// Tests for the relaxed-JSON repair pipeline: leaked quote-token +/// substitution, redundant brace peeling, bare-key quoting (including +/// single/mismatched-quoted keys), and end-to-end recovery of real malformed +/// tool-call payloads captured from local models. +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn repairs_single_quoted_and_mismatched_keys() { + // Captured from `llama3.2:3b` via Ollama: the model loses track of its + // own string delimiters mid-object. + assert_eq!( + recover_relaxed_object(r#"{"name":"get_weather","parameters':{'city':"Paris"}}"#), + Some(json!({ "name": "get_weather", "parameters": { "city": "Paris" } })) + ); + // Single-quoted keys are repaired the same way, as long as the values + // themselves are well-formed. + assert_eq!( + recover_relaxed_object(r#"{'city':"Paris"}"#), + Some(json!({ "city": "Paris" })) + ); + } + + /// Single-quoted *values* are deliberately **not** repaired. + /// + /// A key is a short identifier, so reading `'` as a delimiter there is + /// safe. A value is free text where an apostrophe is ordinary English + /// (`"it's sunny"`), and treating those as delimiters would corrupt real + /// arguments. Such a blob stays unrecovered, the call is marked invalid, + /// and the agent loop hands the model a precise error to retry against — + /// the same path every other unrepairable blob takes. + #[test] + fn single_quoted_values_are_left_unrepaired() { + assert_eq!(recover_relaxed_object(r#"{'city':'Paris'}"#), None); + } + + #[test] + fn an_apostrophe_inside_a_well_formed_key_is_not_a_delimiter() { + // `'` here is followed by ` fine"`, not a colon, so the key survives + // whole rather than being truncated at the apostrophe. + assert_eq!( + recover_relaxed_object(r#"{"it's fine":1,bare:2}"#), + Some(json!({ "it's fine": 1, "bare": 2 })) + ); + } + + #[test] + fn quotes_unquoted_keys() { + assert_eq!( + recover_relaxed_object(r#"{toolkits:["discord"]}"#), + Some(json!({ "toolkits": ["discord"] })) + ); + } + + #[test] + fn quotes_multiple_unquoted_keys_and_bool_value() { + assert_eq!( + recover_relaxed_object(r#"{include_unconnected:true,toolkits:["discord"]}"#), + Some(json!({ "include_unconnected": true, "toolkits": ["discord"] })) + ); + } + + #[test] + fn substitutes_leaked_quote_tokens_in_values() { + assert_eq!( + recover_relaxed_object(r#"{toolkits:[<|">discord<|">]}"#), + Some(json!({ "toolkits": ["discord"] })) + ); + } + + #[test] + fn substitutes_symmetric_leaked_quote_token_variant() { + assert_eq!( + recover_relaxed_object(r#"{toolkits:[<|"|>discord<|"|>]}"#), + Some(json!({ "toolkits": ["discord"] })) + ); + } + + #[test] + fn peels_one_redundant_brace_layer() { + assert_eq!( + recover_relaxed_object(r#"{{"tool":"X","arguments":{"guild_id":"1"}}}"#), + Some(json!({ "tool": "X", "arguments": { "guild_id": "1" } })) + ); + } + + #[test] + fn peels_and_quotes_together() { + assert_eq!( + recover_relaxed_object( + r#"{{tool:"DISCORD_LIST_CHANNELS",arguments:{"guild_id":"1470856511193616498"}}}"# + ), + Some(json!({ + "tool": "DISCORD_LIST_CHANNELS", + "arguments": { "guild_id": "1470856511193616498" } + })) + ); + } + + #[test] + fn recovers_full_composio_execute_with_leaked_quote_tokens() { + assert_eq!( + recover_relaxed_object( + r#"{arguments:{guild_id:<|">1470856511193616498<|">},tool:<|">DISCORD_GET_GUILD_CHANNELS<|">}"# + ), + Some(json!({ + "arguments": { "guild_id": "1470856511193616498" }, + "tool": "DISCORD_GET_GUILD_CHANNELS" + })) + ); + } + + #[test] + fn peels_several_redundant_layers() { + assert_eq!( + recover_relaxed_object(r#"{{{{tool:"X",arguments:{"guild_id":"1"}}}}}"#), + Some(json!({ "tool": "X", "arguments": { "guild_id": "1" } })) + ); + } + + #[test] + fn handles_reordered_relaxed_keys() { + assert_eq!( + recover_relaxed_object(r#"{{arguments:{guild_id:"1"},tool:"X"}}"#), + Some(json!({ "arguments": { "guild_id": "1" }, "tool": "X" })) + ); + } + + #[test] + fn preserves_brace_inside_string_value() { + assert_eq!( + recover_relaxed_object(r#"{{note:"see {ref:1}"}}"#), + Some(json!({ "note": "see {ref:1}" })) + ); + } + + #[test] + fn does_not_quote_array_elements() { + assert_eq!(recover_relaxed_object(r#"{tags:[hi,bye]}"#), None); + } + + #[test] + fn rejects_keyless_nested_object() { + assert_eq!(recover_relaxed_object(r#"{tool:"X",{guild_id:"Y"}}"#), None); + } + + #[test] + fn rejects_non_object_scalar() { + assert_eq!(recover_relaxed_object("42"), None); + assert_eq!(recover_relaxed_object(r#""just a string""#), None); + assert_eq!(recover_relaxed_object("[1,2,3]"), None); + } + + #[test] + fn rejects_unrecoverable_garbage() { + assert_eq!(recover_relaxed_object(r#"{"a":1]"#), None); + assert_eq!(recover_relaxed_object("not json at all"), None); + } + + #[test] + fn already_valid_object_passes_through() { + assert_eq!( + recover_relaxed_object(r#"{"a":1,"b":{"c":2}}"#), + Some(json!({ "a": 1, "b": { "c": 2 } })) + ); + } + + #[test] + fn does_not_unwrap_legitimate_single_object() { + assert_eq!( + recover_relaxed_object(r#"{guild_id:"1",limit:50}"#), + Some(json!({ "guild_id": "1", "limit": 50 })) + ); + } + + #[test] + fn quote_bare_keys_leaves_quoted_keys_untouched() { + assert_eq!(quote_bare_keys(r#"{"a":1,"b":2}"#), r#"{"a":1,"b":2}"#); + } + + #[test] + fn normalize_leaked_quote_tokens_is_noop_without_tokens() { + assert_eq!(normalize_leaked_quote_tokens(r#"{"a":1}"#), r#"{"a":1}"#); + } + + #[test] + fn object_spans_all_respects_strings_and_trailing() { + assert!(object_spans_all(r#"{"a":"}"}"#)); + assert!(!object_spans_all(r#"{"a":1},{"b":2}"#)); + assert!(!object_spans_all(r#"{"a":1}trailing"#)); + } +} diff --git a/crates/tinyagents-harness/src/retry/mod.rs b/crates/tinyagents-harness/src/retry/mod.rs index 090c617e..d93b9a8d 100644 --- a/crates/tinyagents-harness/src/retry/mod.rs +++ b/crates/tinyagents-harness/src/retry/mod.rs @@ -178,7 +178,7 @@ impl RetryPolicy { /// Shared sleep body: logs the decision, then waits when enabled. async fn sleep_for(&self, attempt: usize, backoff: Duration, hint: Option) { if !self.backoff_sleep { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::retry", attempt, backoff_ms = backoff.as_millis() as u64, @@ -187,7 +187,7 @@ impl RetryPolicy { return; } if backoff > Duration::ZERO { - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::retry", attempt, backoff_ms = backoff.as_millis() as u64, @@ -277,7 +277,11 @@ impl RetryPolicy { /// additive: LangGraph adds `uniform(0, 1)` seconds, LangChain applies /// `delay ± 25%` clamped at zero. pub fn backoff_for_attempt_with(&self, attempt: usize, rand01: f64) -> Duration { - let base = (self.initial_backoff_ms as f64) * self.multiplier.powi(attempt as i32); + // `powi` wants `i32`; an `attempt` this large would already dwarf any + // realistic `max_attempts`, so saturate rather than truncate/wrap + // silently (M-13). + let exponent = i32::try_from(attempt).unwrap_or(i32::MAX); + let base = (self.initial_backoff_ms as f64) * self.multiplier.powi(exponent); let jittered = if self.jitter { // Map [0, 1) onto [-1, 1) then scale by the band width. let offset = JITTER_FRACTION * (2.0 * rand01.clamp(0.0, 1.0) - 1.0); @@ -286,7 +290,20 @@ impl RetryPolicy { base }; let capped = jittered.min(self.max_backoff_ms as f64); - Duration::from_millis(capped as u64) + Duration::from_millis(saturating_millis(capped)) + } +} + +/// Converts a millisecond duration held as `f64` to `u64`, saturating a +/// negative or non-finite value to `0` instead of relying on the cast's +/// implicit (if well-defined since Rust 1.45) saturating behavior — the +/// saturation is now spelled out at the call site rather than implicit in a +/// bare `as` cast (M-13). +fn saturating_millis(value: f64) -> u64 { + if value.is_finite() && value > 0.0 { + value as u64 + } else { + 0 } } @@ -333,6 +350,7 @@ pub fn retry_after_hint(error: &TinyAgentsError) -> Option { /// | `Provider` | depends | Classified from [`tinyinference_llm::model::ProviderError::retryable`] — a 429/408/409/5xx is retryable, a 4xx like 401/400 is not. | /// | `Model` | depends | No structured `ProviderError` to read, so the message text is run through [`classify_provider_failure`] — a 5xx / 429 / timeout is retryable, an `invalid api key` or `model not found` is not. | /// | `Tool` | yes | Tool execution may have hit a transient dependency. | +/// | `CallTimeout` | **yes** | A per-call ceiling fired with run time still left; unlike `Timeout`, the run is not out of budget. | /// | `Validation` | **no** | Caller-side schema or policy error; retrying will not help. | /// | `Serialization` | **no** | Malformed data; retrying will not help. | /// | `RecursionLimit` | **no** | Structural loop cap; not transient. | @@ -361,6 +379,18 @@ pub fn is_retryable(err: &TinyAgentsError) -> bool { // guessing. Callers that know better narrow this with // [`RetryPolicy::retry_on`]. TinyAgentsError::Tool(_) => true, + // A1/A3's unified retry vocabulary: `ModelRetry` is explicitly the + // recoverable half (ask the model to try again), `ToolFailed` the + // permanent half. Unlike the generic `Tool(_)` classification above, + // these two carry an explicit author intent rather than arbitrary + // caller-authored text, so retryability follows the variant directly + // instead of guessing. + TinyAgentsError::ModelRetry(_) => true, + TinyAgentsError::ToolFailed(_) => false, + // A per-model-call ceiling firing means this one call wedged, with + // run time still left — retryable, unlike a run-deadline `Timeout` + // (see that variant's own retryability rationale above). + TinyAgentsError::CallTimeout(_) => true, _ => false, } } diff --git a/crates/tinyagents-harness/src/run_queue/README.md b/crates/tinyagents-harness/src/run_queue/README.md index 6e509eca..b3954c77 100644 --- a/crates/tinyagents-harness/src/run_queue/README.md +++ b/crates/tinyagents-harness/src/run_queue/README.md @@ -12,16 +12,32 @@ only the reusable FIFO mechanics for the three lanes an agent runtime can drain at safe iteration boundaries — it has no opinion on what `T` is or when a lane should be drained. +## On the agent loop path (A4) + +A `RunQueueHandle` (`Arc>`) attached with +`RunContext::with_run_queue` is drained by the built-in agent loop: `Steer` +after each tool batch and at a natural finish, `Followup` at a natural finish +when no steer is pending (running one more turn), `Collect` once at run end +onto `AgentRun::collected`. `RunPolicy::queue_mode` (`All` | `OneAtATime`) +sets how many items a boundary takes; each application emits +`AgentEvent::QueuedMessageApplied { lane, count }`. Middleware stops, limit +stops, pauses, and deferrals leave the queue untouched. Details: +`docs/modules/harness/runtime.md`, "Queued steering and follow-ups". + ## Public surface - [`RunQueue`] — the queue itself. `new`/`Default` create an empty queue; `push(lane, item)` appends; `drain(lane)` empties one lane in FIFO order and - returns its contents; `status()` snapshots per-lane depth; `clear()` empties - every lane and returns how many items were dropped. + returns its contents; `take(lane, mode)` is the `QueueMode`-aware form; + `status()` snapshots per-lane depth; `clear()` empties every lane and + returns how many items were dropped. - [`QueueLane`] — which lane an item belongs to: `Steer` (inject at the next safe boundary as an instruction), `Followup` (dispatch as a fresh turn once - the active run completes), `Collect` (inject at the next safe boundary as - additional context). + the active run completes), `Collect` (handed back to the host at run end as + collected context, never injected). +- [`QueueMode`] — how many items `take` (and the loop) consume per call: + `OneAtATime` or `All` (default). +- [`RunQueueHandle`] — `Arc>`, the loop-consumable form. - [`QueueStatus`] — a `Serialize`-able snapshot of per-lane and total pending counts. @@ -30,7 +46,7 @@ a lane should be drained. | File | Role | | --- | --- | | `mod.rs` | `RunQueue` and its private `RunQueueInner` storage. | -| `types.rs` | `QueueLane`, `QueueStatus`. | +| `types.rs` | `QueueLane`, `QueueMode`, `QueueStatus`, `RunQueueHandle`. | | `test.rs` | Per-lane push/drain ordering, status snapshots, `clear`, and lane independence. | ## Operational constraints diff --git a/crates/tinyagents-harness/src/run_queue/mod.rs b/crates/tinyagents-harness/src/run_queue/mod.rs index 7d871849..1c1e300d 100644 --- a/crates/tinyagents-harness/src/run_queue/mod.rs +++ b/crates/tinyagents-harness/src/run_queue/mod.rs @@ -4,12 +4,39 @@ //! the queued payload. TinyAgents owns the reusable FIFO mechanics for the //! three lanes an agent runtime can consume at safe iteration boundaries: //! immediate steering, deferred follow-up work, and collected context. +//! +//! # On the agent loop path (A4) +//! +//! Attach a [`RunQueueHandle`] (an `Arc>`) to a run with +//! [`crate::context::RunContext::with_run_queue`] and the built-in +//! [`crate::agent_loop`] drains it at its safe turn boundaries, taking +//! [`QueueMode::All`] or [`QueueMode::OneAtATime`] items per boundary as +//! [`crate::runtime::RunPolicy::queue_mode`] says: +//! +//! - [`QueueLane::Steer`] — appended to the transcript right after a tool +//! batch's results (never mid-batch), and at a natural finish before any +//! follow-up. A steer that arrives after the model's final answer still +//! gets one more turn. +//! - [`QueueLane::Followup`] — appended only when the model has finished and +//! no steer is pending; the loop runs another turn instead of returning. +//! - [`QueueLane::Collect`] — never enters the transcript; drained once at +//! run end onto [`crate::middleware::AgentRun::collected`]. +//! +//! A middleware stop, limit stop, pause, or deferral is terminal: whatever is +//! still queued stays queued for the host. Every application emits +//! [`crate::events::AgentEvent::QueuedMessageApplied`]. The existing +//! [`crate::steering::SteeringHandle`] control channel (pause/resume/cancel/ +//! inject) is unchanged and independent — `RunQueue` is content injection, +//! not run control. +//! +//! `RunQueue` itself stays generic: hosts may keep using it with any `T` +//! for their own bookkeeping; only a `RunQueue` is loop-consumable. mod types; use tokio::sync::Mutex; -pub use types::{QueueLane, QueueStatus}; +pub use types::{QueueLane, QueueMode, QueueStatus, RunQueueHandle}; /// Thread-safe FIFO queue split into steer, follow-up, and collect lanes. #[derive(Debug)] @@ -62,6 +89,28 @@ impl RunQueue { } } + /// Takes items from `lane` in FIFO order according to `mode`: the oldest + /// item only under [`QueueMode::OneAtATime`], or the whole lane under + /// [`QueueMode::All`]. Returns an empty vec when the lane is empty. + pub async fn take(&self, lane: QueueLane, mode: QueueMode) -> Vec { + match mode { + QueueMode::All => self.drain(lane).await, + QueueMode::OneAtATime => { + let mut inner = self.inner.lock().await; + let items = match lane { + QueueLane::Steer => &mut inner.steers, + QueueLane::Followup => &mut inner.followups, + QueueLane::Collect => &mut inner.collects, + }; + if items.is_empty() { + Vec::new() + } else { + vec![items.remove(0)] + } + } + } + } + /// Returns the current queue depth per lane. pub async fn status(&self) -> QueueStatus { let inner = self.inner.lock().await; diff --git a/crates/tinyagents-harness/src/run_queue/test.rs b/crates/tinyagents-harness/src/run_queue/test.rs index 71a2c94d..80b2d4c7 100644 --- a/crates/tinyagents-harness/src/run_queue/test.rs +++ b/crates/tinyagents-harness/src/run_queue/test.rs @@ -57,3 +57,39 @@ async fn clear_empties_every_lane_and_reports_the_drop_count() { assert_eq!(queue.clear().await, 3); assert_eq!(queue.status().await.total, 0); } + +#[tokio::test] +async fn take_one_at_a_time_pops_only_the_oldest_item() { + let queue = RunQueue::new(); + queue.push(QueueLane::Steer, "first").await; + queue.push(QueueLane::Steer, "second").await; + + assert_eq!( + queue.take(QueueLane::Steer, QueueMode::OneAtATime).await, + vec!["first"] + ); + assert_eq!(queue.status().await.steers, 1); + assert_eq!( + queue.take(QueueLane::Steer, QueueMode::OneAtATime).await, + vec!["second"] + ); + assert!( + queue + .take(QueueLane::Steer, QueueMode::OneAtATime) + .await + .is_empty() + ); +} + +#[tokio::test] +async fn take_all_drains_the_whole_lane() { + let queue = RunQueue::new(); + queue.push(QueueLane::Followup, "first").await; + queue.push(QueueLane::Followup, "second").await; + + assert_eq!( + queue.take(QueueLane::Followup, QueueMode::All).await, + vec!["first", "second"] + ); + assert_eq!(queue.status().await.followups, 0); +} diff --git a/crates/tinyagents-harness/src/run_queue/types.rs b/crates/tinyagents-harness/src/run_queue/types.rs index 63457efd..5cd1936c 100644 --- a/crates/tinyagents-harness/src/run_queue/types.rs +++ b/crates/tinyagents-harness/src/run_queue/types.rs @@ -1,16 +1,64 @@ //! Public types for the active-run queue. +use std::sync::Arc; + +use serde::{Deserialize, Serialize}; +use tinyinference_llm::message::Message; + +/// The shared queue the agent loop drains: a [`RunQueue`][super::RunQueue] +/// of transcript-ready [`Message`]s. +/// +/// Attach one to a run with +/// [`RunContext::with_run_queue`][crate::context::RunContext::with_run_queue] +/// and keep a clone to push into from outside the run. `Steer` and +/// `Followup` items are appended to the transcript verbatim, so push them as +/// [`Message::user`] (or [`Message::system`]) — the host chooses the role. +pub type RunQueueHandle = Arc>; + /// A queue lane consumed by the agent runtime. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] pub enum QueueLane { /// Inject at the next safe iteration boundary as an instruction. Steer, /// Dispatch as a fresh turn after the active run completes. Followup, - /// Inject at the next safe boundary as additional context. + /// Collected context handed back to the host on + /// [`AgentRun::collected`][crate::middleware::AgentRun::collected] at + /// run end; never injected into the transcript. Collect, } +impl QueueLane { + /// Returns a stable, snake_case name for this lane, suitable for logging + /// and event labels (e.g. `"followup"`). + pub fn as_str(self) -> &'static str { + match self { + QueueLane::Steer => "steer", + QueueLane::Followup => "followup", + QueueLane::Collect => "collect", + } + } +} + +/// How many queued items the agent loop takes from a lane at one safe +/// boundary. Mirrors pi's `QueueMode` (`"one-at-a-time" | "all"`). +/// +/// Set on [`crate::runtime::RunPolicy::queue_mode`]; consulted by +/// [`RunQueue::take`][crate::run_queue::RunQueue::take]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum QueueMode { + /// Apply only the oldest queued item per boundary; the rest wait for the + /// next one. Gives the model a chance to react to each instruction + /// separately. + OneAtATime, + /// Apply every item queued in the lane at the boundary, in FIFO order. + /// The default. + #[default] + All, +} + /// Snapshot of the queue depth per lane. #[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)] pub struct QueueStatus { diff --git a/crates/tinyagents-harness/src/runtime/agent.rs b/crates/tinyagents-harness/src/runtime/agent.rs index 9c22ef7e..8c84933d 100644 --- a/crates/tinyagents-harness/src/runtime/agent.rs +++ b/crates/tinyagents-harness/src/runtime/agent.rs @@ -12,7 +12,7 @@ use std::pin::Pin; use std::task::{Context, Poll}; use std::time::Duration; -use futures::{Stream, StreamExt}; +use futures::Stream; use crate::agent_loop::AgentStreamItem; use crate::context::RunContext; @@ -35,7 +35,158 @@ use super::{AgentHarness, HostInvocationBinding, InvocationRuntime}; /// substitute an unhosted or differently-hosted child harness for the /// parent's policy. pub(crate) struct HostInvocationAuthority { - pub(crate) binding: HostInvocationBinding, + pub(crate) binding: std::sync::Arc>, +} + +/// Type-erasure boundary for [`RunContext::host_authority`][crate::context::RunContext]. +/// +/// This is a hand-written alternative to `dyn Any`. `Any::downcast_ref` +/// requires the caller's own generic parameters to be provably `'static`, +/// which the generic agent loop cannot promise: it deliberately keeps +/// working with a borrowed `State`/`Ctx` on the explicit-model path (see +/// `explicit_model_paths_accept_borrowed_state`). [`type_name`][Self::type_name] +/// is callable with no `'static` bound at all (`std::any::type_name` never +/// requires one), so [`host_invocation_binding`] can use it as a fail-closed +/// guard in front of the unavoidable unsafe cast, without forcing `'static` +/// onto the whole generic loop. +/// +/// `type_name` is documented as not a guaranteed-unique identifier, so this +/// is a defensive, best-effort check rather than the same soundness +/// guarantee `TypeId` gives genuinely `'static` types. It still closes the +/// realistic C-1 repro (a hosted context read by a *different* harness): +/// distinct concrete `HostInvocationAuthority` monomorphizations +/// in this crate reliably produce distinct strings. +pub(crate) trait ErasedHostAuthority: Send + Sync { + fn type_name(&self) -> &'static str; +} + +impl ErasedHostAuthority + for HostInvocationAuthority +{ + fn type_name(&self) -> &'static str { + std::any::type_name::() + } +} + +/// Closed, non-leaking classification of a hosted invocation failure. +/// +/// A host reading [`HostedError::kind`] can distinguish "the caller cancelled +/// this" from "a configured limit was exhausted" from "the provider failed" +/// without inspecting [`HostedError::message`] (which stays a fixed, +/// sanitized string per kind — see that field's doc) or attaching a private +/// event listener to reconstruct the same information from the event stream. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum HostedErrorKind { + /// The run was cancelled before completion. + Cancelled, + /// The run exceeded a wall-clock deadline (the run's own, or a per-call + /// ceiling — see [`TinyAgentsError::Timeout`] and + /// [`TinyAgentsError::CallTimeout`]). + Timeout, + /// A configured run limit (model calls, tool calls, recursion depth, a + /// host budget) was exhausted. + LimitExceeded, + /// The host's own policy rejected the invocation (an unresolvable + /// definition, a failed security screen, an unauthorized delegate). + Policy, + /// The model provider failed the call. + Provider, + /// Any other internal failure not covered by a more specific kind. + Internal, +} + +/// The typed failure returned by the hosted entry points +/// ([`AgentHarness::invoke_agent`] and its streaming counterpart) in place of +/// a generic `TinyAgentsError::Model("hosted agent invocation failed")`. +/// +/// This intentionally does not implement `TinyAgentsError`'s "one error type" +/// convention: it is the harness's product-host boundary type, not another +/// case folded into the crate-wide error, and it is deliberately smaller — +/// `message` is a fixed, sanitized string selected by `kind` (never the +/// underlying provider/middleware/budget error text; that stays available to +/// the host only through its own capability bundle's own logging and through +/// the internal (non-hosted) event stream if it chose to attach a listener). +#[derive(Debug)] +pub struct HostedError { + /// Closed classification of the failure. See [`HostedErrorKind`]. + pub kind: HostedErrorKind, + /// Fixed, sanitized message selected by `kind` — never raw provider, + /// middleware, or budget error text. + pub message: String, + /// The accumulated transcript, usage, and executed-tool summary as far as + /// the run got before failing, when available. + pub run: Option>, +} + +impl std::fmt::Display for HostedError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.message) + } +} + +impl std::error::Error for HostedError {} + +/// Classifies a raw loop error into the closed [`HostedErrorKind`] vocabulary +/// a host is allowed to see. +fn classify_hosted_error(error: &TinyAgentsError) -> HostedErrorKind { + match error { + TinyAgentsError::Cancelled => HostedErrorKind::Cancelled, + TinyAgentsError::Timeout(_) | TinyAgentsError::CallTimeout(_) => HostedErrorKind::Timeout, + TinyAgentsError::LimitExceeded(_) | TinyAgentsError::SubAgentDepth(_) => { + HostedErrorKind::LimitExceeded + } + TinyAgentsError::Validation(_) | TinyAgentsError::Steering(_) => HostedErrorKind::Policy, + TinyAgentsError::Provider(_) | TinyAgentsError::Model(_) => HostedErrorKind::Provider, + _ => HostedErrorKind::Internal, + } +} + +/// The fixed, sanitized message for each [`HostedErrorKind`]. Never derived +/// from the underlying error's own text. +fn hosted_error_message(kind: HostedErrorKind) -> &'static str { + match kind { + HostedErrorKind::Cancelled => "hosted agent invocation was cancelled", + HostedErrorKind::Timeout => "hosted agent invocation timed out", + HostedErrorKind::LimitExceeded => "hosted agent invocation exceeded a configured limit", + HostedErrorKind::Policy => "hosted agent invocation was rejected by policy", + HostedErrorKind::Provider => "hosted agent invocation failed at the model provider", + HostedErrorKind::Internal => "hosted agent invocation failed", + } +} + +/// Builds a [`HostedError`] from the raw loop error and whatever partial +/// [`AgentRun`] the loop accumulated before failing. +fn hosted_error(error: &TinyAgentsError, run: AgentRun) -> HostedError { + let kind = classify_hosted_error(error); + HostedError { + kind, + message: hosted_error_message(kind).to_string(), + run: Some(Box::new(run)), + } +} + +/// Reconstructs a crate-wide [`TinyAgentsError`] from a [`HostedError`] for +/// internal callers (recursive hosted delegation) that must keep propagating +/// through the ordinary `Result` = `Result` surface. +/// This is a lossless-enough round trip for control flow: `Cancelled` and +/// `Timeout` map back to their own variants (so cancellation/deadline +/// semantics upstream keep working, e.g. the fallback gate in +/// `invoke_model_resolving`), and the rest become typed but message-generic +/// variants — never worse than what this boundary already returned before +/// `HostedError` existed. +impl From for TinyAgentsError { + fn from(error: HostedError) -> Self { + match error.kind { + HostedErrorKind::Cancelled => TinyAgentsError::Cancelled, + HostedErrorKind::Timeout => TinyAgentsError::Timeout(error.message), + HostedErrorKind::LimitExceeded => TinyAgentsError::LimitExceeded(error.message), + HostedErrorKind::Policy => TinyAgentsError::Validation(error.message), + HostedErrorKind::Provider | HostedErrorKind::Internal => { + TinyAgentsError::Model(error.message) + } + } + } } /// A host-owned turn request. @@ -49,6 +200,9 @@ pub struct AgentTurnRequest { pub agent_id: String, /// Initial transcript supplied by the host. pub messages: Vec, + /// Resolutions for the deferred tool calls left pending on `messages` + /// by a previous hosted turn (A2). See [`Self::with_deferred_results`]. + pub deferred_results: Option, } impl AgentTurnRequest { @@ -60,8 +214,19 @@ impl AgentTurnRequest { Self { agent_id: agent_id.into(), messages, + deferred_results: None, } } + + /// Resumes a hosted turn that stopped with `AgentRun::deferred` set + /// (A2): `messages` should be that run's transcript and `results` must + /// resolve every pending call. The hosted counterpart of + /// [`AgentHarness::resume_deferred`][crate::runtime::AgentHarness::resume_deferred]. + #[must_use] + pub fn with_deferred_results(mut self, results: crate::tool::DeferredToolResults) -> Self { + self.deferred_results = Some(results); + self + } } /// One host-authorized execution of an agent. @@ -137,23 +302,29 @@ impl AgentInvocation { /// observer even when a caller stops listening before a terminal item. The /// invocation's host authority is owned by that context, never by the harness. pub struct AgentStream<'a, State: Send + Sync + 'static, Ctx: Send + Sync> { + // Owns its inputs (the harness borrow or the invocation-local runtime is + // moved into the driving future itself, inside `invoke_stream_with_runner`) + // so nothing outside this field needs to outlive it and no lifetime + // extension is required to store it here. inner: Option + Send + 'a>>>, - // Kept after `inner` so Rust drops the borrowed stream before the overlay - // that owns its harness. See `extend_overlay_stream_lifetime`. - _runtime: Option>>, cancellation: crate::CancellationToken, terminal_observer: std::sync::Arc>>, terminal_observed: bool, - marker: std::marker::PhantomData<(&'a State, Ctx)>, + // `fn() -> Ctx` (rather than bare `Ctx`) keeps this marker `Unpin` + // regardless of `Ctx`, which is what lets `poll_next` use the safe + // `Pin::get_mut` below instead of `get_unchecked_mut`. + #[allow(clippy::type_complexity)] + marker: std::marker::PhantomData<(&'a State, fn() -> Ctx)>, } impl Stream for AgentStream<'_, State, Ctx> { type Item = AgentStreamItem; fn poll_next(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { - // `inner` is pinned independently by `Box`; this projection never moves - // the boxed stream or any other field of `AgentStream`. - let stream = unsafe { self.get_unchecked_mut() }; + // Every field is `Unpin` (`Option>>`, `CancellationToken`, + // an `Arc>`, `bool`, and a `fn()`-based `PhantomData`), so + // `AgentStream` itself is `Unpin` and this projection is safe. + let stream = self.get_mut(); match stream.inner.as_mut() { Some(inner) => match inner.as_mut().poll_next(context) { Poll::Ready(Some(item)) => { @@ -247,7 +418,7 @@ impl Drop for AgentStream<'_, St && let Some(observer) = observer.take() { observer( - AgentRun::new(), + crate::context::TerminalRunSummary::default(), false, Some("hosted stream cancelled before execution began".to_string()), ); @@ -399,7 +570,7 @@ impl AgentHarness, state: &State, - ) -> Result + ) -> std::result::Result where State: 'static, { @@ -413,10 +584,85 @@ impl AgentHarness, state: &State, - ) -> Result + ) -> std::result::Result + where + State: 'static, + { + let (runtime, context, prepared) = self + .prepare_hosted_turn(invocation) + .await + .map_err(|error| hosted_error(&error, AgentRun::new()))?; + let runner = runtime + .as_deref() + .map(InvocationRuntime::harness) + .unwrap_or(self); + + let outcome = runner + .invoke_in_context_collecting_partial(state, context, prepared.messages.clone()) + .await; + match outcome.error { + None => Ok(outcome.run), + Some(error) => Err(hosted_error(&error, outcome.run)), + } + } + + /// Collects a hosted turn through the streaming driver while preserving the + /// parent's exact capability bundle. Recursive streaming delegation uses + /// this rather than the unary entry point so model deltas and delta + /// middleware remain part of the shared parent event stream. + /// + /// Drives [`AgentHarness::invoke_streaming_in_context_collecting_partial`] + /// directly rather than going through [`AgentHarness::invoke_agent_stream`]: + /// the public stream sanitizes every item (see + /// [`sanitize_hosted_stream_item`]), which would throw away the real + /// [`TinyAgentsError`] this method needs to classify into a + /// [`HostedErrorKind`] before its own, separate sanitization. + pub(crate) async fn invoke_agent_streaming_with_capabilities( + &self, + invocation: AgentInvocation, + state: &State, + ) -> std::result::Result where + Ctx: 'static, State: 'static, { + let (runtime, context, prepared) = self + .prepare_hosted_turn(invocation) + .await + .map_err(|error| hosted_error(&error, AgentRun::new()))?; + let runner = runtime + .as_deref() + .map(InvocationRuntime::harness) + .unwrap_or(self); + + let outcome = runner + .invoke_streaming_in_context_collecting_partial( + state, + context, + prepared.messages.clone(), + ) + .await; + match outcome.error { + None => Ok(outcome.run), + Some(error) => Err(hosted_error(&error, outcome.run)), + } + } + + /// Shared setup for both hosted drivers: resolves and authorizes the + /// turn, installs the host authority and terminal observer on `context`, + /// and emits [`ProgressEvent::Started`]. Returns the invocation-local + /// runtime overlay (if any — the caller derives the harness that should + /// actually run the turn from it, since a reference borrowed from it here + /// cannot outlive this function), the prepared `context`, and the + /// prepared turn. + async fn prepare_hosted_turn( + &self, + invocation: AgentInvocation, + ) -> Result<( + Option>>, + RunContext, + PreparedAgentTurn, + )> { let AgentInvocation { host, request, @@ -427,6 +673,9 @@ impl AgentHarness AgentHarness( @@ -445,51 +694,7 @@ impl AgentHarness Ok(outcome.run), - Some(TinyAgentsError::Cancelled) => Err(TinyAgentsError::Cancelled), - Some(TinyAgentsError::Timeout(message)) => Err(TinyAgentsError::Timeout(message)), - Some(_) => Err(TinyAgentsError::Model( - "hosted agent invocation failed".to_string(), - )), - } - } - - /// Collects a hosted turn through the streaming driver while preserving the - /// parent's exact capability bundle. Recursive streaming delegation uses - /// this rather than the unary entry point so model deltas and delta - /// middleware remain part of the shared parent event stream. - pub(crate) async fn invoke_agent_streaming_with_capabilities( - &self, - invocation: AgentInvocation, - state: &State, - ) -> Result - where - Ctx: 'static, - State: 'static, - { - let stream = self - .invoke_agent_stream_with_capabilities(invocation, state) - .await?; - futures::pin_mut!(stream); - while let Some(item) = stream.next().await { - match item { - AgentStreamItem::Completed(run) => return Ok(*run), - AgentStreamItem::Failed { .. } => { - return Err(TinyAgentsError::Model( - "hosted agent invocation failed".to_string(), - )); - } - AgentStreamItem::Event(_) => {} - } - } - Err(TinyAgentsError::Model( - "hosted stream ended without a terminal result".to_string(), - )) + Ok((runtime, context, prepared)) } /// Starts a hosted streaming turn. @@ -530,14 +735,23 @@ impl AgentHarness AgentHarness crate::agent_loop::StreamRunner::Owned(runtime), + None => crate::agent_loop::StreamRunner::Borrowed(self), + }; + let stream = crate::agent_loop::invoke_stream_with_runner( + stream_runner, + state, + context, + prepared.messages.clone(), + ); Ok(AgentStream { - // `runtime` is retained by this stream and is declared after - // `inner`, so it outlives the stream's borrow of its harness. The - // explicit helper records that otherwise non-obvious lifetime - // relationship at the one boundary where the owned hosted - // invocation meets the borrowed stream API. - inner: Some(unsafe { extend_overlay_stream_lifetime(Box::pin(stream)) }), - _runtime: runtime, + inner: Some(Box::pin(stream)), cancellation, terminal_observer, terminal_observed: false, @@ -572,30 +793,28 @@ impl AgentHarness>, request: AgentTurnRequest, context: &RunContext, ) -> Result> { - let cancellation = context.cancellation.clone(); let preparation = self.prepare_agent_turn(host, request, context); - let outcome = match self.host_io_budget(context) { - Some(remaining) => tokio::select! { - biased; - _ = cancellation.cancelled() => Err(TinyAgentsError::Cancelled), - result = tokio::time::timeout(remaining, preparation) => result.map_err(|_| TinyAgentsError::Timeout(format!( + context + .bounded(self.host_io_budget(context), preparation, || { + format!( "host turn preparation for run `{}` exceeded its remaining wall-clock budget", context.run_id() - )))?, - }, - None => tokio::select! { - biased; - _ = cancellation.cancelled() => Err(TinyAgentsError::Cancelled), - result = preparation => result, - }, - }; - outcome.map_err(sanitize_hosted_preparation_error) + ) + }) + .await } /// The wall-clock time left for host I/O (definition lookup, security @@ -645,9 +864,15 @@ impl AgentHarness AgentHarness`), so both collapse to + // `None` here — `resolve_tool_allowlist` in `agent_loop::tools` + // treats that as fail-closed by default (I-9), not as "unrestricted". + let declared_tools: std::collections::HashSet = + definition.tools.into_iter().collect(); + let allowed_tools = (!declared_tools.is_empty()).then_some(declared_tools); Ok(PreparedAgentTurn { binding: HostInvocationBinding { host: host.clone(), agent_id: request.agent_id.clone(), model_pin: definition.model, role: definition.role, - allowed_tools: definition.tools.into_iter().collect(), + allowed_tools, progress: progress.clone(), runtime: None, }, @@ -754,7 +986,7 @@ impl AgentHarness AgentHarness( - stream: Pin + Send + '_>>, -) -> Pin + Send + 'a>> { - // SAFETY: documented above; the owning Arc is retained by AgentStream. - unsafe { std::mem::transmute(stream) } -} - /// Returns this live context's host authorization, if it is a hosted run. /// /// The binding is carried by the non-serializable context rather than the /// reusable harness, so concurrent roots have no shared mutable authority. +/// +/// `host_authority` is `Option>` and is installed +/// only by the hosted entry points in this module, which require +/// `State: 'static, Ctx: 'static` and store exactly +/// `HostInvocationAuthority`. Nothing about [`RunContext`] +/// prevents a caller from handing a hosted context to a *different* harness +/// (a different `State`, or — via [`RunContext::child_with_data`] changing +/// `Ctx`), so the erased type is checked with [`Any::downcast_ref`] rather +/// than assumed. A mismatch fails closed with +/// [`TinyAgentsError::Validation`] instead of reinterpreting memory through +/// the wrong type. Absence of any authority is the ordinary, cheap case (an +/// explicit-model run, or the generic loop when no hosted invocation +/// installed one) and returns `Ok(None)` without touching `Any` at all, so +/// this function itself still only needs `State: 'static, Ctx: 'static` on +/// the (rare) hosted path — its callers already carry that bound. pub(crate) fn host_invocation_binding( context: &RunContext, -) -> Result>> { +) -> Result>>> { let Some(authority) = context.host_authority.as_ref() else { return Ok(None); }; - // `host_authority` is crate-private and is installed only by the hosted - // entry points, which require `State: 'static` and store exactly - // `HostInvocationAuthority`. Explicit-model entry points never - // install it, so they return at the `None` branch without requiring - // `State: 'static` or consulting `Any` at all. Keeping this cast at the - // private hosted-context boundary restores borrowed-state support to the - // generic loop without creating a harness registry or any cross-invocation - // authority channel. - // - // SAFETY: no public API can construct or mutate `host_authority`; its only - // assignment is the hosted `AgentInvocation` path in this module. - // `RunContext::child` clones that same `Arc` only for recursive calls with - // the same `State`. Thus a present authority always points at the concrete - // type requested here for the active harness invocation. + // `RunContext::child` (the only authority-propagating path) requires the + // same `Ctx` as its parent, and `RunContext::child_with_data` (the only + // path that changes `Ctx`) always clears `host_authority` first — so a + // present authority's `Ctx` already matches this call's `Ctx` by + // construction. `State` has no such structural guarantee (nothing + // prevents handing a hosted context to a *different* harness), so it is + // checked here at read time via `ErasedHostAuthority::type_name` (see + // that trait's doc comment for why this, and not `Any`, is used). + let expected = std::any::type_name::>(); + if authority.type_name() != expected { + return Err(TinyAgentsError::Validation( + "host authority type mismatch: this run context was hosted by a different \ + State/Ctx harness than the one reading it" + .to_string(), + )); + } + #[allow(unsafe_code)] + // SAFETY: `context.host_authority` is crate-private and is installed + // only by the hosted entry points in this module, which always store + // exactly `HostInvocationAuthority` for the harness they are + // called on. The `type_name` check above additionally rejects any value + // whose concrete type does not match this call's own `State`/`Ctx` + // before this cast runs, so a mismatched authority never reaches it. let authority = unsafe { &*(std::sync::Arc::as_ptr(authority) as *const HostInvocationAuthority) }; @@ -826,7 +1068,7 @@ pub(crate) fn emit_host_progress( let Ok(Some(binding)) = host_invocation_binding::(context) else { return; }; - let Some(progress) = binding.progress else { + let Some(progress) = binding.progress.as_ref() else { return; }; progress.send_nonterminal(event); @@ -840,7 +1082,9 @@ pub(crate) fn emit_host_progress( /// same generic message so internal detail never reaches a hosted caller. fn sanitize_hosted_preparation_error(error: TinyAgentsError) -> TinyAgentsError { match error { - TinyAgentsError::Cancelled | TinyAgentsError::Timeout(_) => error, + TinyAgentsError::Cancelled + | TinyAgentsError::Timeout(_) + | TinyAgentsError::CallTimeout(_) => error, _ => TinyAgentsError::Model("hosted agent invocation failed".to_string()), } } @@ -854,14 +1098,14 @@ fn sanitize_hosted_preparation_error(error: TinyAgentsError) -> TinyAgentsError /// happen even when the turn ends off the normal async call path. fn spawn_host_finalizer( prepared: PreparedAgentTurn, - run: AgentRun, + run: crate::context::TerminalRunSummary, succeeded: bool, error: Option, ) { if let Ok(handle) = tokio::runtime::Handle::try_current() { handle.spawn(async move { finish_host_turn(prepared, run, succeeded, error).await }); } else { - tinyagents_tracing::warn!( + tracing::warn!( run_id = %prepared.run_id, "[host] no Tokio runtime during terminal cleanup; starting fallback finalizer" ); @@ -885,7 +1129,7 @@ fn spawn_host_finalizer( prepared: PreparedAgentTurn, - run: AgentRun, + run: crate::context::TerminalRunSummary, succeeded: bool, error: Option, ) { @@ -904,7 +1148,7 @@ async fn finish_host_turn( }); } } - let output = run.text().unwrap_or_default(); + let output = run.text.clone().unwrap_or_default(); let mut summary = TurnSummary::new(prepared.thread_id.clone(), &prepared.binding.agent_id) .with_text(&prepared.input_text, &output) .with_usage(run.usage.usage); @@ -921,13 +1165,13 @@ async fn finish_host_turn( "turn_failure" }); if let Err(error) = memory.remember(item).await { - tinyagents_tracing::warn!(%error, "[host] memory sink failed after terminal turn"); + tracing::warn!(%error, "[host] memory sink failed after terminal turn"); } } if let Some(learning) = &prepared.binding.host.learning && let Err(error) = learning.on_turn_complete(&summary).await { - tinyagents_tracing::warn!(%error, "[host] learning sink failed after terminal turn"); + tracing::warn!(%error, "[host] learning sink failed after terminal turn"); } if let Some(store) = &prepared.binding.host.experience { let mut experience = @@ -936,7 +1180,7 @@ async fn finish_host_turn( experience = experience.succeeded(); } if let Err(error) = store.record(&experience).await { - tinyagents_tracing::warn!(%error, "[host] experience store failed after terminal turn"); + tracing::warn!(%error, "[host] experience store failed after terminal turn"); } } } diff --git a/crates/tinyagents-harness/src/runtime/mod.rs b/crates/tinyagents-harness/src/runtime/mod.rs index 48cfc12f..02df06fd 100644 --- a/crates/tinyagents-harness/src/runtime/mod.rs +++ b/crates/tinyagents-harness/src/runtime/mod.rs @@ -29,8 +29,10 @@ mod agent; mod types; -pub use agent::{AgentInvocation, AgentStream, AgentTurnRequest}; -pub(crate) use agent::{emit_host_progress, host_invocation_binding}; +#[cfg(test)] +pub(crate) use agent::HostInvocationAuthority; +pub use agent::{AgentInvocation, AgentStream, AgentTurnRequest, HostedError, HostedErrorKind}; +pub(crate) use agent::{ErasedHostAuthority, emit_host_progress, host_invocation_binding}; pub use types::*; use std::sync::Arc; @@ -55,6 +57,12 @@ impl AgentHarness { policy: RunPolicy::default(), tool_timeouts: None, response_cache: None, + output_validator: None, + deferred_tool_handler: None, + toolset: None, + capabilities: Vec::new(), + capability_base_toolset: None, + loop_driver: None, } } @@ -83,6 +91,16 @@ impl AgentHarness { self } + /// Registers a schema-only external tool the host executes out of band + /// (A2). See [`ToolRegistry::register_external`]. + pub fn register_external_tool( + &mut self, + schema: tinyinference_llm::tool::ToolSchema, + ) -> &mut Self { + self.tools.register_external(schema); + self + } + /// Registers a tool whose execution needs the typed parent run. pub fn register_tool_dispatch( &mut self, @@ -178,6 +196,172 @@ impl AgentHarness { self.response_cache.as_ref() } + /// Registers an [`crate::structured::OutputValidator`] consulted after + /// the final turn's structured extraction succeeds (A3's + /// output-validation retry loop). + /// + /// The validator sees the *already schema-valid* extracted value; a + /// `TinyAgentsError::ModelRetry` it returns is treated exactly like a + /// schema-validation failure — re-asked, bounded by + /// [`RunPolicy::output_retry`]. Only one validator may be installed; + /// calling this again replaces it. Returns `&mut Self` for chaining. + pub fn with_output_validator( + &mut self, + validator: Arc>, + ) -> &mut Self { + self.output_validator = Some(validator); + self + } + + /// Installs an inline [`crate::tool::DeferredToolHandler`] (A2). + /// + /// With a handler present, a tool batch that defers one or more calls + /// (approval-required policy, `ApprovalRequired`/`CallDeferred`, or an + /// external tool) is resolved by calling the handler right there and the + /// loop continues; the caller never sees `AgentRun::deferred`. Without + /// one, the loop exits with the pending requests for the host to resolve + /// and resume later. Only one handler may be installed; calling this + /// again replaces it. Returns `&mut Self` for chaining. + pub fn with_deferred_tool_handler( + &mut self, + handler: Arc, + ) -> &mut Self { + self.deferred_tool_handler = Some(handler); + self + } + + /// Installs a composable [`crate::tool::toolset::ToolSet`] chain (gap + /// B3) as an additional source of tools, consulted alongside + /// [`Self::tools`]. + /// + /// # What this changes + /// + /// - **Advertisement**: the agent loop's per-turn model-visible tool + /// catalogue is built by projecting this toolset's + /// [`crate::tool::toolset::ToolSet::tools`] (re-consulted every turn, + /// so a [`crate::tool::toolset::PreparedToolSet`] or + /// [`crate::tool::toolset::ApprovalRequiredToolSet`] in the chain can + /// vary what is advertised turn to turn) **in addition to** the + /// registry's own `Direct` schemas — a name the toolset does not + /// mention falls back to the registry unchanged. + /// - **Dispatch is not automatically wired to this toolset.** The agent + /// loop's admission path (`agent_loop/tools.rs`) resolves calls through + /// [`Self::tools`] only, exactly as before this field existed. A tool + /// that only the toolset chain exposes must also be reachable through + /// the registry to be *callable* (not just advertised) — bridge it + /// explicitly with + /// [`crate::tool::toolset::ToolSetDispatchBridge`] and + /// [`Self::register_tool_dispatch`]. See that bridge's doc comment for + /// why: it requires `State: 'static, Ctx: 'static`, a bound the loop's + /// generic admission path deliberately does not carry (recursive + /// sub-agent dispatch stays callable with a borrowed, non-`'static` + /// `State`/`Ctx`). + /// + /// A caller building a fresh [`crate::tool::ToolRegistry`] separately + /// (rather than through [`Self::register_tool`]) can pass it here + /// directly — [`crate::tool::ToolRegistry`] implements + /// [`crate::tool::toolset::ToolSet`] — or compose it with other + /// toolsets via [`crate::tool::toolset::CombinedToolSet`]. + /// + /// `None` (never calling this) leaves every existing harness's turn + /// behavior exactly as before this field existed. Returns `&mut Self` + /// for chaining. + pub fn with_toolset( + &mut self, + toolset: Arc>, + ) -> &mut Self { + self.toolset = Some(toolset); + self + } + + /// Returns the installed toolset chain, if any. See + /// [`Self::with_toolset`]. + pub fn toolset(&self) -> Option<&Arc>> { + self.toolset.as_ref() + } + + /// Installs a [`crate::capability::Capability`] bundle (gap G3): its + /// toolset, middleware, and model-request defaults are applied to this + /// harness, and its [`crate::capability::Capability::exposure`]/ + /// [`crate::capability::Capability::defer_loading`] settings are honored + /// by the [`crate::capability::CapabilityToolSet`] this method installs. + /// + /// May be called more than once; every installed capability accumulates + /// (see [`Self::capabilities`]) and [`Self::toolset`] is rebuilt each + /// time from the complete list, so a `defer_loading` capability's + /// [`crate::capability::LoadCapabilityTool`] always covers every deferred + /// capability installed so far, under one shared load state. + /// + /// # What this changes + /// + /// - **Toolset**: composes a fresh + /// [`crate::capability::CapabilityToolSet`] over every installed + /// capability with whatever toolset was already installed via + /// [`Self::with_toolset`] *before* the first `with_capability` call + /// (captured once, in [`Self::capability_base_toolset`]) through + /// [`crate::tool::toolset::CombinedToolSet`]. As with + /// [`Self::with_toolset`], dispatch for a capability's own tools is not + /// automatically bridged into [`Self::tools`] — bridge explicitly with + /// [`crate::tool::toolset::ToolSetDispatchBridge`] and + /// [`Self::register_tool_dispatch`] for a tool that must be callable, + /// not just advertised. The synthetic `load_capability` tool is the one + /// exception: it is registered directly into [`Self::tools`] (it needs + /// no `RunContext`/`State` to execute), so it is callable immediately. + /// - **Middleware**: each capability's middleware is appended, in + /// installation order, via [`Self::push_middleware`]. + /// - **Model defaults**: each capability's + /// [`crate::capability::ModelRequestDefaults`], if set, is applied onto + /// [`Self::policy`] via + /// [`crate::capability::ModelRequestDefaults::apply_to`] — a later + /// capability's set fields win over an earlier one's. + /// + /// Returns `&mut Self` for chaining. + pub fn with_capability( + &mut self, + capability: crate::capability::Capability, + ) -> &mut Self + where + State: 'static, + Ctx: 'static, + { + if self.capabilities.is_empty() { + self.capability_base_toolset = self.toolset.take(); + } + for middleware in capability.middleware.clone() { + self.push_middleware(middleware); + } + if let Some(defaults) = &capability.model_defaults { + defaults.apply_to(&mut self.policy); + } + self.capabilities.push(capability); + + let capability_toolset = + crate::capability::CapabilityToolSet::new(self.capabilities.clone()); + + // The `load_capability` tool needs no `RunContext`/`State` to run, so + // it is registered directly into `self.tools` — the one part of a + // capability's contribution that is callable, not just advertised, + // without a caller-supplied dispatch bridge (see the doc comment + // above). `Self::register_tool` silently replaces a prior + // registration under the same name, so re-registering on every call + // keeps it in sync with the full, still-accumulating capability list. + if let Some(load_tool) = capability_toolset.load_tool() { + self.register_tool(load_tool); + } + + let capability_toolset: Arc> = + Arc::new(capability_toolset); + self.toolset = Some(match &self.capability_base_toolset { + Some(base) => Arc::new(crate::tool::toolset::CombinedToolSet::new(vec![ + base.clone(), + capability_toolset, + ])), + None => capability_toolset, + }); + + self + } + /// Returns a reference to the model registry. pub fn models(&self) -> &ModelRegistry { &self.models @@ -197,6 +381,27 @@ impl AgentHarness { pub fn policy(&self) -> &RunPolicy { &self.policy } + + /// Installs an alternate loop engine (A5), consulted by `invoke*` when + /// [`RunPolicy::execution`] is + /// [`LoopExecution`][crate::runtime::LoopExecution]`::Graph`. See + /// [`crate::agent_loop::phases::LoopDriver`]. Returns `&mut Self` for + /// chaining. + pub fn with_loop_driver( + &mut self, + driver: Arc>, + ) -> &mut Self { + self.loop_driver = Some(driver); + self + } + + /// Returns the installed alternate loop engine, if any. See + /// [`Self::with_loop_driver`]. + pub fn loop_driver( + &self, + ) -> Option<&Arc>> { + self.loop_driver.as_ref() + } } impl Default for AgentHarness { diff --git a/crates/tinyagents-harness/src/runtime/test.rs b/crates/tinyagents-harness/src/runtime/test.rs index 917e8889..15e6409c 100644 --- a/crates/tinyagents-harness/src/runtime/test.rs +++ b/crates/tinyagents-harness/src/runtime/test.rs @@ -975,11 +975,9 @@ async fn host_driven_turn_resolves_and_composes_without_touching_explicit_sdk_de let progress = Arc::new(RecordingProgressSink::new()); let host = crate::host::HostCapabilities::new( Arc::new(StaticContextComposer::new("host system")), - Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )])), + Arc::new(InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["verbose_lookup"]), + ])), Arc::new(AllowAllSecurityGate), Arc::new(FixedModelResolver::new(model.clone())), ) @@ -1064,7 +1062,7 @@ async fn initial_host_model_resolution_is_cancelled_while_the_resolver_is_pendin } result = &mut invocation => panic!("pending resolver unexpectedly finished: {result:?}"), }; - assert!(matches!(error, crate::error::TinyAgentsError::Cancelled)); + assert_eq!(error.kind, crate::runtime::HostedErrorKind::Cancelled); } #[tokio::test] @@ -1100,11 +1098,11 @@ async fn policy_only_deadline_bounds_initial_host_resolution_with_a_timeout_erro ) .await .expect_err("policy deadline must bound a host resolver without a RunConfig timeout"); - assert!(matches!(error, crate::error::TinyAgentsError::Timeout(_))); - assert!( - error.to_string().contains("host model resolution for run `policy-host-resolve-timeout` exceeded its remaining wall-clock budget"), - "timeout must retain its host-resolution and policy-budget shape: {error}" - ); + // `HostedError` intentionally sanitizes the message to a fixed string per + // `kind` (I-6) — the detailed "exceeded its remaining wall-clock budget" + // text is still available on the run's internal `TinyAgentsError` (see + // the non-hosted equivalents of this test), just not leaked here. + assert_eq!(error.kind, crate::runtime::HostedErrorKind::Timeout); } #[tokio::test] @@ -1141,8 +1139,7 @@ async fn per_model_call_limit_bounds_initial_host_resolution() { ) .await .expect_err("per-model-call cap must bound host resolution"); - assert!(matches!(error, crate::error::TinyAgentsError::Timeout(_))); - assert!(error.to_string().contains("per-model-call ceiling")); + assert_eq!(error.kind, crate::runtime::HostedErrorKind::Timeout); } async fn assert_rebound_host_resolution_stops( @@ -1199,6 +1196,7 @@ async fn assert_rebound_host_resolution_stops( .await .expect("rebound resolver must not hang") .expect_err("the rebinding resolver remains pending") + .into() } result = &mut invocation => panic!("rebind resolver unexpectedly finished: {result:?}"), } @@ -1213,10 +1211,15 @@ async fn middleware_rebinding_cancels_a_pending_host_resolver() { #[tokio::test] async fn middleware_rebinding_applies_the_host_resolution_deadline() { + // `assert_rebound_host_resolution_stops` round-trips through the hosted + // entry point (`AgentHarness::invoke_agent`), which now classifies and + // sanitizes via `HostedError` (I-6) before converting back to + // `TinyAgentsError` for this helper's declared return type — so the + // detailed "host model resolution ... remaining wall-clock budget" text + // is intentionally no longer observable here; only the `Timeout` + // classification survives the round trip. let error = assert_rebound_host_resolution_stops(None, Some(5)).await; assert!(matches!(error, crate::error::TinyAgentsError::Timeout(_))); - assert!(error.to_string().contains("host model resolution")); - assert!(error.to_string().contains("remaining wall-clock budget")); } #[tokio::test] @@ -1452,9 +1455,10 @@ async fn hosted_turn_blocks_provider_extension_user_blocks_before_model_submissi ) .await .expect_err("blocked extensions must not reach the provider"); + assert_eq!(error.kind, crate::runtime::HostedErrorKind::Policy); assert_eq!( error.to_string(), - "model error: hosted agent invocation failed" + "hosted agent invocation was rejected by policy" ); assert!(!error.to_string().contains("secret")); assert!(model.requests().is_empty()); @@ -1520,6 +1524,107 @@ async fn hosted_model_resolution_marks_only_root_contexts_as_team_leads() { ); } +/// I-6 regression: a hosted invocation that exhausts a configured run limit +/// must classify as `HostedErrorKind::LimitExceeded`, distinguishable from +/// other hosted failure modes (here, a policy rejection) rather than every +/// non-cancel/timeout failure collapsing into one generic +/// `Model("hosted agent invocation failed")`. +#[tokio::test] +async fn hosted_limit_exceeded_is_distinguishable_from_other_hosted_errors() { + // A model that always requests the same tool call, so the run never + // finishes on its own and must hit `max_model_calls`. + let looping_model = Arc::new(ScriptedModel::new( + std::iter::repeat_with(|| { + let mut response = ModelResponse::assistant(""); + response + .message + .tool_calls + .push(tinyinference_llm::tool::ToolCall::new( + "call", + "noop", + json!({}), + )); + response + }) + .take(8) + .collect(), + )); + let definition = AgentDefinition::new("helper", "Helper", "test helper").with_tools(["noop"]); + let host = crate::host::HostCapabilities::new( + Arc::new(StaticContextComposer::empty()), + Arc::new(InMemoryDefinitionRegistry::new(vec![definition])), + Arc::new(AllowAllSecurityGate), + Arc::new(FixedModelResolver::new(looping_model)), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_tool(Arc::new(NoopTool)); + + let limit_error = harness + .invoke_agent( + AgentInvocation::new( + host, + AgentTurnRequest::new( + "helper", + vec![tinyinference_llm::message::Message::user("go")], + ), + RunContext::new(RunConfig::new("limit-exceeded").with_max_model_calls(1), ()), + ), + &(), + ) + .await + .expect_err("the model-call cap must eventually fail the run"); + assert_eq!( + limit_error.kind, + crate::runtime::HostedErrorKind::LimitExceeded + ); + // The run accumulated before failing is still available. + assert!(limit_error.run.is_some()); + + // A different hosted failure mode (a security-gate denial of the user's + // input) classifies differently, proving `kind` genuinely discriminates + // rather than every non-cancel/timeout error collapsing together. + let denied_definition = + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["noop"]); + let denied_host = crate::host::HostCapabilities::new( + Arc::new(StaticContextComposer::empty()), + Arc::new(InMemoryDefinitionRegistry::new(vec![denied_definition])), + Arc::new(BlockExtensionGate), + Arc::new(FixedModelResolver::new(Arc::new(ScriptedModel::replies( + vec!["unused"], + )))), + ); + let mut denied_harness: AgentHarness<()> = AgentHarness::new(); + denied_harness.register_tool(Arc::new(NoopTool)); + let policy_error = denied_harness + .invoke_agent( + AgentInvocation::new( + denied_host, + AgentTurnRequest::new( + "helper", + vec![tinyinference_llm::message::Message::User( + tinyinference_llm::message::UserMessage { + content: vec![ + tinyinference_llm::message::ContentBlock::ProviderExtension( + json!({"secret": "block me"}), + ), + ], + }, + )], + ), + RunContext::new(RunConfig::new("policy-denied"), ()), + ), + &(), + ) + .await + .expect_err("the security gate must deny this input"); + assert_eq!(policy_error.kind, crate::runtime::HostedErrorKind::Policy); + + assert_ne!( + limit_error.kind, policy_error.kind, + "distinct hosted failure modes must classify to distinct kinds" + ); +} + #[tokio::test] async fn hosted_definition_tool_allowlist_filters_schemas_and_rejects_fabricated_calls() { let mut blocked_call = ModelResponse::assistant(""); @@ -1578,6 +1683,65 @@ async fn hosted_definition_tool_allowlist_filters_schemas_and_rejects_fabricated ); } +/// I-9 regression: a definition that declares **no** tools (an empty list — +/// `AgentDefinition::new` without `with_tools`) must deny every registered +/// tool, not grant the whole catalogue. Before the fix, `HashSet::is_empty()` +/// was read as "unrestricted" instead of "nothing authorized", so a +/// definition whose author simply forgot to declare tools (or a host that +/// failed to populate the field) silently ran with every tool available. +#[tokio::test] +async fn hosted_definition_with_no_declared_tools_denies_every_tool() { + let mut fabricated_call = ModelResponse::assistant(""); + fabricated_call + .message + .tool_calls + .push(tinyinference_llm::tool::ToolCall::new( + "call-1", + "noop", + json!({}), + )); + let model = Arc::new(ScriptedModel::new(vec![ + fabricated_call, + ModelResponse::assistant("recovered"), + ])); + // No `.with_tools(...)`: the definition declares nothing. + let definition = AgentDefinition::new("helper", "Helper", "test helper"); + let host = crate::host::HostCapabilities::new( + Arc::new(StaticContextComposer::empty()), + Arc::new(InMemoryDefinitionRegistry::new(vec![definition])), + Arc::new(AllowAllSecurityGate), + Arc::new(FixedModelResolver::new(model.clone())), + ); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_tool(Arc::new(NoopTool)); + + let run = harness + .invoke_agent( + AgentInvocation::new( + host, + AgentTurnRequest::new( + "helper", + vec![tinyinference_llm::message::Message::user("go")], + ), + RunContext::new(RunConfig::new("empty-allowlist"), ()), + ), + &(), + ) + .await + .expect("the model recovers after its denied call"); + + assert_eq!(run.text().as_deref(), Some("recovered")); + assert!( + run.messages + .iter() + .any(|message| message.text().contains("unknown tool `noop`")), + "a registered tool the definition never declared must be rejected, not silently run" + ); + // No tool schema at all is offered to the provider — the registered + // catalogue is not leaked to a definition that declared nothing. + assert!(model.requests()[0].tools.is_empty()); +} + #[tokio::test] async fn hosted_structured_schema_rejects_hidden_registered_tool_collision() { // `answer` is registered globally but deliberately not allowed for this @@ -1618,9 +1782,10 @@ async fn hosted_structured_schema_rejects_hidden_registered_tool_collision() { .await .expect_err("a hidden registered tool still collides with the schema"); + assert_eq!(error.kind, crate::runtime::HostedErrorKind::Policy); assert_eq!( error.to_string(), - "model error: hosted agent invocation failed" + "hosted agent invocation was rejected by policy" ); assert!(model.requests().is_empty(), "provider was not contacted"); } @@ -1642,11 +1807,9 @@ async fn host_security_denial_returns_a_tool_message_without_executing_the_tool( ])); let host = crate::host::HostCapabilities::new( Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )])), + Arc::new(InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["noop"]), + ])), Arc::new(DenyToolGate), Arc::new(FixedModelResolver::new(model)), ); @@ -1697,11 +1860,9 @@ async fn security_gate_sees_raw_provider_arguments_while_tools_receive_prepared_ let executed = Arc::new(Mutex::new(Vec::new())); let host = crate::host::HostCapabilities::new( Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )])), + Arc::new(InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["injected"]), + ])), gate.clone(), Arc::new(FixedModelResolver::new(model)), ); @@ -1796,11 +1957,9 @@ async fn security_gate_sees_unwrapped_arguments_for_a_bridged_deferred_call() { }); let host = crate::host::HostCapabilities::new( Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )])), + Arc::new(InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["quote"]), + ])), gate.clone(), Arc::new(FixedModelResolver::new(model)), ); @@ -1853,11 +2012,9 @@ async fn denied_tool_calls_release_their_reserved_limit_for_a_later_approval() { ])); let host = crate::host::HostCapabilities::new( Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )])), + Arc::new(InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["noop"]), + ])), Arc::new(DenyThenAllowGate { denials_remaining: AtomicUsize::new(2), }), @@ -2029,11 +2186,9 @@ async fn dropped_host_invocations_finalize_the_actual_partial_run_once() { let progress = Arc::new(RecordingProgressSink::new()); let host = crate::host::HostCapabilities::new( Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )])), + Arc::new(InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["noop"]), + ])), Arc::new(AllowAllSecurityGate), Arc::new(FixedModelResolver::new(model)), ) @@ -2168,11 +2323,9 @@ async fn denied_tool_calls_do_not_enter_terminal_executed_tool_summary() { let learning = Arc::new(RecordingLearning::default()); let host = crate::host::HostCapabilities::new( Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )])), + Arc::new(InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["noop"]), + ])), Arc::new(DenyToolGate), Arc::new(FixedModelResolver::new(Arc::new(ScriptedModel::new(vec![ tool_response, @@ -2324,11 +2477,9 @@ async fn concurrent_roots_keep_every_invocation_capability_bundle_isolated() { }), Arc::new(TaggedDefinitions { trace: trace.clone(), - inner: InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )]), + inner: InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["noop"]), + ]), }), Arc::new(TaggedSecurity { trace: trace.clone(), @@ -2611,11 +2762,9 @@ async fn budget_preflight_estimate_reflects_the_dialect_rewritten_request() { let budget = Arc::new(EstimateRecordingBudget::new()); let host = crate::host::HostCapabilities::new( Arc::new(StaticContextComposer::empty()), - Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( - "helper", - "Helper", - "test helper", - )])), + Arc::new(InMemoryDefinitionRegistry::new(vec![ + AgentDefinition::new("helper", "Helper", "test helper").with_tools(["verbose_lookup"]), + ])), Arc::new(AllowAllSecurityGate), Arc::new(FixedModelResolver::new(model.clone())), ) @@ -2789,9 +2938,10 @@ async fn hard_budget_compression_fails_closed_when_only_system_instructions_rema ) .await .expect_err("hard pressure cannot discard sole system instructions"); + assert_eq!(error.kind, crate::runtime::HostedErrorKind::Policy); assert_eq!( error.to_string(), - "model error: hosted agent invocation failed", + "hosted agent invocation was rejected by policy", "hosted callers receive no internal budget diagnostic" ); assert!(model.requests().is_empty(), "provider was never called"); @@ -2984,3 +3134,98 @@ async fn cached_host_response_does_not_re_record_provider_usage() { assert_eq!(model.requests().len(), 1, "second call is cache-served"); assert_eq!(budget.records.lock().expect("budget lock").len(), 1); } + +#[test] +fn host_invocation_binding_fails_closed_on_a_state_mismatch() { + struct OtherState; + + let host = Arc::new(crate::host::HostCapabilities::new( + Arc::new(StaticContextComposer::empty()), + Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( + "parent", "Parent", "hosted", + )])), + Arc::new(AllowAllSecurityGate), + Arc::new(FixedModelResolver::new(Arc::new(ScriptedModel::replies( + vec!["unused"], + )))), + )); + let mut context: RunContext<()> = RunContext::new(RunConfig::new("state-mismatch"), ()); + context.host_agent_id = Some("parent".to_string()); + context.host_authority = Some(Arc::new( + crate::runtime::HostInvocationAuthority::<(), ()> { + binding: Arc::new(crate::runtime::HostInvocationBinding { + host, + agent_id: "parent".to_string(), + model_pin: None, + role: None, + allowed_tools: None, + progress: None, + runtime: None, + }), + }, + )); + + // Reading it back with the *same* `State`/`Ctx` the authority was + // installed for succeeds. + assert!( + crate::runtime::host_invocation_binding::<(), ()>(&context) + .expect("matching State/Ctx must not be rejected") + .is_some() + ); + + // Reading the same context with a *different* `State` must fail closed + // rather than transmute the wrong `HostInvocationBinding<_>` out of the + // erased authority. + let mismatched = crate::runtime::host_invocation_binding::(&context); + assert!( + matches!( + mismatched, + Err(crate::error::TinyAgentsError::Validation(_)) + ), + "expected a fail-closed Validation error" + ); +} + +/// C-1 regression: `RunContext::child_with_data` (the only primitive that +/// changes `Ctx`) must never propagate host authority, closing the other +/// half of the C-1 repro (a child built with a different `Ctx` type +/// inheriting a parent's hosted authority for the wrong `Ctx`). +#[test] +fn child_with_data_never_propagates_host_authority() { + let host = Arc::new(crate::host::HostCapabilities::new( + Arc::new(StaticContextComposer::empty()), + Arc::new(InMemoryDefinitionRegistry::new(vec![AgentDefinition::new( + "parent", "Parent", "hosted", + )])), + Arc::new(AllowAllSecurityGate), + Arc::new(FixedModelResolver::new(Arc::new(ScriptedModel::replies( + vec!["unused"], + )))), + )); + let mut parent: RunContext<()> = RunContext::new(RunConfig::new("ctx-change-parent"), ()); + parent.host_authority = Some(Arc::new( + crate::runtime::HostInvocationAuthority::<(), ()> { + binding: Arc::new(crate::runtime::HostInvocationBinding { + host, + agent_id: "parent".to_string(), + model_pin: None, + role: None, + allowed_tools: None, + progress: None, + runtime: None, + }), + }, + )); + assert!(parent.host_authority.is_some()); + + // Same-`Ctx` `child` propagates authority. + let same_ctx_child = parent.child(RunConfig::new("same-ctx"), ()).unwrap(); + assert!(same_ctx_child.host_authority.is_some()); + + // Different-`Ctx` `child_with_data` never does, regardless of the + // authority the parent carries. + let different_ctx_child = parent + .child_with_data(RunConfig::new("different-ctx"), "child-data") + .unwrap(); + assert!(different_ctx_child.host_authority.is_none()); +} diff --git a/crates/tinyagents-harness/src/runtime/types.rs b/crates/tinyagents-harness/src/runtime/types.rs index 2cb3c122..5945e498 100644 --- a/crates/tinyagents-harness/src/runtime/types.rs +++ b/crates/tinyagents-harness/src/runtime/types.rs @@ -27,6 +27,7 @@ use crate::limits::RunLimits; use crate::middleware::MiddlewareStack; use crate::model_registry::ModelRegistry; use crate::retry::{FallbackPolicy, RetryPolicy}; +use crate::run_queue::QueueMode; use crate::tool::{ToolRegistry, ToolTimeoutSettings}; use tinyinference_llm::cache::CachePolicy; use tinyinference_llm::model::ResponseFormat; @@ -47,9 +48,17 @@ pub(crate) struct HostInvocationBinding { pub(crate) model_pin: Option, pub(crate) role: Option, /// Canonical names the resolved definition authorizes for this exact run. - /// An empty list retains the legacy unrestricted catalogue; a non-empty - /// list is a host boundary enforced for schemas and dispatch alike. - pub(crate) allowed_tools: HashSet, + /// + /// `None` means the definition declared no tools at all (an empty or + /// absent list) — [`crate::agent_loop`]'s `resolve_tool_allowlist` treats + /// that as fail-closed (deny every tool) by default, controlled by + /// [`HostCapabilities::fail_closed_tool_allowlist`]. `Some(set)` is + /// always the declared set, checked by plain membership: an empty + /// `HashSet` is never stored here (a declared-but-empty list is + /// collapsed to `None` at construction, so "nothing declared" and + /// "declared empty" share one fail-closed code path instead of an empty + /// set silently meaning "unrestricted", as it used to (I-9)). + pub(crate) allowed_tools: Option>, /// Per-turn ordered, nonblocking projection to the optional progress sink. pub(crate) progress: Option, /// The exact invocation-local runtime inherited by authorized children. @@ -253,9 +262,10 @@ pub struct RunPolicy { /// cheapest on tokens and the most demanding on the model, which is why /// it is opt-in only. /// - /// Whatever the dispatcher, a response with no structured calls is still - /// read through every text grammar, because native models narrate calls - /// as text often enough to matter. + /// Under a forced text dialect the answer is always read through every + /// text grammar — parsing text *is* the protocol. Under a native dialect + /// the same read is the fallback for a model that narrated a call as + /// text, gated by [`RunPolicy::text_dialect_recovery`]. pub tool_dialect: ToolDispatcher, /// Maximum consecutive re-prompts when a model signals a tool call it did /// not make: `finish_reason == "tool_calls"` with no structured call and @@ -301,6 +311,200 @@ pub struct RunPolicy { /// to. Admission still validates arguments against the *declared* schema, /// which is never looser than the projected one. pub tool_schemas: Option, + /// Whether the loop parses ``-style text-dialect markup out of + /// an assistant's visible text under a native tool dialect (see + /// [`RunPolicy::tool_dialect`]). A forced text dialect + /// ([`ToolDispatcher::Xml`] / [`ToolDispatcher::Pformat`], or + /// [`ToolDispatcher::Auto`] falling back to Xml for a model without + /// native tool calling) always parses the answer regardless of this + /// policy, since the model can only answer in text. + /// + /// Defaults to [`TextDialectRecovery::Auto`], which only attempts + /// recovery when the resolved model's + /// [`ModelProfile::tool_calling`][tinyinference_llm::model::ModelProfile::tool_calling] + /// is not reported (a model that *does* report native tool calling and + /// still answered in prose was not making a tool call — it was + /// explaining, quoting, or documenting the format, and executing that + /// text as a real call would silently strip visible text the caller + /// asked to see). See [`TextDialectRecovery`]. + pub text_dialect_recovery: TextDialectRecovery, + /// Bounds the output-validation retry loop (A3): how many times the loop + /// re-asks the model after the final turn's structured extraction fails + /// schema validation, or a registered + /// [`crate::structured::OutputValidator`] rejects an otherwise + /// schema-valid value with + /// [`crate::error::TinyAgentsError::ModelRetry`]. See + /// [`OutputRetryPolicy`]. + pub output_retry: OutputRetryPolicy, + /// What the loop does when one turn's tool calls include both a + /// structured-output "schema" call ([`StructuredStrategy::ToolCall`]'s + /// synthetic tool) and one or more genuine function-tool calls (A6). + /// Defaults to [`EndStrategy::Graceful`]. + pub end_strategy: EndStrategy, + /// Forces the [`crate::structured::StructuredStrategy::Prompted`] or + /// [`crate::structured::StructuredStrategy::ToolCallUnion`] mode for a + /// `ResponseFormat::Auto` structured-output request, bypassing + /// [`crate::structured::StructuredStrategy::for_profile`]'s + /// provider-capability heuristic (A6). + /// + /// `None` (the default) preserves the existing `Auto` resolution + /// (`ProviderSchema` or `ToolCall`, chosen from the resolved model's + /// profile). Only consulted for `ResponseFormat::Auto`; an explicit + /// `ResponseFormat::JsonSchema` always uses provider-native mode + /// regardless of this field. + pub structured_strategy_override: Option, + /// How many queued messages the loop takes from a + /// [`crate::run_queue::RunQueue`] lane at each safe boundary (A4): + /// [`QueueMode::All`] (the default) applies every pending item at once, + /// [`QueueMode::OneAtATime`] applies the oldest and leaves the rest for + /// the next boundary. Only consulted when the run's + /// [`RunContext`][crate::context::RunContext] carries a queue. + pub queue_mode: QueueMode, + /// Which engine [`AgentHarness::invoke`][super::AgentHarness::invoke] (and + /// friends) drives the loop with (A5). + /// + /// Defaults to [`LoopExecution::Direct`]: the built-in + /// [`crate::agent_loop`] body, unchanged. Setting + /// [`LoopExecution::Graph`] selects an [`AgentHarness::loop_driver` + /// ][super::AgentHarness::loop_driver] instead — install one with + /// [`AgentHarness::with_loop_driver`][super::AgentHarness::with_loop_driver] + /// (`tinyagents-graph`'s `GraphLoopDriver` is the intended implementor; + /// see `tinyagents_graph::agent_loop`). Selecting `Graph` with no driver + /// installed fails the run with + /// [`crate::error::TinyAgentsError::Validation`] rather than silently + /// falling back to `Direct`. + pub execution: LoopExecution, +} + +/// See [`RunPolicy::execution`]. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum LoopExecution { + /// Drive the run with the built-in [`crate::agent_loop`] body + /// (`run_loop`/`run_loop_body`). The default; behavior-identical to + /// every release before A5. + #[default] + Direct, + /// Drive the run with the installed + /// [`AgentHarness::loop_driver`][super::AgentHarness::loop_driver] + /// instead (a compiled-graph rendition of the loop, in the common case). + Graph, +} + +/// See [`RunPolicy::structured_strategy_override`]. +#[derive(Clone, Debug, PartialEq)] +pub enum StructuredStrategyOverride { + /// Force [`crate::structured::StructuredStrategy::Prompted`]: inject the + /// schema into the system prompt instead of using a provider schema API + /// or a forced tool call. + Prompted { + /// Custom instructions template; `None` uses + /// [`crate::structured::default_prompted_template`]. + template: Option, + }, + /// Force [`crate::structured::StructuredStrategy::ToolCallUnion`]: offer + /// one synthetic tool per `(name, schema)` variant instead of the single + /// schema from the `ResponseFormat`. + ToolCallUnion { + /// The union's variants, in the order their tools are advertised. + variants: Vec<(String, serde_json::Value)>, + }, +} + +/// Resolves the "output tool + function tools in one turn" ambiguity (A6), +/// mirroring Pydantic AI's `end_strategy`. +/// +/// The ambiguity: the model can, in a single turn, both answer (via the +/// structured-output schema call) *and* ask to run further tools. Each +/// strategy answers "what happens to those tool calls, and does the run end +/// this turn?" differently. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum EndStrategy { + /// Run the accompanying function-tool calls (so their side effects still + /// happen and their results are not silently dropped), then finish the + /// run with the structured output already recorded. The default: it + /// never discards a tool call the model asked for, but also never spends + /// an extra model call once the model has already answered. + #[default] + Graceful, + /// Finish the run immediately on the first output-tool call. The + /// accompanying function-tool calls are **not** executed; their + /// `tool_calls` entries are closed with a synthetic "run stopped before + /// this tool call was executed" result so the transcript stays + /// replayable. Use when the structured answer must win even if it means + /// dropping tool calls the model also happened to request. + Early, + /// Ignore the output-tool call this turn (do not record it, do not + /// finish): run the function-tool calls and give the model another turn, + /// exactly as if the output tool had not been called. The run only + /// finishes once a turn produces the output tool with **no** accompanying + /// function-tool calls. Use when function tools must always be allowed to + /// run to completion before an answer is accepted. + Exhaustive, +} + +/// Policy for the output-validation retry loop (A3), mirroring Pydantic AI's +/// `retries={'output': N}`. +/// +/// On the agent loop's final turn, a structured-extraction failure or a +/// registered [`crate::structured::OutputValidator`] rejection no longer +/// immediately fails the run: the error is pushed back to the model as a +/// repair prompt (built from [`Self::message_template`]) and the loop asks +/// again, up to [`Self::max_attempts`] times total for the run. Each retry +/// still counts against [`RunLimits::max_model_calls`] like any other model +/// call — this policy only bounds how many of those calls may be spent on +/// output repair specifically. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct OutputRetryPolicy { + /// How many times the loop may re-ask the model after an output + /// validation failure. `0` disables the retry loop entirely — the first + /// failure fails the run, exactly as before A3. + pub max_attempts: u8, + /// The repair-prompt template pushed to the model as a + /// [`tinyinference_llm::message::Message::user`] turn. `{error}` is + /// replaced with the extraction/validation error text; a template + /// without that placeholder still works (the error is simply omitted) + /// but loses the specific reason. + pub message_template: String, +} + +impl Default for OutputRetryPolicy { + fn default() -> Self { + Self { + max_attempts: 1, + message_template: "{error}\n\nFix the errors and try again.".to_string(), + } + } +} + +/// Policy for recovering ``-style text-dialect tool calls from an +/// assistant's visible text. +/// +/// Some providers/models emit tool calls as XML-ish markup inside ordinary +/// text instead of (or in addition to failing to populate) the provider's +/// native tool-call channel. Recovering that markup lets such a model still +/// drive tools through the same loop as a model with native tool calling. +/// +/// Left unconditional, this is a real correctness hazard: any assistant text +/// that merely *quotes* `` markup — explaining the format to a +/// user, echoing a worked example, or showing it in a fenced code block — +/// gets executed as a real tool call, with the visible text silently +/// stripped and replaced. [`TextDialectRecovery::Auto`] (the default) closes +/// the common case of that hazard by skipping recovery for any model whose +/// resolved profile reports native tool calling; recovery inside fenced code +/// blocks is always skipped regardless of this policy, since a model +/// demonstrating the syntax in a code fence is manifestly not making a call. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum TextDialectRecovery { + /// Never parse text-dialect tool calls. + Off, + /// Always attempt recovery when the provider returned no native tool + /// calls, regardless of the resolved model's advertised capabilities. + On, + /// Attempt recovery only when the resolved model's profile does not + /// report native tool calling (or the profile is unknown). This is the + /// default. + #[default] + Auto, } impl Default for RunPolicy { @@ -328,8 +532,14 @@ impl Default for RunPolicy { // caller, so one stochastic-failure retry is strictly better than a // blank final. truncated_empty_retries: 1, + text_dialect_recovery: TextDialectRecovery::default(), discovery: crate::tool::discover::ToolDiscoveryPolicy::default(), tool_schemas: None, + output_retry: OutputRetryPolicy::default(), + end_strategy: EndStrategy::default(), + structured_strategy_override: None, + queue_mode: QueueMode::default(), + execution: LoopExecution::default(), } } } @@ -377,6 +587,44 @@ pub struct AgentHarness { /// into it. Because it is owned by the harness rather than a single run, a /// repeated identical request can be served from an earlier run's result. pub(crate) response_cache: Option>, + /// Optional validator consulted after the final turn's structured + /// extraction succeeds, driving the output-validation retry loop (A3). + /// See [`crate::structured::OutputValidator`] and + /// [`AgentHarness::with_output_validator`]. + pub(crate) output_validator: Option>>, + /// Optional inline resolver for deferred tool calls (A2). When set, a + /// batch that defers calls is resolved through it and the loop keeps + /// going instead of exiting with `AgentRun::deferred`. See + /// [`crate::tool::DeferredToolHandler`] and + /// [`AgentHarness::with_deferred_tool_handler`]. + pub(crate) deferred_tool_handler: Option>, + /// Optional composable [`crate::tool::toolset::ToolSet`] chain + /// (gap B3) consulted for the model-visible tool catalogue and, when a + /// call is not owned by [`Self::tools`], for dispatch. + /// + /// `None` (the default) preserves every existing harness's behavior + /// unchanged: the loop resolves tools from [`Self::tools`] alone, exactly + /// as before this field existed. Set with + /// [`AgentHarness::with_toolset`]. See that method's doc comment for + /// exactly which turn behavior this changes. + pub(crate) toolset: Option>>, + /// Capability bundles installed via [`AgentHarness::with_capability`] + /// (gap G3), in installation order. Kept so each new `with_capability` + /// call can rebuild [`Self::toolset`]'s + /// [`crate::capability::CapabilityToolSet`] layer from the complete, + /// still-accumulating list rather than nesting one per call. + pub(crate) capabilities: Vec>, + /// The toolset chain that was installed (via [`AgentHarness::with_toolset`], + /// or `None`) before the first [`AgentHarness::with_capability`] call. + /// Captured once so every later `with_capability` rebuild of + /// [`Self::toolset`] keeps composing with it, instead of losing it to + /// the first capability's rebuild. + pub(crate) capability_base_toolset: Option>>, + /// Alternate loop engine selected when [`RunPolicy::execution`] is + /// [`LoopExecution::Graph`] (A5). See + /// [`crate::agent_loop::phases::LoopDriver`] and + /// [`AgentHarness::with_loop_driver`]. + pub(crate) loop_driver: Option>>, } /// The non-serializable mechanics selected for one hosted invocation. diff --git a/crates/tinyagents-harness/src/steering/mod.rs b/crates/tinyagents-harness/src/steering/mod.rs index bcbe6439..17ce0064 100644 --- a/crates/tinyagents-harness/src/steering/mod.rs +++ b/crates/tinyagents-harness/src/steering/mod.rs @@ -59,8 +59,9 @@ use std::collections::{HashSet, VecDeque}; use std::sync::{Arc, Mutex}; use crate::context::RunContext; -use crate::error::{Result, TinyAgentsError}; +use crate::error::Result; use crate::events::AgentEvent; +use crate::ids::RunId; use tinyinference_llm::message::Message; // ── SteeringPolicy ──────────────────────────────────────────────────────────── @@ -99,14 +100,20 @@ impl SteeringPolicy { impl SteeringHandle { /// Builds a handle backed by a fresh, empty queue gated by `policy`. + /// + /// The handle is unbound (empty `run_id`, `is_root = true`) until it is + /// attached to a run via + /// [`RunContext::with_steering`][crate::context::RunContext::with_steering], + /// which binds it to that run's id as the root of its steering tree. pub fn new(policy: SteeringPolicy) -> Self { Self { inner: Arc::new(SteeringInner { queue: Mutex::new(VecDeque::new()), policy, - paused: Mutex::new(None), - checkpoints: Mutex::new(0), }), + run_id: RunId::new(""), + is_root: true, + local: Arc::new(SteeringLocal::default()), } } @@ -116,49 +123,118 @@ impl SteeringHandle { Self::new(SteeringPolicy::allow_all()) } - /// Enqueues `command` for delivery to the running agent loop. + /// Binds this handle to `run_id` as the **root** of its steering tree. + /// + /// Called by [`RunContext::with_steering`][crate::context::RunContext::with_steering] + /// when an orchestrator attaches a handle to a run; every + /// [`SteeringTarget::Root`]-addressed command drains here. + pub(crate) fn bind_root(&self, run_id: RunId) -> Self { + Self { + inner: Arc::clone(&self.inner), + run_id, + is_root: true, + local: Arc::clone(&self.local), + } + } + + /// Derives a handle scoped to a child run. + /// + /// Shares the underlying queue and policy (so an orchestrator holding the + /// root handle can still reach the child by [`SteeringTarget::Run`] or + /// [`SteeringTarget::All`]), but gets its own identity and its own + /// pause/checkpoint state: a command addressed to the parent (or to + /// [`SteeringTarget::Root`]) is invisible to [`SteeringHandle::drain`] on + /// the child, and a pause latched on the child does not latch the parent's. + /// This is what keeps an `Inject`/`Pause` meant for the orchestrator from + /// being consumed by whichever sub-agent happens to reach a checkpoint + /// first (see I-5). + pub(crate) fn for_child(&self, run_id: RunId) -> Self { + Self { + inner: Arc::clone(&self.inner), + run_id, + is_root: false, + local: Arc::new(SteeringLocal::default()), + } + } + + /// Enqueues `command` addressed to [`SteeringTarget::Root`]. /// /// The command becomes visible to the loop at its next steering checkpoint; - /// this method never blocks and does not itself check the policy. + /// this method never blocks and does not itself check the policy. Use + /// [`SteeringHandle::send_to`] to address a specific descendant run, or + /// [`SteeringHandle::send_all`] to reach every run sharing this handle. + pub fn send(&self, command: SteeringCommand) { + self.send_to(SteeringTarget::Root, command); + } + + /// Enqueues `command` addressed to `target`. /// /// Queue accessors recover from a poisoned mutex (a panic in another /// holder) instead of panicking: the queue is a plain `VecDeque` with no /// invariants that a panicking holder could break mid-update. - pub fn send(&self, command: SteeringCommand) { + pub fn send_to(&self, target: SteeringTarget, command: SteeringCommand) { self.inner .queue .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .push_back(command); + .push_back((target, command)); + } + + /// Enqueues `command` addressed to every run sharing this handle + /// ([`SteeringTarget::All`]). + pub fn send_all(&self, command: SteeringCommand) { + self.send_to(SteeringTarget::All, command); } - /// Removes and returns all currently queued commands in FIFO order, leaving - /// the queue empty. Called by the agent loop at each checkpoint. + /// Removes and returns the commands addressed to *this* handle's run (its + /// own [`SteeringTarget::Run`], [`SteeringTarget::Root`] if this handle is + /// the root, or [`SteeringTarget::All`]), leaving commands addressed to + /// other runs in the shared queue for them to drain later. Called by the + /// agent loop at each checkpoint. pub fn drain(&self) -> Vec { let mut queue = self .inner .queue .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - queue.drain(..).collect() + let mut matched = Vec::new(); + let mut remaining = VecDeque::with_capacity(queue.len()); + for (target, command) in queue.drain(..) { + if self.matches(&target) { + matched.push(command); + } else { + remaining.push_back((target, command)); + } + } + *queue = remaining; + matched } - /// Returns `true` when no commands are currently queued. + /// Returns `true` when `target` addresses this handle's run. + fn matches(&self, target: &SteeringTarget) -> bool { + match target { + SteeringTarget::Root => self.is_root, + SteeringTarget::Run(id) => *id == self.run_id, + SteeringTarget::All => true, + } + } + + /// Returns `true` when no commands addressed to this handle's run are + /// currently queued. pub fn is_empty(&self) -> bool { - self.inner - .queue - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_empty() + self.pending() == 0 } - /// Returns the number of commands currently queued. + /// Returns the number of commands currently queued that are addressed to + /// this handle's run. pub fn pending(&self) -> usize { self.inner .queue .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .len() + .iter() + .filter(|(target, _)| self.matches(target)) + .count() } /// Returns the policy gating this handle. @@ -191,7 +267,7 @@ impl SteeringHandle { reason, paused_at_checkpoint: checkpoint, }); - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::steering", checkpoint = state.paused_at_checkpoint, reason = state.reason.as_deref(), @@ -208,7 +284,7 @@ impl SteeringHandle { pub fn resume(&self) -> Option { let cleared = self.lock_paused().take(); if cleared.is_some() { - tinyagents_tracing::debug!(target: "tinyagents::steering", "[steering] pause cleared by resume"); + tracing::debug!(target: "tinyagents::steering", "[steering] pause cleared by resume"); } cleared } @@ -217,7 +293,7 @@ impl SteeringHandle { /// the *current* checkpoint is what a pause records (not the next one). fn advance_checkpoint(&self) -> usize { let mut checkpoints = self - .inner + .local .checkpoints .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); @@ -229,7 +305,7 @@ impl SteeringHandle { /// Locks the pause latch, recovering from poisoning (see /// [`SteeringHandle::send`]). fn lock_paused(&self) -> std::sync::MutexGuard<'_, Option> { - self.inner + self.local .paused .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -248,13 +324,16 @@ impl SteeringHandle { /// /// - When `ctx` has no [`SteeringHandle`], returns /// [`SteeringOutcome::Continue`] without emitting anything. -/// - The batch is **validated in full before anything is applied**. If any -/// command is disallowed, an [`AgentEvent::Steered`] with `accepted = false` -/// is emitted for it and [`TinyAgentsError::Steering`] is returned — with the -/// working transcript and run metadata completely untouched. (It used to -/// validate lazily while applying, so a rejected command at position *n* left -/// commands `0..n` already applied, commands after it dropped, and the run -/// erroring: a partially-steered run and no way to reason about its state.) +/// - Only commands addressed to this run are drained: see +/// [`SteeringHandle::drain`] and [`SteeringHandle::for_child`]. Commands +/// addressed to a different run stay queued for it. +/// - Each command is checked **individually** against the run's +/// [`SteeringPolicy`]. A disallowed command is rejected on its own — an +/// [`AgentEvent::Steered`] with `accepted = false` is emitted for it, and the +/// checkpoint moves on to the next command in the batch — rather than +/// aborting the whole batch or the run. (It used to reject the entire batch, +/// including commands the policy *did* permit, whenever one command in it +/// was disallowed.) /// - [`SteeringCommand::Cancel`] takes precedence: it is applied (emitting an /// accepted event) and the function returns [`SteeringOutcome::Cancel`] /// immediately, ignoring the rest of the batch. @@ -269,9 +348,9 @@ impl SteeringHandle { /// /// # Errors /// -/// Returns [`TinyAgentsError::Steering`] when any drained command is not -/// permitted by the run's [`SteeringPolicy`]. No command in the batch is -/// applied in that case. +/// This function no longer errors on a policy-disallowed command — see above. +/// It returns `Err` only if a future extension needs to signal a checkpoint +/// failure that is not representable as a rejected command. pub fn apply_pending_steering( ctx: &mut RunContext, messages: &mut Vec, @@ -284,35 +363,7 @@ pub fn apply_pending_steering( let checkpoint = handle.advance_checkpoint(); let commands = handle.drain(); - // ── Phase 1: validate the whole batch, mutating nothing ───────────────── - // - // A policy violation must abort the checkpoint *atomically*. Checking as we - // apply means the run dies with some of the batch already in the - // transcript. - if let Some(rejected) = commands - .iter() - .map(SteeringCommand::kind) - .find(|kind| !handle.policy().is_allowed(*kind)) - { - tinyagents_tracing::debug!( - target: "tinyagents::steering", - checkpoint, - command_kind = rejected.as_str(), - batch_size = commands.len(), - "[steering] batch rejected by policy; nothing applied" - ); - ctx.emit(AgentEvent::Steered { - command_kind: rejected.as_str().to_string(), - accepted: false, - }); - return Err(TinyAgentsError::Steering(format!( - "steering command `{}` is not permitted by the run policy", - rejected.as_str() - ))); - } - - // ── Phase 2: apply ────────────────────────────────────────────────────── - tinyagents_tracing::debug!( + tracing::debug!( target: "tinyagents::steering", checkpoint, batch_size = commands.len(), @@ -322,6 +373,23 @@ pub fn apply_pending_steering( for command in commands { let kind = command.kind(); + // Each command is validated on its own: a disallowed command is + // rejected individually (I-5/M-7) rather than voiding the whole + // batch or killing the run. + if !handle.policy().is_allowed(kind) { + tracing::debug!( + target: "tinyagents::steering", + checkpoint, + command_kind = kind.as_str(), + "[steering] command rejected by policy; skipped" + ); + ctx.emit(AgentEvent::Steered { + command_kind: kind.as_str().to_string(), + accepted: false, + }); + continue; + } + match command { SteeringCommand::Pause => { handle.latch_pause(checkpoint, None); diff --git a/crates/tinyagents-harness/src/steering/test.rs b/crates/tinyagents-harness/src/steering/test.rs index 72f1d7ac..b2a4dac1 100644 --- a/crates/tinyagents-harness/src/steering/test.rs +++ b/crates/tinyagents-harness/src/steering/test.rs @@ -17,7 +17,7 @@ use crate::events::AgentEvent; use crate::runtime::AgentHarness; use crate::steering::{ SteeringCommand, SteeringCommandKind, SteeringHandle, SteeringOutcome, SteeringPolicy, - apply_pending_steering, + SteeringTarget, apply_pending_steering, }; use crate::testkit::{EventRecorder, Trajectory}; use tinyinference_llm::message::Message; @@ -39,6 +39,7 @@ fn text_response(text: &str) -> ModelResponse { )], tool_calls: Vec::new(), usage: Some(Usage::new(1, 1)), + origin: None, }, usage: Some(Usage::new(1, 1)), finish_reason: Some("stop".to_string()), @@ -87,6 +88,7 @@ impl ChatModel<()> for RecordingModel { content: Vec::new(), tool_calls: vec![ToolCall::new("c1", "noop", json!({}))], usage: Some(Usage::new(1, 1)), + origin: None, }, usage: Some(Usage::new(1, 1)), finish_reason: Some("tool_calls".to_string()), @@ -237,7 +239,7 @@ fn cancel_wins_over_later_commands() { } #[test] -fn disallowed_command_is_rejected_with_steering_error_and_event() { +fn disallowed_command_is_rejected_with_steered_event_and_the_run_continues() { let recorder = EventRecorder::new(); // Policy permits Pause but not Cancel. let handle = SteeringHandle::new(SteeringPolicy::new().allow(SteeringCommandKind::Pause)); @@ -247,8 +249,10 @@ fn disallowed_command_is_rejected_with_steering_error_and_event() { .with_steering(handle); let mut messages = Vec::new(); - let err = apply_pending_steering(&mut ctx, &mut messages).unwrap_err(); - assert!(matches!(err, TinyAgentsError::Steering(_)), "got {err:?}"); + // A disallowed command is rejected on its own; it no longer fails the + // checkpoint (I-5/M-7). + let outcome = apply_pending_steering(&mut ctx, &mut messages).unwrap(); + assert_eq!(outcome, SteeringOutcome::Continue); assert_eq!( recorder.events(), vec![AgentEvent::Steered { @@ -356,25 +360,27 @@ async fn cancel_terminates_the_run() { } #[tokio::test] -async fn disallowed_command_fails_the_run() { +async fn disallowed_command_is_skipped_and_the_run_still_completes() { let recorder = EventRecorder::new(); // Empty policy: every command is rejected. let handle = SteeringHandle::new(SteeringPolicy::new()); handle.send(SteeringCommand::InjectMessage(Message::user("nope"))); let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("never reached"))); + harness.register_model("mock", Arc::new(MockModel::constant("reached"))); let ctx: RunContext = RunContext::new(RunConfig::new("run-reject"), ()) .with_events(recorder.sink()) .with_steering(handle); - let err = harness + // A disallowed steering command no longer kills the run (I-5/M-7); it is + // rejected individually and the loop continues. + let run = harness .invoke_in_context(&(), ctx, vec![Message::user("start")]) .await - .expect_err("run should fail on disallowed steering"); + .expect("run should complete despite the rejected steering command"); + assert_eq!(run.text(), Some("reached".to_string())); - assert!(matches!(err, TinyAgentsError::Steering(_)), "got {err:?}"); assert!(recorder.events().iter().any(|e| matches!( e, AgentEvent::Steered { command_kind, accepted: false } if command_kind == "inject_message" @@ -405,15 +411,16 @@ fn steering_queue_recovers_from_poisoned_lock() { assert!(handle.is_empty()); } -// ── LOOP-8(a): the batch is validated before anything is applied ────────────── +// ── I-5/M-7: a disallowed command in a batch is rejected individually ───────── #[test] -fn a_rejected_command_leaves_no_earlier_command_applied() { - // Regression test (LOOP-8a): `apply_pending_steering` drained the whole - // batch up front and then validated lazily *while applying*, so a policy - // violation at position 2 left commands 0 and 1 already in the transcript, - // command 3 silently dropped, and the run erroring. The checkpoint must be - // atomic: reject the batch, change nothing. +fn a_rejected_command_in_a_batch_does_not_drop_the_allowed_ones() { + // Regression test (I-5/M-7): `apply_pending_steering` used to validate + // the whole drained batch up front and refuse it entirely — including + // commands the policy *did* permit — the moment one command in it was + // disallowed, and the caller's `?` then killed the run. A command the + // policy disallows must be rejected on its own; every allowed command in + // the same batch still applies, and the checkpoint does not error. let recorder = EventRecorder::new(); let handle = SteeringHandle::new( SteeringPolicy::new() @@ -424,7 +431,7 @@ fn a_rejected_command_leaves_no_earlier_command_applied() { handle.send(SteeringCommand::SetMetadata { metadata: serde_json::json!({"tag": "applied"}), }); - // Not allowed → the whole batch must be refused. + // Not allowed → rejected individually, the rest of the batch still runs. handle.send(SteeringCommand::Cancel); handle.send(SteeringCommand::InjectMessage(Message::user("last"))); @@ -433,25 +440,40 @@ fn a_rejected_command_leaves_no_earlier_command_applied() { .with_steering(handle); let mut messages = Vec::new(); - let err = apply_pending_steering(&mut ctx, &mut messages).unwrap_err(); - assert!(matches!(err, TinyAgentsError::Steering(_)), "got {err:?}"); + let outcome = apply_pending_steering(&mut ctx, &mut messages).unwrap(); + assert_eq!(outcome, SteeringOutcome::Continue); - assert!( - messages.is_empty(), - "an earlier command in a rejected batch was applied: {messages:?}" + assert_eq!( + messages, + vec![Message::user("first"), Message::user("last")], + "allowed commands in the batch should still have applied" ); assert_eq!( ctx.config.metadata, - serde_json::Value::Null, - "metadata was mutated by a rejected batch" + serde_json::json!({"tag": "applied"}), + "the allowed SetMetadata command should still have applied" ); - // Exactly one event, for the offending command. + // Every command gets its own event: accepted, accepted, rejected, accepted. assert_eq!( recorder.events(), - vec![AgentEvent::Steered { - command_kind: "cancel".to_string(), - accepted: false, - }] + vec![ + AgentEvent::Steered { + command_kind: "inject_message".to_string(), + accepted: true, + }, + AgentEvent::Steered { + command_kind: "set_metadata".to_string(), + accepted: true, + }, + AgentEvent::Steered { + command_kind: "cancel".to_string(), + accepted: false, + }, + AgentEvent::Steered { + command_kind: "inject_message".to_string(), + accepted: true, + }, + ] ); } @@ -565,14 +587,18 @@ fn pause_with_is_gated_by_the_same_policy_kind_as_pause() { SteeringCommandKind::Pause ); - // A policy that forbids Pause forbids PauseWith too. + // A policy that forbids Pause forbids PauseWith too: rejected + // individually, and the checkpoint continues rather than the run dying. let handle = SteeringHandle::new(SteeringPolicy::new().allow(SteeringCommandKind::Resume)); handle.send(SteeringCommand::PauseWith { reason: "why".into(), }); let mut ctx: RunContext = RunContext::new(RunConfig::new("r"), ()).with_steering(handle); let mut messages = Vec::new(); - assert!(apply_pending_steering(&mut ctx, &mut messages).is_err()); + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Continue + ); } #[test] @@ -604,3 +630,86 @@ fn pause_with_round_trips_through_json() { let back: SteeringCommand = serde_json::from_value(json).expect("deserialize"); assert_eq!(back, command); } + +// ── I-5: a child run only drains commands addressed to it ───────────────────── + +#[test] +fn root_addressed_command_is_not_consumed_by_a_child() { + // Regression test (I-5): `RunContext::child` used to hand the child a bare + // clone of the parent's `SteeringHandle`, so a command an orchestrator + // addressed to the parent (the default target) could be drained and + // applied by whichever sub-agent reached its checkpoint first. + let handle = SteeringHandle::allow_all(); + let mut parent: RunContext = + RunContext::new(RunConfig::new("parent"), ()).with_steering(handle.clone()); + let child_config = RunConfig::new("child"); + let mut child: RunContext = parent.child(child_config, ()).unwrap(); + + // Addressed to the default target (Root == the parent). + handle.send(SteeringCommand::InjectMessage(Message::user( + "for the parent", + ))); + + let mut child_messages = Vec::new(); + let outcome = apply_pending_steering(&mut child, &mut child_messages).unwrap(); + assert_eq!(outcome, SteeringOutcome::Continue); + assert!( + child_messages.is_empty(), + "child drained a command addressed to the root: {child_messages:?}" + ); + + // The parent's own checkpoint still sees it. + let mut parent_messages = Vec::new(); + apply_pending_steering(&mut parent, &mut parent_messages).unwrap(); + assert_eq!(parent_messages, vec![Message::user("for the parent")]); +} + +#[test] +fn run_addressed_command_reaches_only_that_run() { + let handle = SteeringHandle::allow_all(); + let parent: RunContext = + RunContext::new(RunConfig::new("parent"), ()).with_steering(handle.clone()); + let mut child_a: RunContext = parent.child(RunConfig::new("child-a"), ()).unwrap(); + let mut child_b: RunContext = parent.child(RunConfig::new("child-b"), ()).unwrap(); + + handle.send_to( + SteeringTarget::Run(child_a.run_id().clone()), + SteeringCommand::InjectMessage(Message::user("for child-a only")), + ); + + let mut a_messages = Vec::new(); + apply_pending_steering(&mut child_a, &mut a_messages).unwrap(); + assert_eq!(a_messages, vec![Message::user("for child-a only")]); + + let mut b_messages = Vec::new(); + apply_pending_steering(&mut child_b, &mut b_messages).unwrap(); + assert!( + b_messages.is_empty(), + "a command addressed to child-a leaked into child-b: {b_messages:?}" + ); +} + +#[test] +fn all_addressed_command_is_drained_by_whichever_run_checkpoints_first() { + // `SteeringTarget::All` matches any handle sharing the queue, but delivery + // is still pull-and-consume-once: whichever run reaches its checkpoint + // first drains it, exactly like the pre-routing behaviour for every + // command. It is documented that way (see `SteeringTarget::All`), unlike + // `Root`/`Run(id)` which are exclusive to one run by construction. + let handle = SteeringHandle::allow_all(); + let parent: RunContext = + RunContext::new(RunConfig::new("parent"), ()).with_steering(handle.clone()); + let mut child: RunContext = parent.child(RunConfig::new("child"), ()).unwrap(); + let mut parent = parent; + + handle.send_all(SteeringCommand::InjectMessage(Message::user("broadcast"))); + + let mut child_messages = Vec::new(); + apply_pending_steering(&mut child, &mut child_messages).unwrap(); + assert_eq!(child_messages, vec![Message::user("broadcast")]); + + // Already drained by the child; the parent's checkpoint sees nothing. + let mut parent_messages = Vec::new(); + apply_pending_steering(&mut parent, &mut parent_messages).unwrap(); + assert!(parent_messages.is_empty()); +} diff --git a/crates/tinyagents-harness/src/steering/types.rs b/crates/tinyagents-harness/src/steering/types.rs index c6bc9736..2cc77f5a 100644 --- a/crates/tinyagents-harness/src/steering/types.rs +++ b/crates/tinyagents-harness/src/steering/types.rs @@ -15,8 +15,34 @@ use std::sync::{Arc, Mutex}; use serde::{Deserialize, Serialize}; +use crate::ids::RunId; use tinyinference_llm::message::Message; +/// Which run in the recursion tree a queued [`SteeringCommand`] is addressed +/// to. +/// +/// Every [`SteeringHandle`] clone shares one underlying queue (so an +/// orchestrator can hand a single handle to a deeply nested tree and still +/// reach any run in it), but each level of the tree only *drains* the entries +/// addressed to itself — see [`SteeringHandle::for_child`]. Without this, a +/// command meant for the orchestrating run could be consumed by whichever +/// sub-agent happened to reach a checkpoint first. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum SteeringTarget { + /// The root run of the tree this handle belongs to. This is the default + /// target for [`SteeringHandle::send`]. + Root, + /// A specific run, named by [`RunId`]. + Run(RunId), + /// Matched by any run sharing this handle (root or any descendant). + /// + /// Delivery is still pull-and-consume-once, exactly like every other + /// target: whichever run's checkpoint drains the queue first removes the + /// entry, so `All` is not a broadcast to every run in the tree — it only + /// widens *which* run may claim the command, not how many do. + All, +} + /// A typed runtime control instruction delivered to a running agent loop. /// /// Commands are enqueued on a [`SteeringHandle`] by an orchestrator and drained @@ -229,17 +255,43 @@ pub struct PauseState { /// The handle is std-only — it carries no async runtime dependency. Delivery is /// pull-based: enqueued commands become visible to the loop on its next /// checkpoint, never mid-stream. +/// +/// # Routing +/// +/// A plain `clone()` is a bare alias: it shares this handle's identity +/// (`run_id`/`is_root`) as well as its queue, so it drains exactly the same +/// commands this handle would. When a run spawns a child, +/// [`crate::context::RunContext::child`] calls [`SteeringHandle::for_child`] +/// (not `clone`) so the child only drains commands addressed to it or to +/// [`SteeringTarget::All`] — see that method's docs. #[derive(Clone)] pub struct SteeringHandle { pub(crate) inner: Arc, + /// The identity of the run *this handle instance* drains for. + pub(crate) run_id: RunId, + /// Whether `run_id` is the root of the steering tree, for matching + /// [`SteeringTarget::Root`]. + pub(crate) is_root: bool, + /// This handle's own pause/checkpoint state. Deliberately **not** shared + /// with a parent/child handle derived via [`SteeringHandle::for_child`]: + /// a pause addressed to one run must not latch every run sharing the + /// underlying queue. + pub(crate) local: Arc, } -/// Shared interior of a [`SteeringHandle`]. +/// Shared interior of a [`SteeringHandle`]: the queue and policy every level +/// of a steering tree drains from. pub(crate) struct SteeringInner { - /// FIFO queue of pending commands. - pub(crate) queue: Mutex>, + /// FIFO queue of pending, addressed commands. + pub(crate) queue: Mutex>, /// The allowlist gating which drained commands may be applied. pub(crate) policy: SteeringPolicy, +} + +/// Per-run steering state: **not** shared across a [`SteeringHandle::for_child`] +/// boundary, so a pause or checkpoint count is scoped to the run it belongs to. +#[derive(Default)] +pub(crate) struct SteeringLocal { /// The latched pause, if one is in effect. Survives across checkpoints so a /// [`SteeringCommand::Resume`] delivered in a *later* batch can lift it. pub(crate) paused: Mutex>, diff --git a/crates/tinyagents-harness/src/store/conformance.rs b/crates/tinyagents-harness/src/store/conformance.rs new file mode 100644 index 00000000..8eba43fa --- /dev/null +++ b/crates/tinyagents-harness/src/store/conformance.rs @@ -0,0 +1,399 @@ +//! Reusable storage **conformance** (contract) suites for the harness +//! [`Store`] and [`NamespacedStore`] traits. +//! +//! Mirrors the pattern the graph module established in +//! `tinyagents_graph::testkit::conformance` and the session crate's +//! `tinyagents_session::testkit::conformance`: the same assertions are run +//! against every backend so a defect in one implementation cannot hide behind +//! another, and a downstream author implementing either trait can certify +//! their backend by calling the matching function from a `#[tokio::test]`. +//! +//! # Example +//! +//! ```rust +//! use tinyagents_harness::store::InMemoryStore; +//! use tinyagents_harness::store::conformance::run_store_conformance; +//! +//! # #[tokio::main] +//! # async fn main() { +//! run_store_conformance(&InMemoryStore::new()).await; +//! # } +//! ``` +//! +//! Each function panics with a descriptive message on the first violation. + +use serde_json::json; + +use super::Store; +use super::namespaced::{ + FilterOp, ListNamespacesQuery, Namespace, NamespacedStore, SearchQuery, StoreOp, +}; + +/// Runs the flat [`Store`] contract against `store`. +/// +/// Covers put/get (including overwrite), delete (including deleting an +/// absent key, which must not error), list, and namespace isolation. Any +/// backend that passes this behaves interchangeably as a flat harness +/// [`Store`] — this is what [`crate::store::InMemoryStore`] and +/// [`crate::store::FileStore`] are both certified against. +pub async fn run_store_conformance(store: &S) { + // put + get round-trips the exact value. + store + .put("ns-a", "k1", json!({"v": 1})) + .await + .expect("put k1"); + let got = store.get("ns-a", "k1").await.expect("get k1"); + assert_eq!(got, Some(json!({"v": 1})), "get returns what was put"); + + // A key that was never written is a `None`, not an error. + assert!( + store + .get("ns-a", "missing") + .await + .expect("get of a missing key") + .is_none(), + "get of a never-written key returns None" + ); + + // A namespace that was never written behaves the same way for both get + // and list — no error, just absence. + assert!( + store + .get("ns-never-written", "k1") + .await + .expect("get from a never-written namespace") + .is_none(), + "get from a never-written namespace returns None" + ); + assert!( + store + .list("ns-never-written") + .await + .expect("list of a never-written namespace") + .is_empty(), + "list of a never-written namespace returns an empty Vec" + ); + + // put is an upsert: writing the same key again replaces the value. + store + .put("ns-a", "k1", json!({"v": 2})) + .await + .expect("overwrite k1"); + let got = store + .get("ns-a", "k1") + .await + .expect("get k1 after overwrite"); + assert_eq!(got, Some(json!({"v": 2})), "put overwrites the prior value"); + + // list enumerates every key written to the namespace. + store + .put("ns-a", "k2", json!("second")) + .await + .expect("put k2"); + let mut keys = store.list("ns-a").await.expect("list ns-a"); + keys.sort(); + assert_eq!( + keys, + vec!["k1".to_string(), "k2".to_string()], + "list returns every key written to the namespace" + ); + + // Namespaces are independent: a write to one is invisible from another, + // and listing one namespace never surfaces another's keys. + store + .put("ns-b", "k1", json!("other namespace")) + .await + .expect("put ns-b k1"); + assert_eq!( + store.get("ns-b", "k1").await.expect("get ns-b k1"), + Some(json!("other namespace")), + "a namespace's own write is visible" + ); + assert_eq!( + store + .get("ns-a", "k1") + .await + .expect("get ns-a k1 unaffected"), + Some(json!({"v": 2})), + "writing to ns-b does not affect ns-a's value for the same key" + ); + assert_eq!( + store.list("ns-b").await.expect("list ns-b"), + vec!["k1".to_string()], + "listing ns-b never surfaces ns-a's keys" + ); + + // delete removes exactly the deleted key. + store.delete("ns-a", "k1").await.expect("delete k1"); + assert!( + store + .get("ns-a", "k1") + .await + .expect("get after delete") + .is_none(), + "a deleted key reads back as None" + ); + assert_eq!( + store.list("ns-a").await.expect("list after delete"), + vec!["k2".to_string()], + "delete removes exactly the deleted key, leaving the rest of the namespace intact" + ); + + // Deleting an already-absent key, or a key in a namespace that was never + // written, is a no-op — not an error. + store + .delete("ns-a", "k1") + .await + .expect("delete of an already-absent key must not error"); + store + .delete("ns-never-written", "nope") + .await + .expect("delete from a never-written namespace must not error"); +} + +/// Runs the [`NamespacedStore`] contract against `store`. +/// +/// Covers put/get (including overwrite preserving `created_at_ms`), delete, +/// search (namespace-prefix scoping and field filtering), namespace listing, +/// TTL expiry, and `batch`'s positional-alignment guarantee. Any backend that +/// passes this behaves interchangeably as a [`NamespacedStore`] — this is +/// what [`crate::store::namespaced::InMemoryNamespacedStore`] is certified +/// against. +pub async fn run_namespaced_store_conformance(store: &S) { + let ns_alice = Namespace::new(["users", "alice"]).expect("valid namespace"); + let ns_bob = Namespace::new(["users", "bob"]).expect("valid namespace"); + + // put + get round-trips the exact value, addressed by namespace and key. + store + .put(&ns_alice, "profile", json!({"name": "Alice"})) + .await + .expect("put alice profile"); + let item = store + .get(&ns_alice, "profile") + .await + .expect("get alice profile") + .expect("item is present"); + assert_eq!( + item.value, + json!({"name": "Alice"}), + "get returns what was put" + ); + assert_eq!(item.namespace, ns_alice, "item records its own namespace"); + assert_eq!(item.key, "profile", "item records its own key"); + + // A key that was never written is a `None`, not an error. + assert!( + store + .get(&ns_alice, "missing") + .await + .expect("get of a missing key") + .is_none(), + "get of a never-written key returns None" + ); + + // Overwriting preserves the original creation time and advances the + // update time — an item's identity survives a write, only its content and + // freshness change. + let before = store + .get(&ns_alice, "profile") + .await + .expect("get before overwrite") + .expect("present"); + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + store + .put(&ns_alice, "profile", json!({"name": "Alice", "v": 2})) + .await + .expect("overwrite alice profile"); + let after = store + .get(&ns_alice, "profile") + .await + .expect("get after overwrite") + .expect("present"); + assert_eq!( + after.created_at_ms, before.created_at_ms, + "overwrite preserves the item's original creation time" + ); + assert!( + after.updated_at_ms >= before.updated_at_ms, + "overwrite advances the item's update time" + ); + assert_eq!(after.value["v"], json!(2), "overwrite replaces the value"); + + // Namespaces are independent. + store + .put(&ns_bob, "profile", json!({"name": "Bob"})) + .await + .expect("put bob profile"); + let bob = store + .get(&ns_bob, "profile") + .await + .expect("get bob profile") + .expect("present"); + assert_eq!(bob.value["name"], json!("Bob")); + let alice_again = store + .get(&ns_alice, "profile") + .await + .expect("get alice profile again") + .expect("present"); + assert_eq!( + alice_again.value["name"], + json!("Alice"), + "writing bob's namespace does not affect alice's" + ); + + // search is namespace-prefix scoped: a broader prefix sees both items, a + // narrower one (a specific user) sees only its own. + let under_users = store + .search(SearchQuery { + namespace_prefix: vec!["users".to_string()], + ..Default::default() + }) + .await + .expect("search under the users prefix"); + assert_eq!( + under_users.len(), + 2, + "search finds every item at or beneath the namespace prefix" + ); + let alice_only = store + .search(SearchQuery { + namespace_prefix: vec!["users".to_string(), "alice".to_string()], + ..Default::default() + }) + .await + .expect("search scoped to alice"); + assert_eq!(alice_only.len(), 1, "a narrower prefix excludes bob's item"); + assert_eq!(alice_only[0].key, "profile"); + + // search applies the field filter as a conjunction. + let filtered = store + .search(SearchQuery { + namespace_prefix: vec!["users".to_string()], + filter: [("name".to_string(), FilterOp::Eq(json!("Bob")))] + .into_iter() + .collect(), + ..Default::default() + }) + .await + .expect("search with a field filter"); + assert_eq!( + filtered.len(), + 1, + "the filter selects only the matching item" + ); + assert_eq!(filtered[0].value["name"], json!("Bob")); + + // list_namespaces enumerates namespaces matching a prefix query. + let namespaces = store + .list_namespaces(ListNamespacesQuery { + prefix: Some(vec!["users".to_string()]), + ..Default::default() + }) + .await + .expect("list_namespaces under users"); + assert!( + namespaces.contains(&ns_alice), + "list_namespaces includes alice's namespace" + ); + assert!( + namespaces.contains(&ns_bob), + "list_namespaces includes bob's namespace" + ); + + // delete removes exactly the deleted item, and deleting an already-absent + // item is a no-op rather than an error. + store + .delete(&ns_alice, "profile") + .await + .expect("delete alice profile"); + assert!( + store + .get(&ns_alice, "profile") + .await + .expect("get after delete") + .is_none(), + "a deleted item reads back as None" + ); + store + .delete(&ns_alice, "profile") + .await + .expect("delete of an already-absent item must not error"); + + // TTL: an item written with a very short explicit lifetime is invisible + // to both get and search once it has expired, even though nothing ever + // called `sweep_expired` — expiry is enforced on read. + store + .put_with_ttl(&ns_bob, "ephemeral", json!("soon gone"), Some(0.0001)) + .await + .expect("put with a short ttl"); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert!( + store + .get(&ns_bob, "ephemeral") + .await + .expect("get an expired item") + .is_none(), + "an expired item is invisible to get before it is ever swept" + ); + let after_expiry = store + .search(SearchQuery { + namespace_prefix: vec!["users".to_string(), "bob".to_string()], + ..Default::default() + }) + .await + .expect("search after expiry"); + assert!( + after_expiry.iter().all(|found| found.key != "ephemeral"), + "search does not return an expired item" + ); + + // batch: several operations in one call return one result per operation, + // positionally aligned with the request — the contract every convenience + // method on the trait is built from. + let ns_batch = Namespace::new(["batch"]).expect("valid namespace"); + let ops = vec![ + StoreOp::Put { + namespace: ns_batch.clone(), + key: "one".to_string(), + value: Some(json!(1)), + ttl_minutes: None, + }, + StoreOp::Put { + namespace: ns_batch.clone(), + key: "two".to_string(), + value: Some(json!(2)), + ttl_minutes: None, + }, + StoreOp::Get { + namespace: ns_batch.clone(), + key: "one".to_string(), + refresh_ttl: None, + }, + StoreOp::Get { + namespace: ns_batch.clone(), + key: "two".to_string(), + refresh_ttl: None, + }, + ]; + let mut results = store.batch(&ops).await.expect("batch"); + assert_eq!( + results.len(), + 4, + "batch returns one result per submitted operation" + ); + let two = results + .remove(3) + .into_item() + .expect("op 3 is a Get") + .expect("present"); + let one = results + .remove(2) + .into_item() + .expect("op 2 is a Get") + .expect("present"); + assert_eq!( + (one.value, two.value), + (json!(1), json!(2)), + "batch results are positionally aligned with the request, not just in the same order" + ); +} diff --git a/crates/tinyagents-harness/src/store/mod.rs b/crates/tinyagents-harness/src/store/mod.rs index 7933bf97..38dbafd7 100644 --- a/crates/tinyagents-harness/src/store/mod.rs +++ b/crates/tinyagents-harness/src/store/mod.rs @@ -21,12 +21,16 @@ //! - [`InMemoryStore`] — ephemeral in-process store for tests and examples. //! - [`FileStore`] — file-system-backed store for local development. //! - [`StoreRegistry`] — named bag of stores injected into `RunContext`. +//! - [`conformance`] — contract suites (`run_store_conformance`, +//! `run_namespaced_store_conformance`) any backend of either trait can be +//! certified against. //! //! # Namespace convention //! Use slash-free, lowercase names like `"threads"`, `"events"`, `"cache"`, //! `"artifacts"`. The registry does not enforce a naming scheme, but //! consistent names make multi-store applications easier to audit. +pub mod conformance; pub mod namespaced; mod types; @@ -156,54 +160,67 @@ impl Store for FileStore { Self::sanitize(namespace)?; Self::sanitize(key)?; let path = self.key_path(namespace, key); - if !path.exists() { - return Ok(None); - } - let bytes = fs::read(&path) - .map_err(|e| TinyAgentsError::Validation(format!("store read error: {e}")))?; - let value: Value = serde_json::from_slice(&bytes)?; - Ok(Some(value)) + // Blocking file I/O; offload it via the shared `spawn_blocking` + // helper so a store read never stalls a tokio worker (see I-4). + crate::blocking::run_blocking(move || -> Result> { + if !path.exists() { + return Ok(None); + } + let bytes = fs::read(&path) + .map_err(|e| TinyAgentsError::Validation(format!("store read error: {e}")))?; + let value: Value = serde_json::from_slice(&bytes)?; + Ok(Some(value)) + }) + .await } async fn put(&self, namespace: &str, key: &str, value: Value) -> Result<()> { Self::sanitize(namespace)?; Self::sanitize(key)?; let dir = self.root_dir.join(namespace); - fs::create_dir_all(&dir) - .map_err(|e| TinyAgentsError::Validation(format!("store mkdir error: {e}")))?; - let path = dir.join(format!("{key}.json")); - let bytes = serde_json::to_vec_pretty(&value)?; - // Write to a uniquely named temp file in the same directory, then rename - // over the destination. Rename is atomic on POSIX/Windows for same-dir - // paths, so a reader never observes a partially written file and a crash - // mid-write leaves the previous value intact (as the type docs promise). - let tmp = dir.join(format!( - "{key}.json.tmp.{}.{}", - std::process::id(), - TMP_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed) - )); - fs::write(&tmp, &bytes) - .map_err(|e| TinyAgentsError::Validation(format!("store write error: {e}")))?; - if let Err(e) = fs::rename(&tmp, &path) { - // Best-effort cleanup of the temp file so a failed rename does not - // leak partial files into the namespace directory. - let _ = fs::remove_file(&tmp); - return Err(TinyAgentsError::Validation(format!( - "store rename error: {e}" - ))); - } - Ok(()) + let key = key.to_string(); + crate::blocking::run_blocking(move || -> Result<()> { + fs::create_dir_all(&dir) + .map_err(|e| TinyAgentsError::Validation(format!("store mkdir error: {e}")))?; + let path = dir.join(format!("{key}.json")); + let bytes = serde_json::to_vec_pretty(&value)?; + // Write to a uniquely named temp file in the same directory, then + // rename over the destination. Rename is atomic on POSIX/Windows + // for same-dir paths, so a reader never observes a partially + // written file and a crash mid-write leaves the previous value + // intact (as the type docs promise). + let tmp = dir.join(format!( + "{key}.json.tmp.{}.{}", + std::process::id(), + TMP_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + )); + fs::write(&tmp, &bytes) + .map_err(|e| TinyAgentsError::Validation(format!("store write error: {e}")))?; + if let Err(e) = fs::rename(&tmp, &path) { + // Best-effort cleanup of the temp file so a failed rename does + // not leak partial files into the namespace directory. + let _ = fs::remove_file(&tmp); + return Err(TinyAgentsError::Validation(format!( + "store rename error: {e}" + ))); + } + Ok(()) + }) + .await } async fn delete(&self, namespace: &str, key: &str) -> Result<()> { Self::sanitize(namespace)?; Self::sanitize(key)?; let path = self.key_path(namespace, key); - if path.exists() { - fs::remove_file(&path) - .map_err(|e| TinyAgentsError::Validation(format!("store delete error: {e}")))?; - } - Ok(()) + crate::blocking::run_blocking(move || -> Result<()> { + if path.exists() { + fs::remove_file(&path) + .map_err(|e| TinyAgentsError::Validation(format!("store delete error: {e}")))?; + } + Ok(()) + }) + .await } async fn list(&self, namespace: &str) -> Result> { @@ -513,12 +530,7 @@ impl AppendStore for JsonlAppendStore { Ok(offset) }; - match tokio::runtime::Handle::try_current() { - Ok(handle) => handle.spawn_blocking(work).await.map_err(|e| { - TinyAgentsError::Validation(format!("append store task error: {e}")) - })?, - Err(_) => work(), - } + crate::blocking::run_blocking(work).await } async fn read_from(&self, stream: &str, offset: u64) -> Result> { diff --git a/crates/tinyagents-harness/src/store/namespaced/mod.rs b/crates/tinyagents-harness/src/store/namespaced/mod.rs index fb78d4e4..e2ec37af 100644 --- a/crates/tinyagents-harness/src/store/namespaced/mod.rs +++ b/crates/tinyagents-harness/src/store/namespaced/mod.rs @@ -249,7 +249,7 @@ impl NamespacedStore for InMemoryNamespacedStore { items.retain(|_, item| !item.is_expired(now)); let reclaimed = before - items.len(); if reclaimed > 0 { - tinyagents_tracing::debug!("[store:namespaced] sweep_expired reclaimed={reclaimed}"); + tracing::debug!("[store:namespaced] sweep_expired reclaimed={reclaimed}"); } Ok(reclaimed) } diff --git a/crates/tinyagents-harness/src/stream/frame.rs b/crates/tinyagents-harness/src/stream/frame.rs new file mode 100644 index 00000000..08b75646 --- /dev/null +++ b/crates/tinyagents-harness/src/stream/frame.rs @@ -0,0 +1,429 @@ +//! Durable frame codec for assistant-message streaming. +//! +//! [`ModelStreamItem`]s are the wire-level shape a provider adapter emits; +//! they are not, on their own, durable — a consumer that reconnects mid-turn +//! (or a journal reader replaying a crashed run) needs a compact, self +//! describing record it can persist and fold back into a partial message +//! without re-running the provider stream. +//! +//! [`AssistantFrame`] is that record. [`FrameEncoder`] turns a sequence of +//! [`ModelStreamItem`]s into frames (emitting a periodic +//! [`AssistantFrame::ToolArgsCheckpoint`] for long-running tool-argument +//! streams so a reconnecting reader does not have to replay every single +//! fragment from the start of the block); [`reduce_frames`] folds frames back +//! into a [`PartialAssistantMessage`] — the harness event journal persists +//! frames as they are encoded, and a crashed/reconnecting consumer rebuilds +//! its view by reducing whatever frames it has, including a sequence +//! truncated mid-block. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use tinyinference_llm::message::{AssistantMessage, ContentBlock}; +use tinyinference_llm::model::{BlockDelta, BlockKind, ModelStreamItem}; +use tinyinference_llm::tool::ToolCall; +use tinyinference_llm::usage::Usage; + +/// Number of tool-argument fragments accumulated between automatic +/// [`AssistantFrame::ToolArgsCheckpoint`] snapshots. +/// +/// A checkpoint is a full snapshot (not a delta), so a reader that only has +/// frames from the checkpoint onward — because earlier per-fragment frames +/// were compacted out of the journal — still reduces to the correct partial +/// argument string. +const TOOL_ARGS_CHECKPOINT_INTERVAL: usize = 16; + +/// A compact, self-describing, serializable record of one increment of an +/// in-progress assistant message stream. +/// +/// Frames are the unit the harness event journal persists for a streaming +/// model call. Reducing a sequence of frames with [`reduce_frames`] +/// reconstructs a [`PartialAssistantMessage`] without needing the original +/// provider stream. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case", tag = "type", content = "content")] +pub enum AssistantFrame { + /// A new content block has opened at `index`. + BlockStart { + /// Position of the block within the assistant message. + index: usize, + /// The block's syntactic category. + kind: BlockKind, + }, + /// An incremental fragment for the open block at `index`. + BlockDelta { + /// Position of the block this fragment belongs to. + index: usize, + /// The fragment payload. + delta: BlockDelta, + }, + /// A full snapshot of a tool-call block's accumulated argument JSON, + /// emitted periodically (every [`TOOL_ARGS_CHECKPOINT_INTERVAL`] + /// fragments) so a reader with a truncated frame log can still recover a + /// consistent partial argument string. + ToolArgsCheckpoint { + /// Position of the tool-call block this checkpoint snapshots. + index: usize, + /// The full argument JSON accumulated for this block so far. + json_so_far: String, + }, + /// The block at `index` has closed; `block` is its fully assembled + /// content. + BlockEnd { + /// Position of the closed block. + index: usize, + /// The finished content block. + block: ContentBlock, + }, + /// A usage update. + Usage(Usage), + /// Terminal success: the fully merged response. + Completed { + /// The complete assistant message. + message: AssistantMessage, + /// The provider's reported stop/finish reason, when known. + #[serde(default, skip_serializing_if = "Option::is_none")] + stop_reason: Option, + }, + /// Terminal failure. Carries whatever partial message had accumulated + /// before the failure, mirroring + /// [`tinyinference_llm::model::ProviderError::partial_message`]. + Failed { + /// Human-readable failure message. + message: String, + /// The assistant message accumulated before the failure, when any + /// content had arrived. + #[serde(default, skip_serializing_if = "Option::is_none")] + partial: Option, + /// The stop/finish reason reported before the failure, when known. + #[serde(default, skip_serializing_if = "Option::is_none")] + stop_reason: Option, + }, +} + +// --------------------------------------------------------------------------- +// FrameEncoder +// --------------------------------------------------------------------------- + +/// Turns a sequence of [`ModelStreamItem`]s into durable [`AssistantFrame`]s. +/// +/// Feed items with [`FrameEncoder::push`] as a provider stream produces them; +/// call [`FrameEncoder::into_frames`] (or read [`FrameEncoder::frames`] +/// incrementally) to get the encoded sequence. [`ModelStreamItem::Started`], +/// [`ModelStreamItem::MessageDelta`], and [`ModelStreamItem::ToolCallDelta`] +/// carry no information a block-aware reducer needs beyond what +/// `BlockStart`/`BlockDelta`/`BlockEnd` already carry, so they are not framed +/// — only the block-indexed and terminal items are. +#[derive(Debug, Default)] +pub struct FrameEncoder { + frames: Vec, + /// Per-block running tool-argument JSON and fragment count since the + /// last checkpoint, keyed by block index. + tool_progress: BTreeMap, +} + +impl FrameEncoder { + /// Creates an empty encoder. + pub fn new() -> Self { + Self::default() + } + + /// Folds one stream item, appending zero or more frames. + pub fn push(&mut self, item: &ModelStreamItem) { + match item { + ModelStreamItem::BlockStart { index, kind } => { + if matches!(kind, BlockKind::ToolCall { .. }) { + self.tool_progress.insert(*index, (String::new(), 0)); + } + self.frames.push(AssistantFrame::BlockStart { + index: *index, + kind: kind.clone(), + }); + } + ModelStreamItem::BlockDelta { index, delta } => { + self.frames.push(AssistantFrame::BlockDelta { + index: *index, + delta: delta.clone(), + }); + if let BlockDelta::ToolArgs(fragment) = delta + && let Some((json_so_far, count)) = self.tool_progress.get_mut(index) + { + json_so_far.push_str(fragment); + *count += 1; + if *count >= TOOL_ARGS_CHECKPOINT_INTERVAL { + self.frames.push(AssistantFrame::ToolArgsCheckpoint { + index: *index, + json_so_far: json_so_far.clone(), + }); + *count = 0; + } + } + } + ModelStreamItem::BlockEnd { index, block } => { + self.tool_progress.remove(index); + self.frames.push(AssistantFrame::BlockEnd { + index: *index, + block: block.clone(), + }); + } + ModelStreamItem::UsageDelta(usage) => { + self.frames.push(AssistantFrame::Usage(*usage)); + } + ModelStreamItem::Completed(response) => { + self.frames.push(AssistantFrame::Completed { + message: response.message.clone(), + stop_reason: response.finish_reason.clone(), + }); + } + ModelStreamItem::Failed(message) => { + self.frames.push(AssistantFrame::Failed { + message: message.clone(), + partial: None, + stop_reason: None, + }); + } + ModelStreamItem::ProviderFailed(error) => { + self.frames.push(AssistantFrame::Failed { + message: error.message.clone(), + partial: error.partial_message.clone(), + stop_reason: error.stop_reason.clone(), + }); + } + // No block-boundary information; the compatibility channel is + // fully covered by the block-indexed items above for any + // block-aware adapter. Adapters that only emit the flat + // `MessageDelta`/`ToolCallDelta` shape (no block boundaries) have + // nothing durable to frame here beyond what `Completed`/`Failed` + // already capture. `Deferred` carries no message content either + // (it is a handle to a response that will resolve later via + // `ChatModel::fetch_deferred`, folded by + // `StreamAccumulator::deferred` instead of this block reducer), + // so it is likewise not framed. + ModelStreamItem::Started + | ModelStreamItem::MessageDelta(_) + | ModelStreamItem::ToolCallDelta(_) + | ModelStreamItem::Deferred(_) => {} + } + } + + /// Returns the frames encoded so far without consuming the encoder. + pub fn frames(&self) -> &[AssistantFrame] { + &self.frames + } + + /// Consumes the encoder and returns the full encoded frame sequence. + pub fn into_frames(self) -> Vec { + self.frames + } +} + +/// Encodes a complete slice of [`ModelStreamItem`]s into [`AssistantFrame`]s. +/// +/// A convenience wrapper around [`FrameEncoder`] for callers that already +/// have the full item sequence (tests, post-processing). +pub fn encode_frames(items: &[ModelStreamItem]) -> Vec { + let mut encoder = FrameEncoder::new(); + for item in items { + encoder.push(item); + } + encoder.into_frames() +} + +// --------------------------------------------------------------------------- +// PartialAssistantMessage / reduce_frames +// --------------------------------------------------------------------------- + +/// A block still open (no [`AssistantFrame::BlockEnd`] seen yet) while +/// reducing frames. +#[derive(Clone, Debug, PartialEq)] +enum OpenBlock { + Text(String), + Thinking(String), + ToolCall { + id: Option, + name: Option, + json_so_far: String, + }, +} + +/// The result of folding a (possibly truncated) [`AssistantFrame`] sequence. +/// +/// Reflects exactly what the frames folded in describe: closed blocks are +/// merged into [`Self::content`] in index order, and a block with no +/// [`AssistantFrame::BlockEnd`] yet is exposed via [`Self::open_blocks`] so a +/// reconnecting consumer can still render in-progress content. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct PartialAssistantMessage { + /// Closed content blocks, in block-index order. + pub content: Vec, + /// Blocks that opened but have not closed, as `(index, text_so_far)` for + /// text/thinking blocks or `(index, json_so_far)` for tool-call blocks, + /// in block-index order. + pub open_blocks: Vec<(usize, String)>, + /// Tool calls reconstructed from closed tool-use blocks. Malformed JSON + /// becomes [`ToolCall::invalid`] rather than being dropped. + pub tool_calls: Vec, + /// Most recent usage value seen. + pub usage: Option, + /// Present once a terminal [`AssistantFrame::Completed`] or + /// [`AssistantFrame::Failed`] frame has been folded in. + pub terminal: Option, +} + +/// The terminal outcome folded into a [`PartialAssistantMessage`], when any. +#[derive(Clone, Debug, PartialEq)] +pub enum PartialTerminal { + /// The stream completed successfully; carries the authoritative message. + Completed { + /// The complete assistant message. + message: AssistantMessage, + /// The provider's reported stop/finish reason, when known. + stop_reason: Option, + }, + /// The stream failed; carries the human-readable message and, when the + /// failure was mid-stream, the partial message and stop reason it + /// interrupted. + Failed { + /// Human-readable failure message. + message: String, + /// The assistant message accumulated before the failure, when any. + partial: Option, + /// The stop/finish reason reported before the failure, when known. + stop_reason: Option, + }, +} + +/// Folds a (possibly truncated) [`AssistantFrame`] sequence into a +/// [`PartialAssistantMessage`]. +/// +/// A full sequence — one that ends in [`AssistantFrame::Completed`] or +/// [`AssistantFrame::Failed`] — reduces to a result whose `content` (in the +/// `Completed` case) matches the original [`AssistantMessage`]. A sequence +/// truncated mid-block reduces to a consistent partial: every fully closed +/// block lands in `content`, and the interrupted block's accumulated text or +/// argument JSON (using the most recent [`AssistantFrame::ToolArgsCheckpoint`] +/// as its base, when one was folded in) is exposed via `open_blocks`. +pub fn reduce_frames(frames: &[AssistantFrame]) -> PartialAssistantMessage { + let mut open: BTreeMap = BTreeMap::new(); + let mut closed: BTreeMap = BTreeMap::new(); + let mut tool_calls = Vec::new(); + let mut usage = None; + let mut terminal = None; + + for frame in frames { + match frame { + AssistantFrame::BlockStart { index, kind } => { + let block = match kind { + BlockKind::Text => OpenBlock::Text(String::new()), + BlockKind::Thinking => OpenBlock::Thinking(String::new()), + BlockKind::ToolCall { id, name } => OpenBlock::ToolCall { + id: Some(id.clone()), + name: Some(name.clone()), + json_so_far: String::new(), + }, + }; + open.insert(*index, block); + } + AssistantFrame::BlockDelta { index, delta } => match (open.get_mut(index), delta) { + (Some(OpenBlock::Text(text)), BlockDelta::Text(fragment)) => { + text.push_str(fragment); + } + (Some(OpenBlock::Thinking(text)), BlockDelta::Thinking(fragment)) => { + text.push_str(fragment); + } + (Some(OpenBlock::ToolCall { json_so_far, .. }), BlockDelta::ToolArgs(fragment)) => { + json_so_far.push_str(fragment); + } + _ => {} + }, + AssistantFrame::ToolArgsCheckpoint { index, json_so_far } => { + // A checkpoint is a full snapshot, not a delta: it replaces + // whatever was accumulated so far. A reader that only has + // frames from this checkpoint onward (its matching + // `BlockStart` was pruned from the journal) still needs a + // consistent partial, so a missing entry is created here + // rather than the checkpoint being silently dropped. + match open.get_mut(index) { + Some(OpenBlock::ToolCall { + json_so_far: current, + .. + }) => current.clone_from(json_so_far), + _ => { + open.insert( + *index, + OpenBlock::ToolCall { + id: None, + name: None, + json_so_far: json_so_far.clone(), + }, + ); + } + } + } + AssistantFrame::BlockEnd { index, block } => { + open.remove(index); + if let ContentBlock::Json(value) = &block + && let (Some(id), Some(name)) = ( + value.get("id").and_then(serde_json::Value::as_str), + value.get("name").and_then(serde_json::Value::as_str), + ) + { + let arguments = value + .get("arguments") + .cloned() + .unwrap_or(serde_json::Value::Null); + tool_calls.push(ToolCall::new(id, name, arguments)); + } else { + closed.insert(*index, block.clone()); + } + } + AssistantFrame::Usage(value) => { + usage = Some(*value); + } + AssistantFrame::Completed { + message, + stop_reason, + } => { + terminal = Some(PartialTerminal::Completed { + message: message.clone(), + stop_reason: stop_reason.clone(), + }); + } + AssistantFrame::Failed { + message, + partial, + stop_reason, + } => { + terminal = Some(PartialTerminal::Failed { + message: message.clone(), + partial: partial.clone(), + stop_reason: stop_reason.clone(), + }); + } + } + } + + let content = closed.into_values().collect(); + let open_blocks = open + .into_iter() + .map(|(index, block)| { + let text = match block { + OpenBlock::Text(text) | OpenBlock::Thinking(text) => text, + OpenBlock::ToolCall { json_so_far, .. } => json_so_far, + }; + (index, text) + }) + .collect(); + + PartialAssistantMessage { + content, + open_blocks, + tool_calls, + usage, + terminal, + } +} + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-harness/src/stream/frame/test.rs b/crates/tinyagents-harness/src/stream/frame/test.rs new file mode 100644 index 00000000..e17a6631 --- /dev/null +++ b/crates/tinyagents-harness/src/stream/frame/test.rs @@ -0,0 +1,285 @@ +use serde_json::json; + +use tinyinference_llm::message::{AssistantMessage, ContentBlock}; +use tinyinference_llm::model::{BlockDelta, BlockKind, ModelResponse, ModelStreamItem}; +use tinyinference_llm::usage::Usage; + +use super::*; + +fn interleaved_stream_items() -> Vec { + vec![ + ModelStreamItem::Started, + ModelStreamItem::BlockStart { + index: 0, + kind: BlockKind::Thinking, + }, + ModelStreamItem::BlockDelta { + index: 0, + delta: BlockDelta::Thinking("plan".into()), + }, + ModelStreamItem::BlockEnd { + index: 0, + block: ContentBlock::Thinking { + text: "plan".into(), + signature: None, + }, + }, + ModelStreamItem::BlockStart { + index: 1, + kind: BlockKind::Text, + }, + ModelStreamItem::BlockDelta { + index: 1, + delta: BlockDelta::Text("hel".into()), + }, + ModelStreamItem::BlockDelta { + index: 1, + delta: BlockDelta::Text("lo".into()), + }, + ModelStreamItem::BlockEnd { + index: 1, + block: ContentBlock::Text("hello".into()), + }, + ModelStreamItem::BlockStart { + index: 2, + kind: BlockKind::ToolCall { + id: "call-1".into(), + name: "search".into(), + }, + }, + ModelStreamItem::BlockDelta { + index: 2, + delta: BlockDelta::ToolArgs("{\"q\":".into()), + }, + ModelStreamItem::BlockDelta { + index: 2, + delta: BlockDelta::ToolArgs("1}".into()), + }, + ModelStreamItem::BlockEnd { + index: 2, + block: ContentBlock::Json( + json!({"id": "call-1", "name": "search", "arguments": {"q": 1}}), + ), + }, + ModelStreamItem::UsageDelta(Usage::new(5, 7)), + ModelStreamItem::Completed(ModelResponse { + message: AssistantMessage { + id: Some("msg-1".into()), + content: vec![ + ContentBlock::Thinking { + text: "plan".into(), + signature: None, + }, + ContentBlock::Text("hello".into()), + ], + tool_calls: vec![tinyinference_llm::tool::ToolCall::new( + "call-1", + "search", + json!({"q": 1}), + )], + usage: Some(Usage::new(5, 7)), + origin: None, + }, + usage: Some(Usage::new(5, 7)), + finish_reason: Some("tool_use".into()), + raw: None, + resolved_model: None, + continue_turn: None, + served_from_cache: false, + correlation: None, + resolved_route: None, + }), + ] +} + +#[test] +fn encode_skips_flat_compatibility_items_and_frames_block_items() { + let items = vec![ + ModelStreamItem::Started, + ModelStreamItem::BlockStart { + index: 0, + kind: BlockKind::Text, + }, + ModelStreamItem::BlockDelta { + index: 0, + delta: BlockDelta::Text("hi".into()), + }, + ModelStreamItem::MessageDelta(tinyinference_llm::message::MessageDelta::text("hi")), + ]; + let frames = encode_frames(&items); + assert_eq!( + frames, + vec![ + AssistantFrame::BlockStart { + index: 0, + kind: BlockKind::Text, + }, + AssistantFrame::BlockDelta { + index: 0, + delta: BlockDelta::Text("hi".into()), + }, + ] + ); +} + +#[test] +fn encode_then_reduce_round_trips_to_the_terminal_message() { + let items = interleaved_stream_items(); + let frames = encode_frames(&items); + let partial = reduce_frames(&frames); + + assert!(partial.open_blocks.is_empty(), "every block closed"); + assert_eq!( + partial.content, + vec![ + ContentBlock::Thinking { + text: "plan".into(), + signature: None, + }, + ContentBlock::Text("hello".into()), + ] + ); + assert_eq!(partial.tool_calls.len(), 1); + assert_eq!(partial.tool_calls[0].id, "call-1"); + assert_eq!(partial.tool_calls[0].name, "search"); + assert_eq!(partial.tool_calls[0].arguments, json!({"q": 1})); + assert_eq!(partial.usage, Some(Usage::new(5, 7))); + + let Some(PartialTerminal::Completed { + message, + stop_reason, + }) = partial.terminal + else { + panic!("expected a Completed terminal"); + }; + assert_eq!(stop_reason.as_deref(), Some("tool_use")); + assert_eq!( + tinyinference_llm::message::Message::Assistant(message.clone()).text(), + "hello" + ); + assert_eq!(message.tool_calls.len(), 1); +} + +#[test] +fn reduce_of_a_truncated_sequence_yields_a_consistent_partial() { + // Drop everything from the tool-call block's argument deltas onward: no + // BlockEnd, no Completed. The reducer must still expose the closed + // thinking/text blocks and the in-progress tool-call argument string. + let items = interleaved_stream_items(); + let frames = encode_frames(&items); + let cut = frames + .iter() + .position(|frame| matches!(frame, AssistantFrame::BlockDelta { index: 2, .. })) + .expect("a block-2 delta frame"); + let truncated = &frames[..=cut]; + let partial = reduce_frames(truncated); + + assert_eq!( + partial.content, + vec![ + ContentBlock::Thinking { + text: "plan".into(), + signature: None, + }, + ContentBlock::Text("hello".into()), + ], + "closed blocks are unaffected by the truncation" + ); + assert!(partial.tool_calls.is_empty(), "tool block never closed"); + assert_eq!(partial.terminal, None); + assert_eq!(partial.open_blocks, vec![(2, "{\"q\":".to_string())]); +} + +#[test] +fn checkpoint_is_a_full_snapshot_not_a_delta() { + let frames = vec![ + AssistantFrame::BlockStart { + index: 0, + kind: BlockKind::ToolCall { + id: "call-1".into(), + name: "search".into(), + }, + }, + AssistantFrame::BlockDelta { + index: 0, + delta: BlockDelta::ToolArgs("{\"q\":1".into()), + }, + // A checkpoint replaces the accumulated string; a reader that only + // sees frames from here onward must reduce to the same result as one + // that saw every fragment. + AssistantFrame::ToolArgsCheckpoint { + index: 0, + json_so_far: "{\"q\":1".into(), + }, + AssistantFrame::BlockDelta { + index: 0, + delta: BlockDelta::ToolArgs("}".into()), + }, + ]; + let full = reduce_frames(&frames); + let truncated = reduce_frames(&frames[2..]); + assert_eq!(full.open_blocks, vec![(0, "{\"q\":1}".to_string())]); + assert_eq!(full.open_blocks, truncated.open_blocks); +} + +#[test] +fn automatic_checkpoints_appear_every_configured_interval() { + let mut encoder = FrameEncoder::new(); + encoder.push(&ModelStreamItem::BlockStart { + index: 0, + kind: BlockKind::ToolCall { + id: "call-1".into(), + name: "search".into(), + }, + }); + for _ in 0..TOOL_ARGS_CHECKPOINT_INTERVAL { + encoder.push(&ModelStreamItem::BlockDelta { + index: 0, + delta: BlockDelta::ToolArgs("a".into()), + }); + } + let frames = encoder.into_frames(); + let checkpoint = frames + .iter() + .find_map(|frame| match frame { + AssistantFrame::ToolArgsCheckpoint { json_so_far, .. } => Some(json_so_far.clone()), + _ => None, + }) + .expect("a checkpoint frame after the configured interval"); + assert_eq!(checkpoint, "a".repeat(TOOL_ARGS_CHECKPOINT_INTERVAL)); +} + +#[test] +fn provider_failed_frame_carries_partial_message_and_stop_reason() { + let items = vec![ModelStreamItem::ProviderFailed( + tinyinference_llm::model::ProviderError { + provider: "anthropic".into(), + message: "overloaded".into(), + stop_reason: Some("pause_turn".into()), + partial_message: Some(AssistantMessage { + id: None, + content: vec![ContentBlock::Text("partial".into())], + tool_calls: Vec::new(), + usage: None, + origin: None, + }), + ..Default::default() + }, + )]; + let frames = encode_frames(&items); + let partial = reduce_frames(&frames); + let Some(PartialTerminal::Failed { + message, + partial: partial_message, + stop_reason, + }) = partial.terminal + else { + panic!("expected a Failed terminal"); + }; + assert_eq!(message, "overloaded"); + assert_eq!(stop_reason.as_deref(), Some("pause_turn")); + assert_eq!( + partial_message.unwrap().content, + vec![ContentBlock::Text("partial".into())] + ); +} diff --git a/crates/tinyagents-harness/src/stream/mod.rs b/crates/tinyagents-harness/src/stream/mod.rs index 5e163391..7ae341c8 100644 --- a/crates/tinyagents-harness/src/stream/mod.rs +++ b/crates/tinyagents-harness/src/stream/mod.rs @@ -34,9 +34,14 @@ //! one mode the projection cannot supply — a full state snapshot is graph //! state, so the graph runtime pushes [`StreamChunk::Values`] itself. +pub mod frame; mod project; mod types; +pub use frame::{ + AssistantFrame, FrameEncoder, PartialAssistantMessage, PartialTerminal, encode_frames, + reduce_frames, +}; pub use project::{project_event, project_event_for_modes, projected_mode}; pub use types::*; diff --git a/crates/tinyagents-harness/src/structured/mod.rs b/crates/tinyagents-harness/src/structured/mod.rs index d23f84c8..2ba59963 100644 --- a/crates/tinyagents-harness/src/structured/mod.rs +++ b/crates/tinyagents-harness/src/structured/mod.rs @@ -64,12 +64,72 @@ mod validate; pub use repair::JsonRepair; pub use types::*; +use async_trait::async_trait; use serde::de::DeserializeOwned; use serde_json::Value; +use crate::context::RunContext; use crate::error::{Result, TinyAgentsError}; use tinyinference_llm::model::{ModelProfile, ModelResponse, ResponseFormat}; +// --------------------------------------------------------------------------- +// OutputValidator +// --------------------------------------------------------------------------- + +/// Validates an already schema-valid structured output, driving the +/// output-validation retry loop (A3, `RunPolicy::output_retry`). +/// +/// Registered on a harness via +/// [`crate::runtime::AgentHarness::with_output_validator`]. Called once per +/// final-turn extraction, after [`StructuredExtractor::extract_outcome`] +/// already succeeded — a schema-invalid value never reaches the validator; it +/// retries through the same loop for the extraction-failure reason instead. +/// +/// Returning `Err(TinyAgentsError::ModelRetry(message))` asks the agent loop +/// to push `message` back to the model as a repair prompt and try again +/// (bounded by [`crate::runtime::RunPolicy::output_retry`]'s +/// `max_attempts`); any other `Err` variant fails the run immediately, +/// exactly like an error from any other fallible call in the loop. Mirrors +/// Pydantic AI's `@agent.output_validator`. +/// +/// # Example +/// +/// ```rust +/// use async_trait::async_trait; +/// use tinyagents_harness::context::RunContext; +/// use tinyagents_harness::error::{Result, TinyAgentsError}; +/// use tinyagents_harness::structured::OutputValidator; +/// +/// struct NonEmpty; +/// +/// #[async_trait] +/// impl OutputValidator<()> for NonEmpty { +/// async fn validate( +/// &self, +/// _ctx: &mut RunContext<()>, +/// _state: &(), +/// output: &serde_json::Value, +/// ) -> Result<()> { +/// if output.get("answer").and_then(|v| v.as_str()).is_none_or(str::is_empty) { +/// return Err(TinyAgentsError::ModelRetry( +/// "`answer` must be a non-empty string".to_string(), +/// )); +/// } +/// Ok(()) +/// } +/// } +/// ``` +#[async_trait] +pub trait OutputValidator: Send + Sync { + /// Validates `output`. See the trait docs for how `Err` is handled. + async fn validate( + &self, + ctx: &mut RunContext, + state: &State, + output: &Value, + ) -> Result<()>; +} + // --------------------------------------------------------------------------- // Strategy selection // --------------------------------------------------------------------------- @@ -133,6 +193,23 @@ impl StructuredStrategy { /// ); /// ``` pub fn for_profile(profile: Option<&ModelProfile>) -> StructuredStrategy { + // A profile's `default_structured_mode` is an explicit authoring + // decision about *this* model — it wins over the generic + // capability-based inference below, which only guesses from + // `native_structured_output`/`tool_calling`. + if let Some(mode) = profile.and_then(|p| p.default_structured_mode) { + return match mode { + tinyinference_llm::model::StructuredMode::Native => { + StructuredStrategy::ProviderSchema + } + tinyinference_llm::model::StructuredMode::Tool => StructuredStrategy::ToolCall, + tinyinference_llm::model::StructuredMode::Prompted => { + StructuredStrategy::Prompted { + template: profile.and_then(|p| p.prompted_output_template.clone()), + } + } + }; + } match profile { Some(p) if p.native_structured_output && p.json_schema => { StructuredStrategy::ProviderSchema @@ -190,17 +267,40 @@ impl StructuredExtractor { strategy, schema_name: schema_name.into(), schema, + variants: Vec::new(), + } + } + + /// Creates a [`StructuredStrategy::ToolCallUnion`] extractor over + /// `variants` (`(name, schema)` pairs, one per synthetic tool). + /// + /// `schema_name` is used only to label errors when *no* variant matched; + /// it need not be one of the variant names. + pub fn new_union(schema_name: impl Into, variants: Vec<(String, Value)>) -> Self { + Self { + strategy: StructuredStrategy::ToolCallUnion, + schema_name: schema_name.into(), + schema: Value::Null, + variants, } } /// Returns the JSON Schema document this extractor was configured with. /// /// Used for local validation and for echoing the schema back into a - /// [`ResponseFormat`] when re-requesting structured output. + /// [`ResponseFormat`] when re-requesting structured output. Meaningless + /// for [`StructuredStrategy::ToolCallUnion`] (use [`Self::variants`]). pub fn schema(&self) -> &Value { &self.schema } + /// Returns the `(name, schema)` variants this + /// [`StructuredStrategy::ToolCallUnion`] extractor was configured with. + /// Empty for every other strategy. + pub fn variants(&self) -> &[(String, Value)] { + &self.variants + } + /// Extracts a [`StructuredOutput`] from `response` using the configured /// strategy. /// @@ -227,12 +327,19 @@ impl StructuredExtractor { /// /// See strategy descriptions above. pub fn extract(&self, response: &ModelResponse) -> Result { - let output = match self.strategy { - StructuredStrategy::ProviderSchema => self.extract_provider_schema(response)?, - StructuredStrategy::ToolCall => self.extract_tool_call(response)?, - }; - validate::validate_value(&self.schema, &output.value, &self.instance_root())?; - Ok(output) + match &self.strategy { + StructuredStrategy::ProviderSchema | StructuredStrategy::Prompted { .. } => { + let output = self.extract_provider_schema(response)?; + validate::validate_value(&self.schema, &output.value, &self.instance_root())?; + Ok(output) + } + StructuredStrategy::ToolCall => { + let output = self.extract_tool_call(response)?; + validate::validate_value(&self.schema, &output.value, &self.instance_root())?; + Ok(output) + } + StructuredStrategy::ToolCallUnion => self.extract_tool_call_union(response), + } } /// Extracts without failing: records the error instead of raising it. @@ -253,10 +360,11 @@ impl StructuredExtractor { value: Some(output.value), raw: response.clone(), error: None, + variant: output.variant, }, Err(error) => { let error = error.to_string(); - tinyagents_tracing::debug!( + tracing::debug!( "[structured] extraction failed for schema '{}': {error}", self.schema_name ); @@ -264,6 +372,7 @@ impl StructuredExtractor { value: None, raw: response.clone(), error: Some(error), + variant: None, } } } @@ -322,7 +431,7 @@ impl StructuredExtractor { ))); }; if repair.is_repaired() { - tinyagents_tracing::debug!( + tracing::debug!( "[structured] schema '{}': recovered the value with repair `{}`", self.schema_name, repair.as_str() @@ -331,6 +440,7 @@ impl StructuredExtractor { Ok(StructuredOutput { value, raw_text: Some(raw), + variant: None, }) } @@ -358,7 +468,7 @@ impl StructuredExtractor { && let Some((value, repair)) = repair::parse_lenient(raw) { if repair.is_repaired() { - tinyagents_tracing::debug!( + tracing::debug!( "[structured] schema '{}': recovered tool-call arguments with repair `{}`", self.schema_name, repair.as_str() @@ -367,12 +477,67 @@ impl StructuredExtractor { return Ok(StructuredOutput { value, raw_text: Some(raw.to_string()), + variant: None, }); } Ok(StructuredOutput { value: call.arguments.clone(), raw_text: None, + variant: None, + }) + } + + /// [`StructuredStrategy::ToolCallUnion`] extraction: scans the response's + /// tool calls for the first one whose name matches a variant, validates + /// its arguments against *that variant's* schema (running the same + /// repair ladder [`Self::extract_tool_call`] does for unparseable + /// provider arguments), and records the matched variant name. + fn extract_tool_call_union(&self, response: &ModelResponse) -> Result { + let variant_names: Vec<&str> = self.variants.iter().map(|(n, _)| n.as_str()).collect(); + let call = response + .tool_calls() + .iter() + .find(|tc| variant_names.contains(&tc.name.as_str())) + .ok_or_else(|| { + TinyAgentsError::Validation(format!( + "schema '{}': no tool call matching any of the union's variants {:?} was \ + found in response", + self.schema_name, variant_names + )) + })?; + let (variant_name, variant_schema) = self + .variants + .iter() + .find(|(name, _)| name == &call.name) + .expect("matched call name came from variant_names"); + + let (value, raw_text) = if let Some(raw) = call.arguments.as_str() + && let Some((value, repair)) = repair::parse_lenient(raw) + { + if repair.is_repaired() { + tracing::debug!( + "[structured] union variant '{}': recovered tool-call arguments with \ + repair `{}`", + variant_name, + repair.as_str() + ); + } + (value, Some(raw.to_string())) + } else { + (call.arguments.clone(), None) + }; + + validate::validate_value( + variant_schema, + &value, + &format!("union variant '{variant_name}'"), + )?; + + Ok(StructuredOutput { + value, + raw_text, + variant: Some(variant_name.clone()), }) } } @@ -402,9 +567,28 @@ pub fn response_format_for_strategy( ) -> ResponseFormat { match strategy { StructuredStrategy::ProviderSchema => ResponseFormat::json_schema(name, schema), - StructuredStrategy::ToolCall => ResponseFormat::Text, + // Both send the structure through a channel other than the + // provider's native schema field: `ToolCall` through a forced tool + // call, `Prompted` through instructions plus free text the repair + // ladder parses. Neither wants the provider attempting its own + // (possibly conflicting) schema enforcement on top. + StructuredStrategy::ToolCall + | StructuredStrategy::Prompted { .. } + | StructuredStrategy::ToolCallUnion => ResponseFormat::Text, } } +/// The default instructions [`StructuredStrategy::Prompted`] injects ahead of +/// the schema when no custom `template` is configured. +/// +/// Mirrors Pydantic AI's `PromptedOutput` default wording: state the +/// requirement, then let the schema (appended separately by the caller) speak +/// for itself. +pub fn default_prompted_template() -> &'static str { + "Respond with a single JSON object that conforms exactly to this JSON Schema. \ + Do not include any text before or after the JSON object, and do not wrap it in \ + a code fence." +} + #[cfg(test)] mod test; diff --git a/crates/tinyagents-harness/src/structured/repair.rs b/crates/tinyagents-harness/src/structured/repair.rs index 4f4e7122..4e0a7a64 100644 --- a/crates/tinyagents-harness/src/structured/repair.rs +++ b/crates/tinyagents-harness/src/structured/repair.rs @@ -6,10 +6,12 @@ //! assistant text, and a single malformed brace on the final turn discarded the //! whole run — every tool call and token already spent. Meanwhile the crate //! already carried a repair ladder for the *other* place a model emits JSON: -//! tool-call arguments, repaired by the protocol crate's -//! [`recover_object`][rj]. Structured output got none of it. +//! tool-call arguments, repaired by +//! [`recover_tool_arguments`][rta] over +//! [`relaxed_json`][rj]. Structured output got none of it. //! -//! [rj]: tinytools_agent::repair::json::recover_object +//! [rta]: tinyinference_llm::providers::openai +//! [rj]: tinyinference_llm::providers::openai::relaxed_json //! //! # The ladder //! @@ -23,7 +25,7 @@ //! | `Strict` | nothing — the input was already valid | — | //! | `CodeFence` | ```` ```json … ``` ```` wrappers | ubiquitous | //! | `Slice` | prose around the value (`Here is the JSON: {…}`) | — | -//! | `Relaxed` | unquoted keys, doubled braces, leaked chat-template quote tokens | [`recover_object`][rj] | +//! | `Relaxed` | unquoted keys, doubled braces, leaked chat-template quote tokens | [`relaxed_json`][rj] | //! | `Closed` | truncated output: unterminated strings and unclosed brackets | LangChain `parse_partial_json` | //! //! # What it deliberately does not do @@ -94,35 +96,29 @@ pub fn parse_lenient(raw: &str) -> Option<(Value, JsonRepair)> { if unfenced != trimmed && let Ok(value) = serde_json::from_str::(unfenced) { - tinyagents_tracing::debug!( - "[structured::repair] recovered JSON by removing a markdown code fence" - ); + tracing::debug!("[structured::repair] recovered JSON by removing a markdown code fence"); return Some((value, JsonRepair::CodeFence)); } if let Some(sliced) = slice_json_span(unfenced) && let Ok(value) = serde_json::from_str::(sliced) { - tinyagents_tracing::debug!( + tracing::debug!( "[structured::repair] recovered JSON by slicing it out of surrounding text" ); return Some((value, JsonRepair::Slice)); } - // Reuses the protocol crate's repair ladder rather than a second, + // Reuses the crate's existing relaxed-JSON repairs rather than a second, // divergent implementation. It only yields objects, which is the shape a // JSON-Schema structured output almost always declares. if let Some(value) = tinytools_agent::repair::json::recover_object(unfenced) { - tinyagents_tracing::debug!( - "[structured::repair] recovered JSON through the relaxed-JSON repairs" - ); + tracing::debug!("[structured::repair] recovered JSON through the relaxed-JSON repairs"); return Some((value, JsonRepair::Relaxed)); } if let Some(value) = close_truncated(unfenced) { - tinyagents_tracing::debug!( - "[structured::repair] recovered JSON by closing a truncated value" - ); + tracing::debug!("[structured::repair] recovered JSON by closing a truncated value"); return Some((value, JsonRepair::Closed)); } diff --git a/crates/tinyagents-harness/src/structured/test.rs b/crates/tinyagents-harness/src/structured/test.rs index 2dade776..75d556dc 100644 --- a/crates/tinyagents-harness/src/structured/test.rs +++ b/crates/tinyagents-harness/src/structured/test.rs @@ -45,6 +45,65 @@ fn auto_strategy_uses_provider_schema_with_native_structured_output() { ); } +#[test] +fn default_structured_mode_native_wins_over_capability_inference() { + // A profile that could otherwise infer `ToolCall` (tool_calling: true, + // no native structured output) is overridden by an explicit + // `default_structured_mode`. + let profile = ModelProfile { + tool_calling: true, + native_structured_output: false, + default_structured_mode: Some(tinyinference_llm::model::StructuredMode::Native), + ..ModelProfile::default() + }; + assert_eq!( + StructuredStrategy::for_profile(Some(&profile)), + StructuredStrategy::ProviderSchema + ); +} + +#[test] +fn default_structured_mode_tool_wins_over_capability_inference() { + let profile = ModelProfile { + native_structured_output: true, + json_schema: true, + default_structured_mode: Some(tinyinference_llm::model::StructuredMode::Tool), + ..ModelProfile::default() + }; + assert_eq!( + StructuredStrategy::for_profile(Some(&profile)), + StructuredStrategy::ToolCall + ); +} + +#[test] +fn default_structured_mode_prompted_carries_the_profiles_template() { + let profile = ModelProfile { + default_structured_mode: Some(tinyinference_llm::model::StructuredMode::Prompted), + prompted_output_template: Some("Answer using this schema:".to_string()), + ..ModelProfile::default() + }; + assert_eq!( + StructuredStrategy::for_profile(Some(&profile)), + StructuredStrategy::Prompted { + template: Some("Answer using this schema:".to_string()) + } + ); +} + +#[test] +fn default_structured_mode_prompted_without_a_template_carries_none() { + let profile = ModelProfile { + default_structured_mode: Some(tinyinference_llm::model::StructuredMode::Prompted), + prompted_output_template: None, + ..ModelProfile::default() + }; + assert_eq!( + StructuredStrategy::for_profile(Some(&profile)), + StructuredStrategy::Prompted { template: None } + ); +} + #[test] fn provider_schema_parses_json_text() { let extractor = @@ -107,6 +166,7 @@ fn structured_output_parse_deserialises() { let output = StructuredOutput { value: json!({"value": "hello"}), raw_text: None, + variant: None, }; let parsed: Answer = output.parse().unwrap(); assert_eq!(parsed.value, "hello"); diff --git a/crates/tinyagents-harness/src/structured/types.rs b/crates/tinyagents-harness/src/structured/types.rs index 4fc00b04..56233d0b 100644 --- a/crates/tinyagents-harness/src/structured/types.rs +++ b/crates/tinyagents-harness/src/structured/types.rs @@ -19,16 +19,37 @@ use tinyinference_llm::model::ModelResponse; /// from the raw response text. /// * [`ToolCall`] – an artificial tool was exposed to the model; the structured /// value is read from the matching tool-call's `arguments` field. +/// * [`Prompted`] – for a model with no native schema or tool-calling support: +/// the schema is injected into the system prompt as instructions instead of +/// a provider API field, and extraction falls back to the same repair +/// ladder as [`ProviderSchema`]. Mirrors Pydantic AI's `PromptedOutput`. +/// * [`ToolCallUnion`] – one synthetic tool per schema variant; extraction +/// matches whichever variant's tool the model actually called and records +/// which one (see [`StructuredOutput::variant`]). /// /// [`ProviderSchema`]: StructuredStrategy::ProviderSchema /// [`ToolCall`]: StructuredStrategy::ToolCall -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +/// [`Prompted`]: StructuredStrategy::Prompted +/// [`ToolCallUnion`]: StructuredStrategy::ToolCallUnion +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum StructuredStrategy { /// Parse the JSON from the model's text response (provider-native mode). ProviderSchema, /// Read the arguments of a matching tool call. ToolCall, + /// Provider-native/tool-calling structured output is unavailable: the + /// schema is described in the system prompt instead, and extraction + /// parses the response text through the same repair ladder as + /// [`Self::ProviderSchema`]. + Prompted { + /// Custom instructions template injected ahead of the schema; `None` + /// uses [`super::default_prompted_template`]. + template: Option, + }, + /// A union output type: the model may satisfy the request by calling any + /// one of several synthetic tools, one per schema variant. + ToolCallUnion, } // --------------------------------------------------------------------------- @@ -41,13 +62,18 @@ pub enum StructuredStrategy { /// text that was parsed (useful for debugging or provider-native mode). /// /// [`ModelResponse`]: tinyinference_llm::model::ModelResponse -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] pub struct StructuredOutput { /// The extracted JSON value. pub value: Value, /// The raw assistant text that was parsed, when applicable. #[serde(default, skip_serializing_if = "Option::is_none")] pub raw_text: Option, + /// Which schema variant matched, for + /// [`StructuredStrategy::ToolCallUnion`]. `None` for every other + /// strategy. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub variant: Option, } // --------------------------------------------------------------------------- @@ -80,6 +106,10 @@ pub struct StructuredOutcome { /// verbatim: it names the schema and, for a validation failure, the exact /// failing instance path. pub error: Option, + /// Which schema variant matched, when extraction succeeded under + /// [`StructuredStrategy::ToolCallUnion`]. Mirrors + /// [`StructuredOutput::variant`]. + pub variant: Option, } impl StructuredOutcome { @@ -139,6 +169,11 @@ pub struct StructuredExtractor { /// The JSON Schema document. **Enforced**: every extracted value is checked /// against it by [`super::validate`] before it is returned, so a /// well-formed value of the wrong shape is a reported error rather than - /// silent garbage in `run.structured`. + /// silent garbage in `run.structured`. Unused (empty object) for + /// [`StructuredStrategy::ToolCallUnion`], which validates each match + /// against its own entry in [`Self::variants`] instead. pub(crate) schema: Value, + /// `(name, schema)` pairs for [`StructuredStrategy::ToolCallUnion`], one + /// per synthetic tool the model may call. Empty for every other strategy. + pub(crate) variants: Vec<(String, Value)>, } diff --git a/crates/tinyagents-harness/src/subagent/README.md b/crates/tinyagents-harness/src/subagent/README.md deleted file mode 100644 index ae05bb8c..00000000 --- a/crates/tinyagents-harness/src/subagent/README.md +++ /dev/null @@ -1,108 +0,0 @@ -# harness::subagent - -First-class sub-agents with recursion-depth tracking — the harness's flagship -"agents calling agents" surface. It lets one agent run another agent as a -child of itself, the in-harness analogue of `crate::graph::subgraph`'s -recursion on the graph side. - -## Design - -Three cooperating types: - -- **`SubAgent`** wraps an `AgentHarness` and runs it - as a child run one recursion level deeper than its caller. `invoke` / - `invoke_with_events` are the standalone entry points; `invoke_in_parent` - and `invoke_hosted_in_parent` thread a live `RunContext` so the child - inherits the parent's cancellation, event sink, stores, and (for the - hosted variant) host delegation authority. -- **`SubAgentTool`** adapts a `SubAgent` into a typed - `crate::tool::ToolDispatch`, so a parent agent can call another agent as an - ordinary tool call. Registered via - `crate::tool::ToolRegistry::register_dispatch` (not the plain - `tinytools::Tool` path — `SubAgentToolDeclaration::execute` refuses direct - calls and points at the typed dispatcher). -- **`SubAgentSession`** keeps a single `SubAgent` alive across - multiple turns, reusing the same harness while accumulating the transcript - — the post-completion, human-in-the-loop *reuse* primitive, distinct from - `crate::steering`'s mid-run *steering*. - -### Depth tracking - -Every run carries a `depth` in its `RunConfig` (top-level = `0`). Invoking a -sub-agent at `parent_depth` creates the child at `parent_depth + 1`, capped -by the child harness's `RunLimits::max_depth` (default -`RunLimits::DEFAULT_MAX_DEPTH` = `8`). Exceeding the cap fails fast with -`TinyAgentsError::SubAgentDepth` *before* any model call — `child_config` -computes and checks this on every invocation path. - -### Observability - -Every invocation brackets the child run with -`AgentEvent::SubAgentStarted`/`SubAgentCompleted`. Invoking through -`invoke_with_events`/`invoke_in_parent`/a `SubAgentSession` routes the -child's own events onto the shared parent sink, so a parent observer sees -the whole nested run tree; `SubAgentSession` additionally emits -`AgentEvent::SubAgentReused` on every send after the first. - -## Public surface - -- `SubAgent::new` / `with_system_prompt` / `name` / `description` / - `harness` — construction and accessors. -- `SubAgent::invoke` / `invoke_with_events` / `invoke_in_parent` / - `invoke_hosted_in_parent` — the four invocation entry points, differing in - how much of the live parent context they thread through (see the doc - comment on each for exactly what's inherited). -- `SubAgentTool::new` / `with_tool_name` / `with_parameters` / - `invoke_in_parent_context` — the typed-parent dispatch adapter; the last - method is also what its `ToolDispatch::execute` impl calls. -- `SubAgentSession::new` / `from_subagent` / `with_events` / - `with_parent_depth` / `subagent` / `transcript` / `turns` / `reset` / - `send` — reusable multi-turn session over one `SubAgent`. -- `ChildDataPolicy` — explicit parent→child application-data transform - (`new`, `child_data`); required by `SubAgentTool::new` so data inheritance - is never a silent `Default`. -- `SUBAGENT_INPUT_FIELD` — the argument key (`"input"`) a `SubAgentTool` - reads the child's user prompt from. - -## Files - -| File | Role | -| --- | --- | -| `mod.rs` | All impls: constructors, the four invoke paths, `SubAgentTool` dispatch, `SubAgentSession::send`, the private `SubAgentToolDeclaration`. | -| `types.rs` | Public struct/const definitions, re-exported via `pub use types::*`. | -| `test.rs` | `ChildDataPolicy`, `SubAgentTool` dispatch (data inheritance, depth-cap-as-recoverable-result, cancellation), `invoke_in_parent` propagation, depth-cap-before-model-work, `SubAgentSession` reuse. | - -## Key invariants - -- **Depth is checked before any model call.** `child_config` computes and - validates the child depth up front, so a run that would exceed - `max_depth` never reaches the network. -- **`invoke_in_parent` rejects a hosted parent.** A `RunContext` carrying - `host_authority` must go through `invoke_hosted_in_parent`, which resolves - the parent's delegate allowlist; falling through the generic path would - silently discard that authority (`TinyAgentsError::Validation`). -- **A hosted child cannot select an alternate host authority.** It always - re-enters through the exact capability bundle installed on the parent - (`crate::runtime::host_invocation_binding`); the child harness supplies - durable mechanics only. -- **A depth-cap or run-limit failure inside `SubAgentTool` is a recoverable - tool result, not a hard error** — `invoke_in_parent_context` turns - `SubAgentDepth`/`LimitExceeded`/`Timeout` into a `tinytools::ToolResult` - error the parent orchestrator can read as "delegated agent hit a limit," - not a thrown error that aborts the parent run. -- **`SubAgentSession` reuses the same harness and `Arc` across - every send** — nothing is reconstructed; only the transcript and turn - counter mutate. `reset()` clears the transcript without touching the - underlying harness. -- Run/thread ids minted here always suffix a process-unique sequence - (`crate::ids::next_seq`) so two invocations of the same sub-agent — or two - `SubAgentSession`s reusing it — never collide on run id. - -## Relation to neighbouring modules - -Built on `crate::runtime::AgentHarness` (the child agent loop), -`crate::context::RunContext` (the live run being extended into a child), and -`crate::tool::ToolDispatch` (the typed dispatch seam `SubAgentTool` -implements). See `crate::steering` for the complementary mid-run interruption -mechanism, and `crate::graph::subgraph` for the equivalent recursion primitive -on the graph side of the crate. diff --git a/crates/tinyagents-harness/src/summarization/README.md b/crates/tinyagents-harness/src/summarization/README.md index f847c521..b4b4792a 100644 --- a/crates/tinyagents-harness/src/summarization/README.md +++ b/crates/tinyagents-harness/src/summarization/README.md @@ -32,6 +32,17 @@ loop. tool results, and reasoning blocks that `Message::text()` drops — into summarizable text, so `ConcatSummarizer`'s default output isn't a column of bare role labels for a tool-driven run. +- **Compaction** (`compaction.rs`, `pub mod compaction`) is the durable, + rule-driven layer: `find_cut_point` (token-budget cut points, repaired via + `pairing.rs`), `summarize_with_split` (split-turn summarization + merge), + `Summarizer::summarize_request`/`SummaryRequest::previous_summary` + (iterative summaries), `CompactionRecord`/`CompactionSink`/ + `CompactionReason`, `CompactionContext`/`CompactionDecision` (the + `before_compaction` hook), and `OverflowClassifier` (table-driven provider + overflow detection). `ContextCompressionMiddleware` uses all of this for + its `before_model` (threshold) and `wrap_model` (overflow → compact → + retry) paths. Full contract: + `docs/modules/harness/compaction.md`. ## Public surface @@ -84,7 +95,9 @@ loop. | `pairing.rs` | Tool-call-pairing-safe cut-point repair. | | `render.rs` | `render_message_for_summary`. | | `trim.rs` | `trim_messages`/`trim_messages_with`/`trim_messages_to_token_budget_with`. | +| `compaction.rs` | `find_cut_point`, `summarize_with_split`, `OverflowClassifier`, `CompactionContext`/`CompactionDecision`. | | `test.rs` | Coverage for token estimation, trim strategies, pairing repair, policy triggering/planning, and `ConcatSummarizer`. | +| `compaction/test.rs` | Coverage for cut points, split-turn merge, iterative summaries, `OverflowClassifier`. | ## Key invariants diff --git a/crates/tinyagents-harness/src/summarization/compaction.rs b/crates/tinyagents-harness/src/summarization/compaction.rs new file mode 100644 index 00000000..d0776647 --- /dev/null +++ b/crates/tinyagents-harness/src/summarization/compaction.rs @@ -0,0 +1,494 @@ +//! Durable, rule-driven compaction: token-budget cut points, split-turn +//! summarization, and overflow classification. +//! +//! This is the harness's port of pi's `compaction.ts` / +//! `overflow.ts` (`docs/runtime-comparison/pi.md` §4.5): where +//! [`super::types::SummarizationPolicy`] decides *when* to compact and splits +//! by a fixed `keep_last` message count, this module adds a *token-budget* +//! cut point ([`find_cut_point`]), the "a single turn is itself too big to +//! summarize in one call" case ([`summarize_with_split`]), and a table-driven +//! classifier for turning a provider's context-overflow error into a typed +//! [`OverflowInfo`] ([`OverflowClassifier`]) so a caller can drive an +//! overflow → compact → retry loop +//! ([`crate::middleware::ContextCompressionMiddleware`]). + +use std::sync::Arc; + +use tinyinference_llm::message::Message; + +use crate::error::{Result, TinyAgentsError}; + +use super::pairing::find_safe_cutoff_point; +use super::trim::partition_system; +use super::types::{CompactionReason, Summarizer, SummaryRecord, SummaryRequest}; + +// --------------------------------------------------------------------------- +// Cut points +// --------------------------------------------------------------------------- + +/// A validated cut point into a message slice: everything before +/// [`Self::index`] (in the non-system message slice the cut was computed +/// over) is old enough to fold into a summary; everything from `index` on is +/// kept verbatim. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CutPoint { + /// Index, into the non-system message slice, of the first message that + /// stays verbatim. Always a safe boundary: never inside an assistant + /// tool-call turn and its answering tool results (see + /// [`super::pairing::find_safe_cutoff_point`]). + pub index: usize, + /// Estimated tokens of the messages that would be folded into a summary + /// (`non_system[..index]`). + pub tokens_before: u64, + /// Estimated tokens of the messages kept verbatim (`non_system[index..]`). + pub tokens_after: u64, +} + +/// Finds a cut point that keeps at least `keep_recent_tokens` worth of the +/// most recent messages verbatim, walking `messages` newest-first. +/// +/// Port of pi's `findCutPoint` (`compaction.ts:370`). System messages are +/// excluded from consideration (partitioned out first, as every other +/// pairing-aware operation in this module does) and are always implicitly +/// kept by the caller. Returns `None` when there is nothing to cut: either +/// `messages` has no non-system content, or the whole non-system slice +/// already fits within `keep_recent_tokens` (nothing old enough to +/// summarize). +/// +/// # Safety +/// +/// The raw token-budget candidate is repaired with +/// [`find_safe_cutoff_point`] before being returned, so the result never +/// splits an assistant tool-call turn from the tool results answering it — +/// the same invariant [`super::types::SummarizationPolicy::plan`] enforces +/// for its count-based split. `keep_recent_tokens` is therefore a *minimum* +/// retained budget, not an exact one: the repaired boundary may keep +/// (never drop) a few extra tokens to preserve pairing. +/// +/// # Example +/// +/// ``` +/// use tinyagents_harness::summarization::{estimate_tokens, find_cut_point}; +/// use tinyinference_llm::message::Message; +/// +/// let messages = vec![ +/// Message::user("hello"), +/// Message::assistant("hi there"), +/// Message::user("what's the weather?"), +/// ]; +/// let cut = find_cut_point(&messages, 4, |m| estimate_tokens(&m.text())).unwrap(); +/// assert!(cut.index > 0); +/// ``` +pub fn find_cut_point( + messages: &[Message], + keep_recent_tokens: u64, + estimator: impl Fn(&Message) -> u64, +) -> Option { + let (_system, non_system) = partition_system(messages); + if non_system.is_empty() { + return None; + } + + // Walk newest-first, accumulating tokens until the budget would be + // exceeded; `idx` lands on the oldest message still inside the budget. + let mut acc = 0u64; + let mut idx = 0usize; + for i in (0..non_system.len()).rev() { + let tokens = estimator(&non_system[i]); + if acc + tokens > keep_recent_tokens { + idx = i + 1; + break; + } + acc += tokens; + idx = i; + } + + let safe_idx = find_safe_cutoff_point(&non_system, idx); + if safe_idx == 0 { + // Everything fits (or pairing repair pulled the cut all the way back + // to the start) — nothing old enough to compact. + return None; + } + + let tokens_before: u64 = non_system[..safe_idx].iter().map(&estimator).sum(); + let tokens_after: u64 = non_system[safe_idx..].iter().map(&estimator).sum(); + + Some(CutPoint { + index: safe_idx, + tokens_before, + tokens_after, + }) +} + +// --------------------------------------------------------------------------- +// Split-turn summarization +// --------------------------------------------------------------------------- + +/// Summarizes `messages` with `summarizer`, splitting into two halves and +/// merging their summaries when `messages` alone estimates above +/// `max_turn_tokens` — the "a single turn is too big for one summarization +/// call" case a fixed-size compaction batch can otherwise hit (a turn with a +/// huge tool result, for instance). +/// +/// `previous_summary`, when set, is threaded to the *first* half's +/// [`SummaryRequest::previous_summary`] only — the second half has no +/// predecessor of its own within this split, and +/// [`Summarizer::merge`] is what reconciles the two halves into one summary +/// that itself becomes the next call's `previous_summary`. +/// +/// When `messages` fits under `max_turn_tokens`, or no interior message +/// index is a safe split boundary (see +/// [`super::pairing::find_safe_cutoff_point`] — this can happen when the +/// whole turn is a single indivisible tool-call/tool-result pair), the whole +/// slice is summarized in one call instead of forcing an unsafe split. +pub async fn summarize_with_split( + summarizer: &dyn Summarizer, + messages: &[Message], + max_turn_tokens: u64, + previous_summary: Option, + estimator: impl Fn(&Message) -> u64, +) -> Result { + if messages.is_empty() { + return Err(TinyAgentsError::Validation( + "cannot summarize an empty turn".into(), + )); + } + + let total: u64 = messages.iter().map(&estimator).sum(); + if total <= max_turn_tokens || messages.len() < 2 { + return summarizer + .summarize_request(&SummaryRequest { + messages: messages.to_vec(), + previous_summary, + }) + .await; + } + + let midpoint = messages.len() / 2; + let split = find_safe_cutoff_point(messages, midpoint); + if split == 0 || split >= messages.len() { + // No safe interior boundary — fall back to one call rather than + // breaking tool-call pairing to force a split. + return summarizer + .summarize_request(&SummaryRequest { + messages: messages.to_vec(), + previous_summary, + }) + .await; + } + + let (first_half, second_half) = messages.split_at(split); + + let first_summary = summarizer + .summarize_request(&SummaryRequest { + messages: first_half.to_vec(), + previous_summary, + }) + .await?; + let second_summary = summarizer + .summarize_request(&SummaryRequest { + messages: second_half.to_vec(), + previous_summary: None, + }) + .await?; + + summarizer.merge(&[first_summary, second_summary]).await +} + +// --------------------------------------------------------------------------- +// before_compaction hook +// --------------------------------------------------------------------------- + +/// What [`super::types::CompactionRecord`]-producing code hands to a +/// `before_compaction` hook so it can decide whether to proceed, decline, or +/// substitute its own summary — pi's `AgentHarness` compaction operation +/// (`docs/runtime-comparison/pi.md` §4.5). +#[derive(Clone, Debug)] +pub struct CompactionContext { + /// Why this compaction is about to run. + pub reason: CompactionReason, + /// Estimated tokens of the transcript immediately before compaction. + pub tokens_before: u64, + /// Number of messages that would be folded into the summary. + pub to_summarize_count: usize, + /// Number of messages that would be kept verbatim. + pub to_keep_count: usize, +} + +/// A `before_compaction` hook's decision. +#[derive(Clone, Debug)] +pub enum CompactionDecision { + /// Run the compaction as planned. + Proceed, + /// Skip this compaction; leave the transcript untouched. The caller that + /// triggered compaction is responsible for deciding what happens next + /// (for the overflow → compact → retry path, a decline means the + /// original provider error propagates instead of a retry). + Decline, + /// Run the compaction, but install this text as the summary instead of + /// calling the configured [`Summarizer`]. + UseSummary(String), +} + +// --------------------------------------------------------------------------- +// Overflow classification +// --------------------------------------------------------------------------- + +/// Best-effort structured detail extracted from a classified overflow error: +/// the token count the request attempted to send and the provider's context +/// limit, when either is recoverable from the error text. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct OverflowInfo { + /// Tokens the failed request attempted to send, when the provider's + /// message reports it. + pub requested: Option, + /// The provider's context-window limit, when the provider's message + /// reports it. + pub limit: Option, +} + +/// A cheap, structured view of a model-call failure handed to an +/// [`OverflowClassifier`] pattern: the provider's raw message plus whatever +/// structured detail is available. +#[derive(Clone, Copy, Debug)] +pub struct OverflowProbe<'a> { + /// The provider's human-readable error message, verbatim. + pub message: &'a str, + /// The provider's error code or type, when reported + /// (e.g. `"context_length_exceeded"`). + pub code: Option<&'a str>, + /// The transport HTTP status, when the failure came from an HTTP + /// response. + pub status: Option, +} + +type OverflowMatchFn = Arc) -> Option + Send + Sync>; + +#[derive(Clone)] +struct OverflowPattern { + label: &'static str, + matches: OverflowMatchFn, +} + +impl std::fmt::Debug for OverflowPattern { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OverflowPattern") + .field("label", &self.label) + .finish_non_exhaustive() + } +} + +/// A table-driven matcher that classifies a [`TinyAgentsError`] as a +/// provider context-window overflow, or `None` for every other failure. +/// +/// Port of pi's `isContextOverflow` (`overflow.ts:135`), generalized to a +/// pattern table so a host can register a pattern for a provider or local +/// server this crate does not ship a built-in for (see [`Self::with_pattern`]). +/// [`Self::default`] ships patterns for OpenAI's `context_length_exceeded`, +/// Anthropic's `"prompt is too long"`, a generic `"maximum context length"` +/// phrasing several providers share, local llama.cpp's `n_ctx` messages, and +/// a fallback for an HTTP 400/413 whose body mentions the context window. +/// +/// # Example +/// +/// ``` +/// use tinyagents_harness::error::TinyAgentsError; +/// use tinyagents_harness::summarization::OverflowClassifier; +/// +/// let classifier = OverflowClassifier::default(); +/// let err = TinyAgentsError::Model( +/// "This model's maximum context length is 8192 tokens. \ +/// However, your messages resulted in 9000 tokens.".to_string(), +/// ); +/// let info = classifier.classify(&err).expect("classified as overflow"); +/// assert_eq!(info.limit, Some(8192)); +/// assert_eq!(info.requested, Some(9000)); +/// ``` +#[derive(Clone, Debug)] +pub struct OverflowClassifier { + patterns: Vec, +} + +impl Default for OverflowClassifier { + fn default() -> Self { + Self::with_builtins() + } +} + +impl OverflowClassifier { + /// An empty classifier with no patterns — every error classifies as + /// `None`. Prefer [`OverflowClassifier::default()`] for the built-in + /// provider patterns; use this only to build a classifier with entirely + /// custom patterns. + pub fn empty() -> Self { + Self { + patterns: Vec::new(), + } + } + + /// Registers an additional pattern, checked after every pattern already + /// registered (built-ins first when starting from [`Self::default`]). + /// + /// `matcher` returns `Some(info)` when the probe indicates an overflow + /// (`info`'s fields may both be `None` when no numeric detail is + /// recoverable — the match itself is still meaningful) and `None` + /// otherwise. + pub fn with_pattern( + mut self, + label: &'static str, + matcher: impl Fn(&OverflowProbe<'_>) -> Option + Send + Sync + 'static, + ) -> Self { + self.patterns.push(OverflowPattern { + label, + matches: Arc::new(matcher), + }); + self + } + + /// The labels of every registered pattern, in check order. Exposed for + /// tests and diagnostics. + pub fn pattern_labels(&self) -> Vec<&'static str> { + self.patterns.iter().map(|p| p.label).collect() + } + + /// Classifies `error`, returning `Some(OverflowInfo)` when a registered + /// pattern (or [`TinyAgentsError::ContextOverflow`] directly) matches. + pub fn classify(&self, error: &TinyAgentsError) -> Option { + match error { + // Already a typed overflow — trust it directly, and try the + // pattern table over its message only to fill in numeric detail. + TinyAgentsError::ContextOverflow { message, .. } => { + let probe = OverflowProbe { + message, + code: None, + status: None, + }; + Some(self.match_probe(&probe).unwrap_or_default()) + } + TinyAgentsError::Provider(provider_error) => { + let probe = OverflowProbe { + message: &provider_error.message, + code: provider_error.code.as_deref(), + status: provider_error.status, + }; + self.match_probe(&probe) + } + TinyAgentsError::Model(message) => { + let probe = OverflowProbe { + message, + code: None, + status: None, + }; + self.match_probe(&probe) + } + _ => None, + } + } + + fn match_probe(&self, probe: &OverflowProbe<'_>) -> Option { + self.patterns + .iter() + .find_map(|pattern| (pattern.matches)(probe)) + } +} + +/// The built-in pattern table: OpenAI, Anthropic, a generic phrasing, local +/// llama.cpp, and an HTTP-status fallback, checked in that order. +impl OverflowClassifier { + /// Builds the classifier [`Default`] returns. A free function so + /// `Default::default()` and any caller rebuilding the built-in set (for + /// example after calling [`Self::empty`]) share one definition. + fn with_builtins() -> Self { + Self::empty() + .with_pattern("openai", |probe| { + let code_hit = probe.code == Some("context_length_exceeded"); + let message_hit = probe.message.contains("context_length_exceeded"); + if !(code_hit || message_hit) { + return None; + } + let numbers = extract_numbers(probe.message); + Some(OverflowInfo { + limit: numbers.first().copied(), + requested: numbers.get(1).copied(), + }) + }) + .with_pattern("anthropic", |probe| { + if !probe.message.to_lowercase().contains("prompt is too long") { + return None; + } + let numbers = extract_numbers(probe.message); + Some(OverflowInfo { + requested: numbers.first().copied(), + limit: numbers.get(1).copied(), + }) + }) + .with_pattern("llama_cpp", |probe| { + if !probe.message.contains("n_ctx") { + return None; + } + let numbers = extract_numbers(probe.message); + Some(OverflowInfo { + limit: numbers.first().copied(), + requested: numbers.get(1).copied(), + }) + }) + .with_pattern("generic", |probe| { + if !probe + .message + .to_lowercase() + .contains("maximum context length") + { + return None; + } + let numbers = extract_numbers(probe.message); + Some(OverflowInfo { + limit: numbers.first().copied(), + requested: numbers.get(1).copied(), + }) + }) + .with_pattern("http_body", |probe| { + let status_hit = matches!(probe.status, Some(400) | Some(413)); + if !status_hit { + return None; + } + let lower = probe.message.to_lowercase(); + let body_hit = lower.contains("context") + || lower.contains("too long") + || lower.contains("token limit"); + if !body_hit { + return None; + } + Some(OverflowInfo::default()) + }) + } +} + +/// Extracts every run of ASCII digits from `text` as a `u64`, in order of +/// appearance, tolerating `,` thousands separators inside a run (`"9,000"` → +/// `9000`). Best-effort: used only to fill optional +/// [`OverflowInfo`] detail, never to decide whether an error is an overflow. +fn extract_numbers(text: &str) -> Vec { + let mut numbers = Vec::new(); + let mut current = String::new(); + for ch in text.chars() { + if ch.is_ascii_digit() { + current.push(ch); + } else if ch == ',' && !current.is_empty() { + continue; + } else if !current.is_empty() { + if let Ok(n) = current.parse::() { + numbers.push(n); + } + current.clear(); + } + } + if !current.is_empty() + && let Ok(n) = current.parse::() + { + numbers.push(n); + } + numbers +} + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-harness/src/summarization/compaction/test.rs b/crates/tinyagents-harness/src/summarization/compaction/test.rs new file mode 100644 index 00000000..695e00f1 --- /dev/null +++ b/crates/tinyagents-harness/src/summarization/compaction/test.rs @@ -0,0 +1,389 @@ +//! Tests for cut-point discovery, split-turn summarization, and overflow +//! classification. + +use std::sync::Mutex; + +use async_trait::async_trait; +use serde_json::json; +use tinyinference_llm::message::{AssistantMessage, Message}; +use tinyinference_llm::tool::ToolCall; + +use super::*; +use crate::error::TinyAgentsError; +use crate::summarization::{ConcatSummarizer, SummaryRecord, tool_pairing_is_intact}; +use crate::token_estimation::estimate_message_tokens; + +fn assistant_calling(ids: &[&str]) -> Message { + Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: ids + .iter() + .map(|id| ToolCall::new(*id, "lookup", json!({"q": "rust"}))) + .collect(), + usage: None, + origin: None, + }) +} + +// --------------------------------------------------------------------------- +// find_cut_point +// --------------------------------------------------------------------------- + +#[test] +fn find_cut_point_never_splits_a_tool_pair() { + // `[user, assistant(tool_calls=[c1]), tool(c1), assistant("done")]`. The + // budget below is chosen (from the messages' own estimated weights) so + // the *naive* newest-first walk lands the cut exactly on `tool(c1)` — + // precisely the split the repair exists to prevent. + let user = Message::user("weather?"); + let call = assistant_calling(&["c1"]); + let result = Message::tool("c1", "sunny and warm today, 21 degrees"); + let done = Message::assistant("It's sunny and warm."); + let non_system = vec![user, call, result.clone(), done.clone()]; + + let budget = estimate_message_tokens(&result) + estimate_message_tokens(&done); + let cut = find_cut_point(&non_system, budget, estimate_message_tokens) + .expect("some cut point should be found"); + + // The naive (unrepaired) boundary would be index 2 (`tool(c1)` itself); + // confirm the repair actually moved it, not that it happened to already + // be safe. + assert_ne!( + cut.index, 2, + "test setup did not land the naive cut on the tool result" + ); + assert!(!matches!(non_system[cut.index], Message::Tool(_))); + + let kept = &non_system[cut.index..]; + assert!( + tool_pairing_is_intact(kept), + "cut point {} orphans a tool pair: {kept:?}", + cut.index + ); +} + +#[test] +fn find_cut_point_respects_keep_recent_tokens() { + let messages = vec![ + Message::user("one"), + Message::user("two"), + Message::user("three"), + Message::user("four"), + ]; + // A generous budget should keep everything (nothing old enough to cut). + let generous = find_cut_point(&messages, 10_000, estimate_message_tokens); + assert!(generous.is_none()); + + // A tiny budget keeps only the most recent message(s), and whatever it + // keeps meets or exceeds the requested budget (a cut point is a floor, + // not an exact count). + let tiny = find_cut_point(&messages, 1, estimate_message_tokens) + .expect("a tiny budget should still find a cut point"); + assert!(tiny.index > 0); + assert!(tiny.tokens_after >= 1); + assert_eq!( + tiny.tokens_before + tiny.tokens_after, + messages.iter().map(estimate_message_tokens).sum::() + ); +} + +#[test] +fn find_cut_point_none_when_nothing_to_summarize() { + let messages = vec![Message::system("sys"), Message::user("hi")]; + assert!(find_cut_point(&messages, 10_000, estimate_message_tokens).is_none()); +} + +#[test] +fn find_cut_point_none_for_only_system_messages() { + let messages = vec![Message::system("sys")]; + assert!(find_cut_point(&messages, 0, estimate_message_tokens).is_none()); +} + +// --------------------------------------------------------------------------- +// summarize_with_split +// --------------------------------------------------------------------------- + +/// Records every call `summarize_request`/`merge` received, so tests can +/// assert both the split happened and what was threaded through it. +#[derive(Default)] +struct RecordingSummarizer { + requests: Mutex>, + merges: Mutex, +} + +#[async_trait] +impl Summarizer for RecordingSummarizer { + async fn summarize(&self, messages: &[Message]) -> Result { + self.summarize_request(&SummaryRequest::new(messages.to_vec())) + .await + } + + async fn summarize_request(&self, request: &SummaryRequest) -> Result { + self.requests.lock().unwrap().push(request.clone()); + ConcatSummarizer.summarize(&request.messages).await + } + + async fn merge(&self, summaries: &[SummaryRecord]) -> Result { + *self.merges.lock().unwrap() += 1; + // Delegate to the default concatenation merge via ConcatSummarizer's + // inherited default (ConcatSummarizer never overrides `merge`). + ConcatSummarizer.merge(summaries).await + } +} + +#[tokio::test] +async fn split_turn_summarizes_whole_slice_when_under_budget() { + let summarizer = RecordingSummarizer::default(); + let messages = vec![Message::user("a"), Message::user("b")]; + let total: u64 = messages.iter().map(estimate_message_tokens).sum(); + + summarize_with_split( + &summarizer, + &messages, + total + 100, + None, + estimate_message_tokens, + ) + .await + .unwrap(); + + assert_eq!(summarizer.requests.lock().unwrap().len(), 1); + assert_eq!(*summarizer.merges.lock().unwrap(), 0); +} + +#[tokio::test] +async fn split_turn_splits_and_merges_when_over_budget() { + let summarizer = RecordingSummarizer::default(); + let big = "word ".repeat(50); + let messages = vec![ + Message::user(format!("first half {big}")), + Message::user(format!("second half {big}")), + ]; + let total: u64 = messages.iter().map(estimate_message_tokens).sum(); + + let merged = summarize_with_split( + &summarizer, + &messages, + total / 2, + Some("previous run's summary".to_string()), + estimate_message_tokens, + ) + .await + .unwrap(); + + assert_eq!(summarizer.requests.lock().unwrap().len(), 2); + assert_eq!(*summarizer.merges.lock().unwrap(), 1); + // The merged text carries content from both halves. + let text = merged.summary.text(); + assert!(text.contains("first half")); + assert!(text.contains("second half")); +} + +#[tokio::test] +async fn split_turn_threads_previous_summary_to_first_half_only() { + let summarizer = RecordingSummarizer::default(); + let big = "word ".repeat(50); + let messages = vec![ + Message::user(format!("alpha {big}")), + Message::user(format!("beta {big}")), + ]; + let total: u64 = messages.iter().map(estimate_message_tokens).sum(); + + summarize_with_split( + &summarizer, + &messages, + total / 2, + Some("prior".to_string()), + estimate_message_tokens, + ) + .await + .unwrap(); + + let requests = summarizer.requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].previous_summary.as_deref(), Some("prior")); + assert_eq!(requests[1].previous_summary, None); +} + +// --------------------------------------------------------------------------- +// Iterative summaries (SummaryRequest.previous_summary) +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn iterative_summary_receives_the_previous_summary() { + let summarizer = RecordingSummarizer::default(); + let messages = vec![Message::user("new turn")]; + + summarizer + .summarize_request(&SummaryRequest::new(messages).with_previous_summary("earlier context")) + .await + .unwrap(); + + let requests = summarizer.requests.lock().unwrap(); + assert_eq!( + requests[0].previous_summary.as_deref(), + Some("earlier context") + ); +} + +#[tokio::test] +async fn default_summarize_request_ignores_previous_summary() { + // `ConcatSummarizer` never overrides `summarize_request`, so the default + // trait method's back-compat delegation to `summarize` applies: the + // previous summary is accepted but not threaded into the (no-LLM) output. + let request = SummaryRequest::new(vec![Message::user("hi")]) + .with_previous_summary("ignored by ConcatSummarizer"); + let record = ConcatSummarizer.summarize_request(&request).await.unwrap(); + assert!(record.summary.text().contains("hi")); +} + +// --------------------------------------------------------------------------- +// OverflowClassifier +// --------------------------------------------------------------------------- + +#[test] +fn classifies_typed_context_overflow_directly() { + let classifier = OverflowClassifier::default(); + let err = TinyAgentsError::ContextOverflow { + provider: "openai".to_string(), + model: Some("gpt-test".to_string()), + message: "too long".to_string(), + }; + assert!(classifier.classify(&err).is_some()); +} + +#[test] +fn classifies_openai_context_length_exceeded_by_code() { + let classifier = OverflowClassifier::default(); + let err = TinyAgentsError::Provider(Box::new(tinyinference_llm::model::ProviderError { + provider: "openai".to_string(), + model: None, + status: Some(400), + code: Some("context_length_exceeded".to_string()), + message: "This model's maximum context length is 8192 tokens. However, your messages \ + resulted in 9000 tokens." + .to_string(), + retryable: false, + retry_after_ms: None, + raw: None, + partial_message: None, + stop_reason: None, + })); + let info = classifier.classify(&err).expect("classified as overflow"); + assert_eq!(info.limit, Some(8192)); + assert_eq!(info.requested, Some(9000)); +} + +#[test] +fn classifies_anthropic_prompt_is_too_long() { + let classifier = OverflowClassifier::default(); + let err = + TinyAgentsError::Model("prompt is too long: 210000 tokens > 200000 maximum".to_string()); + let info = classifier.classify(&err).expect("classified as overflow"); + assert_eq!(info.requested, Some(210_000)); + assert_eq!(info.limit, Some(200_000)); +} + +#[test] +fn classifies_generic_maximum_context_length_phrasing() { + let classifier = OverflowClassifier::default(); + let err = TinyAgentsError::Model( + "request exceeds the maximum context length of 4096 tokens (sent 5000)".to_string(), + ); + let info = classifier.classify(&err).expect("classified as overflow"); + assert_eq!(info.limit, Some(4096)); + assert_eq!(info.requested, Some(5000)); +} + +#[test] +fn classifies_local_llama_cpp_n_ctx_messages() { + let classifier = OverflowClassifier::default(); + let err = + TinyAgentsError::Model("context size exceeded (n_ctx = 4096, tokens = 4300)".to_string()); + let info = classifier.classify(&err).expect("classified as overflow"); + assert_eq!(info.limit, Some(4096)); + assert_eq!(info.requested, Some(4300)); +} + +#[test] +fn classifies_http_400_and_413_bodies_mentioning_the_context_window() { + let classifier = OverflowClassifier::default(); + let err_400 = TinyAgentsError::Provider(Box::new(tinyinference_llm::model::ProviderError { + provider: "custom".to_string(), + model: None, + status: Some(400), + code: None, + message: "request too long for the context window".to_string(), + retryable: false, + retry_after_ms: None, + raw: None, + partial_message: None, + stop_reason: None, + })); + assert!(classifier.classify(&err_400).is_some()); + + let err_413 = TinyAgentsError::Provider(Box::new(tinyinference_llm::model::ProviderError { + provider: "custom".to_string(), + model: None, + status: Some(413), + code: None, + message: "payload too large: context exceeded".to_string(), + retryable: false, + retry_after_ms: None, + raw: None, + partial_message: None, + stop_reason: None, + })); + assert!(classifier.classify(&err_413).is_some()); +} + +#[test] +fn does_not_classify_unrelated_provider_errors() { + let classifier = OverflowClassifier::default(); + let err = TinyAgentsError::Provider(Box::new(tinyinference_llm::model::ProviderError { + provider: "openai".to_string(), + model: None, + status: Some(429), + code: Some("rate_limit_exceeded".to_string()), + message: "rate limit exceeded, please retry later".to_string(), + retryable: true, + retry_after_ms: Some(1000), + raw: None, + partial_message: None, + stop_reason: None, + })); + assert!(classifier.classify(&err).is_none()); + assert!( + classifier + .classify(&TinyAgentsError::Tool("boom".into())) + .is_none() + ); +} + +#[test] +fn with_pattern_extends_the_classifier() { + let classifier = OverflowClassifier::empty().with_pattern("acme", |probe| { + probe + .message + .contains("ACME_CONTEXT_FULL") + .then_some(OverflowInfo::default()) + }); + assert_eq!(classifier.pattern_labels(), vec!["acme"]); + let err = TinyAgentsError::Model("ACME_CONTEXT_FULL: cannot proceed".to_string()); + assert!(classifier.classify(&err).is_some()); + assert!( + classifier + .classify(&TinyAgentsError::Model("something else".to_string())) + .is_none() + ); +} + +#[test] +fn default_classifier_has_the_documented_built_in_patterns() { + let labels = OverflowClassifier::default().pattern_labels(); + assert_eq!( + labels, + vec!["openai", "anthropic", "llama_cpp", "generic", "http_body"] + ); +} diff --git a/crates/tinyagents-harness/src/summarization/mod.rs b/crates/tinyagents-harness/src/summarization/mod.rs index b0f2e4a0..a7102835 100644 --- a/crates/tinyagents-harness/src/summarization/mod.rs +++ b/crates/tinyagents-harness/src/summarization/mod.rs @@ -19,11 +19,16 @@ //! All policy decisions are explicit data types, never hidden behaviour. Callers //! choose when to call, what to pass, and how to handle the result. +pub mod compaction; pub mod pairing; mod render; mod trim; mod types; +pub use compaction::{ + CompactionContext, CompactionDecision, CutPoint, OverflowClassifier, OverflowInfo, + OverflowProbe, find_cut_point, summarize_with_split, +}; pub use pairing::{ advance_past_orphan_tools, find_safe_cutoff_point, is_tool_calling_assistant, retract_orphan_tool_calls, tool_pairing_is_intact, @@ -253,7 +258,7 @@ impl SummarizationPolicy { let requested_split = non_system.len() - self.keep_last; let split = find_safe_cutoff_point(&non_system, requested_split); if split != requested_split { - tinyagents_tracing::debug!( + tracing::debug!( "[summarization::plan] keep_last={} moved split {requested_split} -> {split} to preserve tool-call pairing", self.keep_last ); diff --git a/crates/tinyagents-harness/src/summarization/pairing.rs b/crates/tinyagents-harness/src/summarization/pairing.rs index bd4a9987..0f181065 100644 --- a/crates/tinyagents-harness/src/summarization/pairing.rs +++ b/crates/tinyagents-harness/src/summarization/pairing.rs @@ -105,14 +105,14 @@ pub fn find_safe_cutoff_point(messages: &[Message], cutoff_index: usize) -> usiz for index in (0..cutoff_index).rev() { let declared = declared_call_ids(&messages[index]); if !declared.is_empty() && declared.intersection(&orphan_ids).next().is_some() { - tinyagents_tracing::debug!( + tracing::debug!( "[summarization::pairing] cutoff {cutoff_index} split a tool pair; moving back to {index} to keep the assistant tool-call turn" ); return index; } } - tinyagents_tracing::debug!( + tracing::debug!( "[summarization::pairing] cutoff {cutoff_index} has no matching assistant tool-call turn; advancing to {past_run} to drop unpairable tool results" ); past_run @@ -132,7 +132,7 @@ pub fn advance_past_orphan_tools(messages: &[Message], cutoff_index: usize) -> u index += 1; } if index != cutoff_index { - tinyagents_tracing::debug!( + tracing::debug!( "[summarization::pairing] dropped {} leading orphan tool result(s) at cutoff {cutoff_index}", index - cutoff_index ); @@ -155,7 +155,7 @@ pub fn retract_orphan_tool_calls(messages: &[Message], end_index: usize) -> usiz end -= 1; } if end != end_index.min(messages.len()) { - tinyagents_tracing::debug!( + tracing::debug!( "[summarization::pairing] retracted retained prefix end from {end_index} to {end} to drop unanswered assistant tool call(s)" ); } diff --git a/crates/tinyagents-harness/src/summarization/render.rs b/crates/tinyagents-harness/src/summarization/render.rs index 579dbf7f..1b1844a9 100644 --- a/crates/tinyagents-harness/src/summarization/render.rs +++ b/crates/tinyagents-harness/src/summarization/render.rs @@ -41,6 +41,7 @@ fn role_label(message: &Message) -> &'static str { Message::User(_) => "user", Message::Assistant(_) => "assistant", Message::Tool(_) => "tool", + Message::Custom(_) => "custom", } } @@ -85,6 +86,9 @@ fn render_content(content: &[ContentBlock]) -> Vec { "{}", elide(&value.to_string()) )), + ContentBlock::Audio(_) => Some("